[Pytorch][Bug] DCP Checkpoint Loading Fixes for FSDP2 with QuantizedModelInit#2974
[Pytorch][Bug] DCP Checkpoint Loading Fixes for FSDP2 with QuantizedModelInit#2974vthumbe1503 wants to merge 9 commits into
Conversation
vthumbe1503
changed the title
Greptile SummaryThis PR fixes DCP sync and async checkpoint loading for MXFP8/NVFP4 and async loading for all quantization recipes when using FSDP2 with
Confidence Score: 4/5Safe to merge once the None-data guard is added to _make_float8_tensor_in_reduce_ex. In transformer_engine/pytorch/tensor/float8_tensor.py — specifically _make_float8_tensor_in_reduce_ex. Important Files Changed
Flowchart%%{init: {'theme': 'neutral'}}%%
flowchart TD
A[DCP Checkpoint Load] --> B{Sync or Async?}
B -->|Sync| C[FSDP2 compares untyped_storage]
C --> D["QuantizedTensor.untyped_storage()\nreturns empty 0-byte storage"]
D --> E[Storage != BF16 sharded param\n→ not same_tensor → copy path taken]
E --> F[copy_ dispatches to QuantizedTensor\n→ quantize_ or dequantize]
B -->|Async| G[DCP staging serializes\nQuantizedTensor to CPU]
G --> H["__reduce_ex__ called\n(module-level _make_*_in_reduce_ex)"]
H --> I[fp8/fp4_dtype serialized as int\nQuantizer.dtype serialized as int]
I --> J["torch.load(weights_only=True)\nuses add_safe_globals whitelist"]
J --> K[Reconstruct QuantizedTensor\nvia _make_*_in_reduce_ex]
G --> L{Float8Tensor}
G --> M{MXFP8Tensor / NVFP4Tensor}
L --> N[Serialize FP8 buffers directly\nno CPU-dequantize fallback]
M --> O["dequantize() on CPU:\nbounce to CUDA → tex.dequantize\n→ move result back to CPU"]
K --> P[QuantizedTensor on CUDA\nready for training]
O --> P
Reviews (6): Last reviewed commit: "[pre-commit.ci] auto fixes from pre-comm..." | Re-trigger Greptile |
![greptile-apps[bot]](https://avatars.githubusercontent.com/in/867647?s=60&v=4){kind=small}
| def untyped_storage(self) -> torch.UntypedStorage: | ||
| """Return an empty UntypedStorage on the tensor's device. | ||
|
|
||
| ``QuantizedTensor`` is a ``_make_wrapper_subclass`` and has no real | ||
| backing storage of its own; the actual bytes live in the inner | ||
| buffers (e.g. ``_rowwise_data`` / ``_columnwise_data``) which are | ||
| an implementation detail of the quantization scheme. Need to define | ||
| this method to avoid DCP staging errors with FSDP2. | ||
| """ | ||
| return torch.UntypedStorage(0, device=self.device) |
There was a problem hiding this comment.
QuantizedTensor.untyped_storage() now returns a freshly allocated zero-byte storage every call. Code in module/_common.py:128 compares tensors[0].untyped_storage().nbytes() against expected size to decide between a no-op view and an out-of-place torch.cat. With 0 bytes returned, that condition is always true, silently disabling the in-place fast path for any QuantizedTensor through ConcatMerge.forward. More critically, utils.py:403-412 in SplitAlongDim.backward uses data_ptr() for noop detection — if all zero-size CUDA allocations return data_ptr() == 0, every QuantizedTensor pair incorrectly appears co-located, setting noop_ok = True and crashing on ret.set_() against a 0-byte storage.
There was a problem hiding this comment.
The correct behavior for these functions is to fall back to the slow path for QuantizedTensor s, unless it has a dedicated implementation to handle quantized data.
There was a problem hiding this comment.
Yeah, while I don't think we use QuantizedTensors in the SplitAlongDim ever, the concat sounds plausible to be hit.
There was a problem hiding this comment.
Need to resolve this comment after going thoroughly over noop_cat consequences on Quantizedtensors
| elif recipe_name == "NVFP4BlockScaling": | ||
| # NVFP4 DCP load goes through a dequant + quant, so neec to relax tolerances | ||
| torch.testing.assert_close( | ||
| loaded_output, | ||
| ref_output, | ||
| rtol=0.125, | ||
| atol=0.25, | ||
| msg=lambda x: f"NVFP4BlockScaling: Fresh model loaded from DCP checkpoint produces different output: {x}", | ||
| ) |
There was a problem hiding this comment.
| elif recipe_name == "NVFP4BlockScaling": | |
| # NVFP4 DCP load goes through a dequant + quant, so neec to relax tolerances | |
| torch.testing.assert_close( | |
| loaded_output, | |
| ref_output, | |
| rtol=0.125, | |
| atol=0.25, | |
| msg=lambda x: f"NVFP4BlockScaling: Fresh model loaded from DCP checkpoint produces different output: {x}", | |
| ) | |
| elif recipe_name == "NVFP4BlockScaling": | |
| # NVFP4 DCP load goes through a dequant + quant, so need to relax tolerances | |
| torch.testing.assert_close( | |
| loaded_output, | |
| ref_output, | |
| rtol=0.125, | |
| atol=0.25, | |
| msg=lambda x: f"NVFP4BlockScaling: Fresh model loaded from DCP checkpoint produces different output: {x}", | |
| ) |
| elif recipe_name == "NVFP4BlockScaling": | ||
| # NVFP4 DCP load goes through a dequant + quant, so neec to relax tolerances | ||
| torch.testing.assert_close( | ||
| out2, | ||
| out1, | ||
| rtol=0.125, | ||
| atol=0.25, | ||
| msg=lambda x: f"NVFP4BlockScaling: Training step after DCP load produces different output: {x}", | ||
| ) |
There was a problem hiding this comment.
| elif recipe_name == "NVFP4BlockScaling": | |
| # NVFP4 DCP load goes through a dequant + quant, so neec to relax tolerances | |
| torch.testing.assert_close( | |
| out2, | |
| out1, | |
| rtol=0.125, | |
| atol=0.25, | |
| msg=lambda x: f"NVFP4BlockScaling: Training step after DCP load produces different output: {x}", | |
| ) | |
| elif recipe_name == "NVFP4BlockScaling": | |
| # NVFP4 DCP load goes through a dequant + quant, so need to relax tolerances | |
| torch.testing.assert_close( | |
| out2, | |
| out1, | |
| rtol=0.125, | |
| atol=0.25, | |
| msg=lambda x: f"NVFP4BlockScaling: Training step after DCP load produces different output: {x}", | |
| ) |
vthumbe1503
changed the title
|
/te-ci L1 pytorch |
| msg=lambda x: f"Fresh model loaded from DCP checkpoint produces different output: {x}", | ||
| ) | ||
| elif recipe_name == "NVFP4BlockScaling": | ||
| # NVFP4 DCP load goes through a dequant + quant, so neec to relax tolerances |
There was a problem hiding this comment.
Why do we need dequant + quant here?
| # When a CPU copy of a quantized tensor is requested (e.g. by | ||
| # torch DCP staging via ``x.new_empty(..., device="cpu")``), we | ||
| # save the high-precision values in a plain CPU dense tensor. | ||
| # For the DCP load path, we will re-quantize the high-precision values. |
There was a problem hiding this comment.
This fix seems ad hoc to me. It's not obvious why qtensor.new_empty(..., device="cpu") returns a quantized tensor while qtensor.new_empty(..., device="cuda") returns a plain tensor. I wonder if it would be cleaner to just return a plain tensor in all cases. Thoughts:
- It's uncomfortable how
new_emptyandempty_likewould have different behavior. I suppose we could interpretempty_likeas "make a tensor that matches the input" andnew_emptyas "call torch.empty with defaults taken from input", but that would be a private interpretation that no one else follows. - Would this affect FSDP or CPU offloading?
- Given the weirdness, would it be worthwhile raising a warning if
new_emptyis called outside of DCP?
| # torch DCP staging via ``x.new_empty(..., device="cpu")``), we | ||
| # save the high-precision values in a plain CPU dense tensor. | ||
| # For the DCP load path, we will re-quantize the high-precision values. | ||
| target_size = torch.Size(size) if len(size) > 0 else tensor.size() |
There was a problem hiding this comment.
An empty size is valid and it corresponds to a tensor with 1 entry (for the same reason 2^0=1).
>>> import torch
>>> x = torch.ones(123).new_empty([])
>>> print(x.numel())
1
| target_size = torch.Size(size) if len(size) > 0 else tensor.size() | |
| target_size = size |
| def untyped_storage(self) -> torch.UntypedStorage: | ||
| """Return an empty UntypedStorage on the tensor's device. | ||
|
|
||
| ``QuantizedTensor`` is a ``_make_wrapper_subclass`` and has no real | ||
| backing storage of its own; the actual bytes live in the inner | ||
| buffers (e.g. ``_rowwise_data`` / ``_columnwise_data``) which are | ||
| an implementation detail of the quantization scheme. Need to define | ||
| this method to avoid DCP staging errors with FSDP2. | ||
| """ | ||
| return torch.UntypedStorage(0, device=self.device) |
There was a problem hiding this comment.
The correct behavior for these functions is to fall back to the slow path for QuantizedTensor s, unless it has a dedicated implementation to handle quantized data.
| # differences vs the manual dequantize-then-allgather path. | ||
| if isinstance(param, NVFP4Tensor): | ||
| tols = dict(atol=5e-4, rtol=5e-3) | ||
| tols = dict(atol=0.125, rtol=0.25) |
There was a problem hiding this comment.
Why are the tolerances so much bigger? Is it also due to the dequant+quant path? If so, the comment above is no longer relevant and should be replaced with a better one (but I would still like an explanation why we cannot just load the nvfp4 values from the checkpoint).
| # When a CPU copy of a quantized tensor is requested (e.g. by | ||
| # torch DCP staging via ``x.new_empty(..., device="cpu")``), we | ||
| # save the high-precision values in a plain CPU dense tensor. | ||
| # For the DCP load path, we will re-quantize the high-precision values. |
There was a problem hiding this comment.
Ok, I see now why you want to dequantize. I don't think this is needed though - we should be able to create the QuantlizedTensor on the CPU and save it, no? I remember that the CPU offloading of the activations faced similar problem and already had to support some CPU ops on the QuantizedTensor anyway.
Signed-off-by: Varun Thumbe <vthumbe@nvidia.com>
vthumbe1503
force-pushed
the
fsdp2_dcp_laod_fix
branch
from
3589ffa to
4197bee
Compare
for more information, see https://pre-commit.ci
|
/te-ci L1 pytorch |
| # allow-listed for ``torch.load(weights_only=True)`` (used | ||
| # internally by DCP async-staging) to accept the stream. | ||
| _TE_DType, | ||
| getattr, |
There was a problem hiding this comment.
getattr whitelisted for weights_only=True loading
getattr is registered as a safe global at module import time. add_safe_globals is process-wide in PyTorch, so any torch.load(…, weights_only=True) call made anywhere in a session that has imported transformer_engine.pytorch — including checkpoint loads for entirely different models — now has getattr available to the pickle stream. A malicious checkpoint loaded elsewhere could use getattr to access sensitive attributes of any already-constructed object reachable from the whitelisted globals (e.g. getattr(Float8Quantizer_instance, 'amax_reduction_group') to obtain a process group, or to build callable gadget chains). The weights_only=True flag is specifically a defence against untrusted pickle payloads; adding a general-purpose reflective accessor defeats that defence.
A targeted alternative: serialize _fp8_dtype as its integer value (int(self._fp8_dtype)) and reconstruct it in _make_in_reduce_ex via TE_DType(int_value), then add TE_DType to safe globals instead of getattr. This preserves the weights_only invariant without whitelisting a reflective accessor.
Signed-off-by: Varun Thumbe <vthumbe@nvidia.com>
Signed-off-by: Varun Thumbe <vthumbe@nvidia.com>
Signed-off-by: vthumbe1503 <vthumbe@nvidia.com>
for more information, see https://pre-commit.ci
Signed-off-by: Varun Thumbe <vthumbe@nvidia.com>
for more information, see https://pre-commit.ci
3 participants
Description
Fixes DCP Sync and Async checkpoint loading for MXFP8/NVFP4.
Fixes DCP Async checkpoint loading for all Quantization recipes
Fixes NVFP4 allgather + dequant numerical errors for fsdp2. Turns out this was due to us not setting the amax reduction group as the fsdp group in the quantizer
Fixes # (issue)
Type of change
Changes
Please list the changes introduced in this PR:
DCP Sync Checkpoint loading
DCP Async Checkpointing
Checklist: