…ms_coalesced (deepspeedai#8073)

This PR fixes the `_allgather_params_coalesced` method in
`partition_parameters.py`. The change ensures that each `flat_tensor` is
created with the correct data type by referencing the corresponding
parameter in `param_list`, rather than always using the first
parameter's data type.

Fix deepspeedai#8072.

### Problem

`_allgather_params_coalesced` allocates all output buffers using the
dtype of the first parameter in param_list:

```python
# before
for psize in partition_sizes:
    flat_tensor = torch.empty(tensor_size, dtype=param_list[0].ds_tensor.dtype, ...)
```

This assumed every persistent parameter shares the same dtype. The
assumption was incidentally maintained before 0.19.2 because
`_configure_distributed_model` called `module.bfloat16()`
unconditionally, normalising all persistent parameters (including PEFT
LoRA adapters) to a uniform dtype.

PR deepspeedai#8066 "Mixed-precision:
per-policy param/buffer dtype cast (preserve fp32 buffers)" (commit
b919284) correctly stopped casting ZeRO-Init model params, but exposed
the latent bug: PEFT's default `autocast_adapter_dtype=True` keeps LoRA
adapters in fp32 even when the base model is bf16.
`persistent_parameters` therefore ends up with mixed dtypes (bf16
base-model params + fp32 LoRA params), and the mismatch between a bf16
output buffer and a fp32 input tensor raises:
> TypeError: output tensor must have the same type as input tensor

- Reported in: deepspeedai#8072
- Downstream tracking: huggingface/trl#6089

### Solution

Allocate each output buffer with the dtype of its own parameter:

```python
# after
for i, psize in enumerate(partition_sizes):
    flat_tensor = torch.empty(tensor_size, dtype=param_list[i].ds_tensor.dtype, ...)
```

This removes the shared-dtype assumption at the source rather than
relying on upstream callers to normalise dtypes before calling
`_allgather_params_coalesced`.

### Changes

- Corrected tensor data type selection in `_allgather_params_coalesced`
to use the data type of each parameter in `param_list`, ensuring proper
handling of mixed data types.

Signed-off-by: Albert Villanova del Moral <8515462+albertvillanova@users.noreply.github.com>

…ffers) (deepspeedai#8066)

## Summary
- Add `data_types.param_dtype` and `data_types.buffer_dtype` (both
default `None`), mirroring FSDP `MixedPrecisionPolicy`.
- Replace the blanket `module.half()` / `module.bfloat16()` in
`_configure_distributed_model` with a targeted cast: parameters go to
`param_dtype`; floating buffers keep their loaded dtype unless
`buffer_dtype` is explicitly set.

## Motivation
The blanket cast downcasts every floating buffer, including the rotary
`inv_freq` buffer that HF/FSDP2 keep in fp32. On long contexts the bf16
`inv_freq` loses precision, RoPE angles drift, and logits/grads diverge
from the FSDP2 reference. Preserving fp32 buffers by default fixes this;
`buffer_dtype` is the escape hatch to reproduce the legacy behavior.

## Behavior
- `param_dtype` unset -> derived from the fp16/bf16 enabled flag (legacy
param behavior).
- `buffer_dtype` unset -> buffers keep their loaded dtype (e.g. fp32
`inv_freq`).
- `buffer_dtype` set -> buffers force-cast (legacy blanket-cast parity).

## Test plan
- [ ] `param_dtype=bf16`, `buffer_dtype` unset -> params bf16,
`inv_freq` stays fp32.
- [ ] `buffer_dtype=bf16` -> buffers downcast (legacy parity).
- [ ] bf16/fp16 run with neither key set behaves as before except fp32
buffers preserved.
- [ ] 8B / 32B ZeRO-3 long-context run -> grad_norm tracks the FSDP2
reference.

Made with [Cursor](https://cursor.com)

---------

Signed-off-by: Olatunji Ruwase <tunji.ruwase@snowflake.com>
Signed-off-by: Stas Bekman <stas@stason.org>
Co-authored-by: Cursor <cursoragent@cursor.com>
Co-authored-by: Olatunji Ruwase <tjruwase@gmail.com>
Co-authored-by: Stas Bekman <stas00@users.noreply.github.com>
Co-authored-by: Stas Bekman <stas@stason.org>

…ffers) (deepspeedai#8066)

## Summary
- Add `data_types.param_dtype` and `data_types.buffer_dtype` (both
default `None`), mirroring FSDP `MixedPrecisionPolicy`.
- Replace the blanket `module.half()` / `module.bfloat16()` in
`_configure_distributed_model` with a targeted cast: parameters go to
`param_dtype`; floating buffers keep their loaded dtype unless
`buffer_dtype` is explicitly set.

## Motivation
The blanket cast downcasts every floating buffer, including the rotary
`inv_freq` buffer that HF/FSDP2 keep in fp32. On long contexts the bf16
`inv_freq` loses precision, RoPE angles drift, and logits/grads diverge
from the FSDP2 reference. Preserving fp32 buffers by default fixes this;
`buffer_dtype` is the escape hatch to reproduce the legacy behavior.

## Behavior
- `param_dtype` unset -> derived from the fp16/bf16 enabled flag (legacy
param behavior).
- `buffer_dtype` unset -> buffers keep their loaded dtype (e.g. fp32
`inv_freq`).
- `buffer_dtype` set -> buffers force-cast (legacy blanket-cast parity).

## Test plan
- [ ] `param_dtype=bf16`, `buffer_dtype` unset -> params bf16,
`inv_freq` stays fp32.
- [ ] `buffer_dtype=bf16` -> buffers downcast (legacy parity).
- [ ] bf16/fp16 run with neither key set behaves as before except fp32
buffers preserved.
- [ ] 8B / 32B ZeRO-3 long-context run -> grad_norm tracks the FSDP2
reference.

Made with [Cursor](https://cursor.com)

---------

Signed-off-by: Olatunji Ruwase <tunji.ruwase@snowflake.com>
Signed-off-by: Stas Bekman <stas@stason.org>
Co-authored-by: Cursor <cursoragent@cursor.com>
Co-authored-by: Olatunji Ruwase <tjruwase@gmail.com>
Co-authored-by: Stas Bekman <stas00@users.noreply.github.com>
Co-authored-by: Stas Bekman <stas@stason.org>
Signed-off-by: nathon-lee <leejianwoo@gmail.com>

…ms_coalesced (deepspeedai#8073)

This PR fixes the `_allgather_params_coalesced` method in
`partition_parameters.py`. The change ensures that each `flat_tensor` is
created with the correct data type by referencing the corresponding
parameter in `param_list`, rather than always using the first
parameter's data type.

Fix deepspeedai#8072.

### Problem

`_allgather_params_coalesced` allocates all output buffers using the
dtype of the first parameter in param_list:

```python
# before
for psize in partition_sizes:
    flat_tensor = torch.empty(tensor_size, dtype=param_list[0].ds_tensor.dtype, ...)
```

This assumed every persistent parameter shares the same dtype. The
assumption was incidentally maintained before 0.19.2 because
`_configure_distributed_model` called `module.bfloat16()`
unconditionally, normalising all persistent parameters (including PEFT
LoRA adapters) to a uniform dtype.

PR deepspeedai#8066 "Mixed-precision:
per-policy param/buffer dtype cast (preserve fp32 buffers)" (commit
b919284) correctly stopped casting ZeRO-Init model params, but exposed
the latent bug: PEFT's default `autocast_adapter_dtype=True` keeps LoRA
adapters in fp32 even when the base model is bf16.
`persistent_parameters` therefore ends up with mixed dtypes (bf16
base-model params + fp32 LoRA params), and the mismatch between a bf16
output buffer and a fp32 input tensor raises:
> TypeError: output tensor must have the same type as input tensor

- Reported in: deepspeedai#8072
- Downstream tracking: huggingface/trl#6089

### Solution

Allocate each output buffer with the dtype of its own parameter:

```python
# after
for i, psize in enumerate(partition_sizes):
    flat_tensor = torch.empty(tensor_size, dtype=param_list[i].ds_tensor.dtype, ...)
```

This removes the shared-dtype assumption at the source rather than
relying on upstream callers to normalise dtypes before calling
`_allgather_params_coalesced`.

### Changes

- Corrected tensor data type selection in `_allgather_params_coalesced`
to use the data type of each parameter in `param_list`, ensuring proper
handling of mixed data types.

Signed-off-by: Albert Villanova del Moral <8515462+albertvillanova@users.noreply.github.com>
Signed-off-by: nathon-lee <leejianwoo@gmail.com>

…ffers) (deepspeedai#8066)

## Summary
- Add `data_types.param_dtype` and `data_types.buffer_dtype` (both
default `None`), mirroring FSDP `MixedPrecisionPolicy`.
- Replace the blanket `module.half()` / `module.bfloat16()` in
`_configure_distributed_model` with a targeted cast: parameters go to
`param_dtype`; floating buffers keep their loaded dtype unless
`buffer_dtype` is explicitly set.

## Motivation
The blanket cast downcasts every floating buffer, including the rotary
`inv_freq` buffer that HF/FSDP2 keep in fp32. On long contexts the bf16
`inv_freq` loses precision, RoPE angles drift, and logits/grads diverge
from the FSDP2 reference. Preserving fp32 buffers by default fixes this;
`buffer_dtype` is the escape hatch to reproduce the legacy behavior.

## Behavior
- `param_dtype` unset -> derived from the fp16/bf16 enabled flag (legacy
param behavior).
- `buffer_dtype` unset -> buffers keep their loaded dtype (e.g. fp32
`inv_freq`).
- `buffer_dtype` set -> buffers force-cast (legacy blanket-cast parity).

## Test plan
- [ ] `param_dtype=bf16`, `buffer_dtype` unset -> params bf16,
`inv_freq` stays fp32.
- [ ] `buffer_dtype=bf16` -> buffers downcast (legacy parity).
- [ ] bf16/fp16 run with neither key set behaves as before except fp32
buffers preserved.
- [ ] 8B / 32B ZeRO-3 long-context run -> grad_norm tracks the FSDP2
reference.

Made with [Cursor](https://cursor.com)

---------

Signed-off-by: Olatunji Ruwase <tunji.ruwase@snowflake.com>
Signed-off-by: Stas Bekman <stas@stason.org>
Co-authored-by: Cursor <cursoragent@cursor.com>
Co-authored-by: Olatunji Ruwase <tjruwase@gmail.com>
Co-authored-by: Stas Bekman <stas00@users.noreply.github.com>
Co-authored-by: Stas Bekman <stas@stason.org>

…ms_coalesced (deepspeedai#8073)

This PR fixes the `_allgather_params_coalesced` method in
`partition_parameters.py`. The change ensures that each `flat_tensor` is
created with the correct data type by referencing the corresponding
parameter in `param_list`, rather than always using the first
parameter's data type.

Fix deepspeedai#8072.

### Problem

`_allgather_params_coalesced` allocates all output buffers using the
dtype of the first parameter in param_list:

```python
# before
for psize in partition_sizes:
    flat_tensor = torch.empty(tensor_size, dtype=param_list[0].ds_tensor.dtype, ...)
```

This assumed every persistent parameter shares the same dtype. The
assumption was incidentally maintained before 0.19.2 because
`_configure_distributed_model` called `module.bfloat16()`
unconditionally, normalising all persistent parameters (including PEFT
LoRA adapters) to a uniform dtype.

PR deepspeedai#8066 "Mixed-precision:
per-policy param/buffer dtype cast (preserve fp32 buffers)" (commit
b919284) correctly stopped casting ZeRO-Init model params, but exposed
the latent bug: PEFT's default `autocast_adapter_dtype=True` keeps LoRA
adapters in fp32 even when the base model is bf16.
`persistent_parameters` therefore ends up with mixed dtypes (bf16
base-model params + fp32 LoRA params), and the mismatch between a bf16
output buffer and a fp32 input tensor raises:
> TypeError: output tensor must have the same type as input tensor

- Reported in: deepspeedai#8072
- Downstream tracking: huggingface/trl#6089

### Solution

Allocate each output buffer with the dtype of its own parameter:

```python
# after
for i, psize in enumerate(partition_sizes):
    flat_tensor = torch.empty(tensor_size, dtype=param_list[i].ds_tensor.dtype, ...)
```

This removes the shared-dtype assumption at the source rather than
relying on upstream callers to normalise dtypes before calling
`_allgather_params_coalesced`.

### Changes

- Corrected tensor data type selection in `_allgather_params_coalesced`
to use the data type of each parameter in `param_list`, ensuring proper
handling of mixed data types.

Signed-off-by: Albert Villanova del Moral <8515462+albertvillanova@users.noreply.github.com>