Summary
Under FSDP1 with use_orig_params=False, the wrapping policy decides whether AdamW8bit keeps a parameter's optimizer state in FP32 or 8-bit. The optimizer applies its size threshold to the local flattened tensor, so grouping several small parameters into one flat parameter can switch their state from FP32 to 8-bit.
Model, gradients, and optimizer hyperparameters stay identical while the resulting updates change.
System Info
Environment: bitsandbytes 0.50.2, commit 08a9956b, PyTorch 2.13.0+cu130, two ranks, FSDP1 FULL_SHARD.
The results below are from the pinned version. Source inspection on 2026-09-10 found the same calculation in 83364904; the full experiment has not been rerun on that commit.
The linked self-contained example was also run on Ubuntu 26.04 LTS (x86_64), Python 3.11.15, PyTorch 2.13.0+cpu, and bitsandbytes 0.50.2, with two local CPU/Gloo processes. It uses FSDP1, FULL_SHARD, and use_orig_params=False. The current-source checkout supplies the example; its native optimizer run uses the installed bitsandbytes 0.50.2 library.
Reproduction
The fixture holds four original parameters of 2,048 elements each and sets min_8bit_size=4096. It compares two wrapping policies while keeping the initial values, original parameter order, full FP32 gradients, optimizer settings, and update boundaries fixed:
| Layout of the four small parameters |
Local size seen by the optimizer |
State chosen by AdamW8bit |
| One FSDP unit per parameter |
1,024 per tensor |
FP32 |
| All four in one FSDP unit |
4,096 combined |
8-bit |
A separate 8,192-element parameter stays in its own FSDP unit and uses 8-bit state under both layouts.
The test also compares a fixed projection of the parameter update after the second step: the relative L2 difference is 2.99e-4 before the reference repair and 0 after it. The recorded state-eligibility comparison also becomes exact. Full-precision AdamW shows zero difference in the same checks. Reproduced on a second host.
A self-contained native/reference comparison is available in the example added by PR #2081. From a checkout of that PR, run:
torchrun --standalone --nproc-per-node=2 examples/fsdp1_original_parameter_state.py --device cpu
This runs both wrapping layouts and both optimizer paths. In the CPU run, the native full update-vector maximum difference was 6.15e-5; the reference difference was zero. This full-vector metric is separate from the pinned CUDA projection reported above.
Expected behavior
If original-parameter optimizer semantics are supported, changing only FSDP1 wrapping should preserve state precision and quantization boundaries for the same original parameters. If eligibility is intentionally defined by the local flat shard instead, that limitation should be documented explicitly.
Root cause
Optimizer2State.init_state applies the size threshold to the tensor FSDP hands it. That tensor is the local shard of a flattened group, not an original model parameter.
Both the precision decision and the quantization-block layout therefore inherit the wrapping policy. In this fixture each original parameter is below the threshold, but the combined local flat parameter reaches it exactly.
Proposed fix
Base the precision decision and the quantization-block boundaries on the original parameters' identities and sizes, and carry that mapping through flattening and sharding.
The tested reference repair keeps the coarse wrapping, gathers each original parameter's gradient, runs AdamW8bit with state kept in per-original-parameter shadow tensors, and scatters the updates back. The large control parameter keeps its 8-bit state.
This demonstrates the behavior to preserve rather than a mergeable implementation. The shadow state is replicated, and the metadata interface and an efficient sharded implementation are still open.
Expected behavior: preserve per-original-parameter state eligibility when only the FSDP1 wrapping policy changes, or document that min_8bit_size is intentionally a threshold on local flat shards. The reference repair above is a proof of concept for discussing that choice.
Summary
Under FSDP1 with
use_orig_params=False, the wrapping policy decides whether AdamW8bit keeps a parameter's optimizer state in FP32 or 8-bit. The optimizer applies its size threshold to the local flattened tensor, so grouping several small parameters into one flat parameter can switch their state from FP32 to 8-bit.Model, gradients, and optimizer hyperparameters stay identical while the resulting updates change.
System Info
Environment: bitsandbytes 0.50.2, commit 08a9956b, PyTorch 2.13.0+cu130, two ranks, FSDP1
FULL_SHARD.The results below are from the pinned version. Source inspection on 2026-09-10 found the same calculation in 83364904; the full experiment has not been rerun on that commit.
The linked self-contained example was also run on Ubuntu 26.04 LTS (x86_64), Python 3.11.15, PyTorch 2.13.0+cpu, and bitsandbytes 0.50.2, with two local CPU/Gloo processes. It uses
FSDP1,FULL_SHARD, anduse_orig_params=False. The current-source checkout supplies the example; its native optimizer run uses the installed bitsandbytes 0.50.2 library.Reproduction
The fixture holds four original parameters of 2,048 elements each and sets
min_8bit_size=4096. It compares two wrapping policies while keeping the initial values, original parameter order, full FP32 gradients, optimizer settings, and update boundaries fixed:A separate 8,192-element parameter stays in its own FSDP unit and uses 8-bit state under both layouts.
The test also compares a fixed projection of the parameter update after the second step: the relative L2 difference is
2.99e-4before the reference repair and0after it. The recorded state-eligibility comparison also becomes exact. Full-precision AdamW shows zero difference in the same checks. Reproduced on a second host.A self-contained native/reference comparison is available in the example added by PR #2081. From a checkout of that PR, run:
This runs both wrapping layouts and both optimizer paths. In the CPU run, the native full update-vector maximum difference was
6.15e-5; the reference difference was zero. This full-vector metric is separate from the pinned CUDA projection reported above.Expected behavior
If original-parameter optimizer semantics are supported, changing only FSDP1 wrapping should preserve state precision and quantization boundaries for the same original parameters. If eligibility is intentionally defined by the local flat shard instead, that limitation should be documented explicitly.
Root cause
Optimizer2State.init_stateapplies the size threshold to the tensor FSDP hands it. That tensor is the local shard of a flattened group, not an original model parameter.Both the precision decision and the quantization-block layout therefore inherit the wrapping policy. In this fixture each original parameter is below the threshold, but the combined local flat parameter reaches it exactly.
Proposed fix
Base the precision decision and the quantization-block boundaries on the original parameters' identities and sizes, and carry that mapping through flattening and sharding.
The tested reference repair keeps the coarse wrapping, gathers each original parameter's gradient, runs AdamW8bit with state kept in per-original-parameter shadow tensors, and scatters the updates back. The large control parameter keeps its 8-bit state.
This demonstrates the behavior to preserve rather than a mergeable implementation. The shadow state is replicated, and the metadata interface and an efficient sharded implementation are still open.
Expected behavior: preserve per-original-parameter state eligibility when only the FSDP1 wrapping policy changes, or document that
min_8bit_sizeis intentionally a threshold on local flat shards. The reference repair above is a proof of concept for discussing that choice.