fix(models): audit fp32 protected tensors#2598
Open
yuhezhang-ai wants to merge 8 commits into
Open
Conversation
akoumpa
added
the
r0.5.0
Contributor
|
/nvskills-ci |
Contributor
Author
|
/ok to test 704a39c |
Signed-off-by: Yuhe Zhang <yuhez@nvidia.com>
Signed-off-by: Yuhe Zhang <yuhez@nvidia.com>
Signed-off-by: Yuhe Zhang <yuhez@nvidia.com>
Signed-off-by: Yuhe Zhang <yuhez@nvidia.com>
Signed-off-by: Yuhe Zhang <yuhez@nvidia.com>
Signed-off-by: Yuhe Zhang <yuhez@nvidia.com>
Signed-off-by: Yuhe Zhang <yuhez@nvidia.com>
yuhezhang-ai
force-pushed
the
yuhez/fix/fp32-protected-tensor-audit
branch
from
b2ec859 to
b035704
Compare
Signed-off-by: Yuhe Zhang <yuhez@nvidia.com>
Contributor
Author
|
/ok to test 9e1db7b |
Contributor
Author
|
/nvskills-ci |
Contributor
Author
|
/claude review |
![claude[bot]](https://avatars.githubusercontent.com/in/1236702?s=60&v=4){kind=small}
| assert original_fn is not None, "apply_fsdp2_sharding_recursively not found in module globals" | ||
|
|
||
| def _fsdp_by_dtype(module, mesh, mp_policy, offload_policy=None, *args, **kwargs): | ||
| def _fsdp_by_dtype( |
Contributor
There was a problem hiding this comment.
LGTM — clean PR. The fp32 holder pattern, reshard threading, and state-dict adapter routing all look correct. Good test coverage across models and distributed paths.
This file contains hidden or bidirectional Unicode text that may be interpreted or compiled differently than what appears below. To review, open the file in an editor that reveals hidden Unicode characters.
Learn more about bidirectional Unicode characters
Sign up for free
to join this conversation on GitHub.
Already have an account?
Sign in to comment
2 participants
Add this suggestion to a batch that can be applied as a single commit.This suggestion is invalid because no changes were made to the code.Suggestions cannot be applied while the pull request is closed.Suggestions cannot be applied while viewing a subset of changes.Only one suggestion per line can be applied in a batch.Add this suggestion to a batch that can be applied as a single commit.Applying suggestions on deleted lines is not supported.You must change the existing code in this line in order to create a valid suggestion.Outdated suggestions cannot be applied.This suggestion has been applied or marked resolved.Suggestions cannot be applied from pending reviews.Suggestions cannot be applied on multi-line comments.Suggestions cannot be applied while the pull request is queued to merge.Suggestion cannot be applied right now. Please check back later.
Summary
A_log/dt_bias/Dinto fp32 holder modules, with HF-compatible state-dict adapter routing for the public checkpoint keys.Why
model.torch_dtype=bf16and broad dtype casts can round small tensors that are part of a model's numerical/checkpoint contract. Trainable fp32 params that live directly on mixed modules also cannot be isolated cleanly by FSDP, so this follows the_fp32_paramsholder pattern used by the stacked Qwen GDN branch.The real 2-GPU FSDP smoke exposed one more holder-specific issue: a callable holder that returns its parameter can hand the parent module a sharded DTensor. The fix keeps storage sharded/fp32, but returns the full fp32 tensor value to the caller.
Notes
This PR is stacked on
yuhez/fix/qwen-gdn-fp32-precision. After that branch lands, this should be rebased ontomain.Validation
uv run --no-sync pytest tests/unit_tests/distributed/test_parallelizer_utils.py tests/unit_tests/distributed/test_fp32_compute_contract.py -q(25 passed)30 passed)uv run --no-sync pytest tests/unit_tests/models/nemotron_v3/test_nemotron_v3_state_dict_adapter.py tests/unit_tests/distributed/test_fp32_compute_contract.py -q(30 passed)9 passed)5 passed, 1 skipped)11 passed)fp32-holder-fsdp-smoke-c938fb8-r3, job12863187, exit0:0ruff checkon touched filesgit diff --check