FSDP2 calibration with hf_ptq.py [1/2] - #1563
Merged
Merged
Conversation
Contributor
|
Note Reviews pausedIt looks like this branch is under active development. To avoid overwhelming you with review comments due to an influx of new commits, CodeRabbit has automatically paused this review. You can configure this behavior by changing the Use the following commands to manage reviews:
Use the checkboxes below for quick actions:
📝 WalkthroughWalkthroughAdds FSDP2-aware distributed loading, sharding, and checkpoint export paths into hf_ptq; implements parallel safetensors subset loading, fsdp2 wrap/sharding utilities, DTensor gather/writeback bookkeeping changes, rank-aware export synchronization, README and SLURM docs updates, and regression/unit tests. ChangesFSDP2 PTQ Consolidation
Estimated code review effort🎯 4 (Complex) | ⏱️ ~45 minutes Suggested reviewers
🚥 Pre-merge checks | ✅ 5 | ❌ 1❌ Failed checks (1 warning)
✅ Passed checks (5 passed)
✨ Finishing Touches🧪 Generate unit tests (beta)
Comment |
sugunav14
marked this pull request as draft
May 28, 2026 23:23
|
Auto-sync is disabled for draft pull requests in this repository. Workflows must be run manually. Contributors can view more details about this message here. |
![coderabbitai[bot]](https://avatars.githubusercontent.com/in/347564?s=60&v=4){kind=small}
Contributor
There was a problem hiding this comment.
Actionable comments posted: 4
🤖 Prompt for all review comments with AI agents
Verify each finding against current code. Fix only still-valid issues, skip the
rest with a brief reason, keep changes minimal, and validate.
Inline comments:
In `@examples/llm_ptq/example_utils.py`:
- Around line 68-78: The else branch for when getattr(args, "use_fsdp2", False)
is false is overwriting args.device with None which breaks non-FSDP device
handling (e.g., --device cpu); change the else branch in that block so it sets
args.rank = 0, args.world_size = 1 and args.is_main = True but does NOT modify
args.device (leave existing args.device intact or only set it if undefined), so
downstream calls like get_model() receive the requested device; keep reference
to getattr(args, "use_fsdp2", False), args.device, args.rank, args.world_size,
args.is_main and get_model() to locate the code to change.
- Around line 183-201: Modify create_fsdp2_calibration_loop to accept an is_main
(default False) parameter and thread it into the inner calibrate closure; inside
calibrate only wrap the dataloader with tqdm(desc="Calibrating") when is_main is
True and otherwise iterate the dataloader without tqdm to avoid per-rank
progress bars, and use print_rank_0 to emit any high-level start/finish messages
on rank 0 if needed; update references to the calibrate closure accordingly so
callers can pass is_main=True on rank 0.
- Around line 115-180: The FSDP2 load path in load_and_prepare_fsdp2_model
ignores the --attn_implementation setting, so pass the attn implementation
through to HuggingFace calls and the caller: include args.attn_implementation
(when args is not None) in AutoConfig.from_pretrained,
AutoModelForCausalLM.from_pretrained and AutoModelForCausalLM.from_config by
forwarding it as the attention backend argument expected by HF (e.g.,
attn_implementation or equivalent kwarg), and ensure examples/llm_ptq/hf_ptq.py
passes args.attn_implementation into load_and_prepare_fsdp2_model when invoking
it; update references in this function (hf_config creation, from_pretrained
call, from_config call) and the call site to propagate the flag unchanged.
In `@examples/llm_ptq/hf_ptq.py`:
- Around line 1430-1440: Add a fail-fast validation after argument parsing to
reject the incompatible combination of --use_fsdp2 and auto-quantize: detect
when args.use_fsdp2 is true and the auto-quantize flag/setting (the CLI flag
that triggers mtq.auto_quantize(), e.g., args.auto_quantize_bits or equivalent)
is present/non-zero, log/raise a clear error and exit non-zero before calling
fsdp2_shard() or mtq.auto_quantize(); reference the parser option "--use_fsdp2",
the fsdp2_shard() call and mtq.auto_quantize() so the check runs immediately
after parsing and prevents entering the frozen-parameter path.
🪄 Autofix (Beta)
Fix all unresolved CodeRabbit comments on this PR:
- Push a commit to this branch (recommended)
- Create a new PR with the fixes
ℹ️ Review info
⚙️ Run configuration
Configuration used: Path: .coderabbit.yaml
Review profile: CHILL
Plan: Enterprise
Run ID: 250df2b3-92b6-427a-ad39-e2fb1869f878
📒 Files selected for processing (7)
examples/llm_ptq/README.mdexamples/llm_ptq/example_utils.pyexamples/llm_ptq/fsdp2.yamlexamples/llm_ptq/hf_ptq.pyexamples/llm_ptq/multinode_ptq.pymodelopt/torch/quantization/utils/core_utils.pymodelopt/torch/utils/distributed.py
💤 Files with no reviewable changes (2)
- examples/llm_ptq/fsdp2.yaml
- examples/llm_ptq/multinode_ptq.py
Codecov Report❌ Patch coverage is Additional details and impacted files@@ Coverage Diff @@
## main #1563 +/- ##
==========================================
+ Coverage 75.79% 77.90% +2.11%
==========================================
Files 518 519 +1
Lines 58658 58896 +238
==========================================
+ Hits 44459 45885 +1426
+ Misses 14199 13011 -1188
Flags with carried forward coverage won't be shown. Click here to find out more. ☔ View full report in Codecov by Harness. 🚀 New features to boost your workflow:
|
sugunav14
marked this pull request as ready for review
May 29, 2026 16:47
Contributor
There was a problem hiding this comment.
Caution
Some comments are outside the diff and can’t be posted inline due to platform limitations.
⚠️ Outside diff range comments (1)
examples/llm_ptq/example_utils.py (1)
97-112:⚠️ Potential issue | 🟠 Major | ⚡ Quick winReject encoder-decoder configs before entering the FSDP2 loader.
load_and_prepare_fsdp2_model()always usesAutoModelForCausalLM, so families like T5/BART/Whisper currently slip through validation and fail later during model construction instead of with the intended earlyNotImplementedError.Suggested fix
if getattr(config, "quantization_config", None) is not None: issues.append("pack-quantized / compressed-tensors checkpoints") + if getattr(config, "is_encoder_decoder", False): + issues.append("encoder-decoder models (FSDP2 v1 only supports causal LMs)") if getattr(args, "specdec_offline_dataset", None) is not None: issues.append("speculative decoding (--specdec_offline_dataset)")🤖 Prompt for AI Agents
Verify each finding against current code. Fix only still-valid issues, skip the rest with a brief reason, keep changes minimal, and validate. In `@examples/llm_ptq/example_utils.py` around lines 97 - 112, The validation block that collects issues must also reject encoder-decoder configs so load_and_prepare_fsdp2_model() doesn't call AutoModelForCausalLM on incompatible families; detect encoder-decoder models by checking config.is_encoder_decoder (or config.model_type in known encoder-decoder types) and append a descriptive message to issues (e.g., "encoder-decoder models (T5/BART/Whisper)"), keeping the existing checks (vila, is_nemotron_vl, _is_multimodal_config, quantization_config, specdec_offline_dataset, low_memory_mode) and raising the same NotImplementedError if issues is non-empty before any use of AutoModelForCausalLM or model construction.
♻️ Duplicate comments (4)
examples/llm_ptq/hf_ptq.py (1)
1552-1558:⚠️ Potential issue | 🟠 Major | ⚡ Quick winReject
--use_fsdp2together with--auto_quantize_bits.
load_and_prepare_fsdp2_model()goes throughfsdp2_shard(), which freezes every parameter beforeauto_quantize()runs. The help text already says this combo is unsupported, so error out here instead of failing deep in the quantization flow.Suggested fix
if args.use_fsdp2 and os.environ.get("RANK") is None: parser.error("--use_fsdp2 requires launching with torchrun") + if args.use_fsdp2 and args.auto_quantize_bits is not None: + parser.error("--use_fsdp2 is not supported with --auto_quantize_bits") if args.cpu_offload and not args.use_fsdp2: parser.error("--cpu_offload requires --use_fsdp2")🤖 Prompt for AI Agents
Verify each finding against current code. Fix only still-valid issues, skip the rest with a brief reason, keep changes minimal, and validate. In `@examples/llm_ptq/hf_ptq.py` around lines 1552 - 1558, The CLI currently allows combining --use_fsdp2 and --auto_quantize_bits which later breaks because load_and_prepare_fsdp2_model -> fsdp2_shard freezes parameters before auto_quantize runs; add an explicit check after parsing args that if args.use_fsdp2 and args.auto_quantize_bits are both truthy, call parser.error with a clear message rejecting this combination (mentioning --use_fsdp2 and --auto_quantize_bits) so the program exits early instead of failing inside load_and_prepare_fsdp2_model / fsdp2_shard / auto_quantize.modelopt/torch/utils/distributed.py (1)
439-458:⚠️ Potential issue | 🟡 Minor | ⚡ Quick winHide the calibration progress bar on non-main ranks.
tqdmwill render once per rank here, so multi-node calibration still spams overlapping progress bars. Thread anis_mainflag into this helper and disable the bar elsewhere.Suggested fix
-def fsdp_aware_forward_loop(wrapped_model, dataloader, device=None): +def fsdp_aware_forward_loop(wrapped_model, dataloader, device=None, is_main: bool = True): @@ - for batch in tqdm(dataloader, desc="Calibrating"): + for batch in tqdm(dataloader, desc="Calibrating", disable=not is_main):As per coding guidelines
**/*.py: Useprint_rank_0orwarn_rank_0when possible to avoid noisy logs and guard shared side effects against race conditions between ranks in distributed processing.🤖 Prompt for AI Agents
Verify each finding against current code. Fix only still-valid issues, skip the rest with a brief reason, keep changes minimal, and validate. In `@modelopt/torch/utils/distributed.py` around lines 439 - 458, The calibration progress bar in fsdp_aware_forward_loop's inner function calibrate renders on every rank; add an is_main (or rank0) boolean parameter to fsdp_aware_forward_loop and thread it into calibrate, then pass disable=not is_main (or only instantiate tqdm when is_main) so tqdm only shows on the main rank; additionally wrap any shared-side-effect logs with print_rank_0/warn_rank_0 and ensure any state changes that could race between ranks are guarded by the same is_main check.examples/llm_ptq/example_utils.py (2)
74-78:⚠️ Potential issue | 🟠 Major | ⚡ Quick winDon't overwrite the requested device in the non-FSDP path.
Setting
args.device = Nonemakes--device cpua no-op and can later feedNoneintoget_model().Suggested fix
else: args.rank = 0 args.world_size = 1 - args.device = None args.is_main = True🤖 Prompt for AI Agents
Verify each finding against current code. Fix only still-valid issues, skip the rest with a brief reason, keep changes minimal, and validate. In `@examples/llm_ptq/example_utils.py` around lines 74 - 78, The code sets args.device = None in the non-FSDP branch which overwrites any user-requested device and can pass None into get_model(); change this to preserve the provided device (do not assign None) — either remove the args.device = None assignment or set args.device = args.device if already defined, ensuring args.device retains "--device" input before calling get_model() and other functions that expect a valid device.
115-123:⚠️ Potential issue | 🟠 Major | ⚡ Quick winPreserve
--attn_implementationon the FSDP2 load path.Under
--use_fsdp2, this helper ignores the CLI flag entirely, so config/model loading diverges fromget_model()and can pick a different attention backend or memory profile.Suggested fix
def load_and_prepare_fsdp2_model( ckpt_path: str, device: torch.device, rank: int, args=None, trust_remote_code: bool = False, + attn_implementation: str | None = None, mp_policy=None, cpu_offload: bool = False, ): @@ - hf_config = AutoConfig.from_pretrained(ckpt_path, trust_remote_code=trust_remote_code) + config_kwargs = {"trust_remote_code": trust_remote_code} + if attn_implementation is not None: + config_kwargs["attn_implementation"] = attn_implementation + hf_config = AutoConfig.from_pretrained(ckpt_path, **config_kwargs) @@ src_model = AutoModelForCausalLM.from_pretrained( ckpt_path, torch_dtype="auto", trust_remote_code=trust_remote_code, + attn_implementation=attn_implementation, low_cpu_mem_usage=True, ) @@ model = AutoModelForCausalLM.from_config( - hf_config, torch_dtype=dtype, trust_remote_code=trust_remote_code + hf_config, + torch_dtype=dtype, + trust_remote_code=trust_remote_code, + attn_implementation=attn_implementation, )Please also pass
args.attn_implementationfromexamples/llm_ptq/hf_ptq.py.Also applies to: 146-168
🤖 Prompt for AI Agents
Verify each finding against current code. Fix only still-valid issues, skip the rest with a brief reason, keep changes minimal, and validate. In `@examples/llm_ptq/example_utils.py` around lines 115 - 123, The FSDP2 load path in load_and_prepare_fsdp2_model currently ignores the CLI flag for attention backend; update the function to accept and propagate args.attn_implementation (or a dedicated attn_implementation param) when building the model/config so the same attention implementation used by get_model() is preserved, and ensure callers (examples/llm_ptq/hf_ptq.py) pass args.attn_implementation into load_and_prepare_fsdp2_model; also mirror this change in the related FSDP2 code paths around the block at 146-168 so the attention backend and memory profile remain consistent.
🤖 Prompt for all review comments with AI agents
Verify each finding against current code. Fix only still-valid issues, skip the
rest with a brief reason, keep changes minimal, and validate.
Outside diff comments:
In `@examples/llm_ptq/example_utils.py`:
- Around line 97-112: The validation block that collects issues must also reject
encoder-decoder configs so load_and_prepare_fsdp2_model() doesn't call
AutoModelForCausalLM on incompatible families; detect encoder-decoder models by
checking config.is_encoder_decoder (or config.model_type in known
encoder-decoder types) and append a descriptive message to issues (e.g.,
"encoder-decoder models (T5/BART/Whisper)"), keeping the existing checks (vila,
is_nemotron_vl, _is_multimodal_config, quantization_config,
specdec_offline_dataset, low_memory_mode) and raising the same
NotImplementedError if issues is non-empty before any use of
AutoModelForCausalLM or model construction.
---
Duplicate comments:
In `@examples/llm_ptq/example_utils.py`:
- Around line 74-78: The code sets args.device = None in the non-FSDP branch
which overwrites any user-requested device and can pass None into get_model();
change this to preserve the provided device (do not assign None) — either remove
the args.device = None assignment or set args.device = args.device if already
defined, ensuring args.device retains "--device" input before calling
get_model() and other functions that expect a valid device.
- Around line 115-123: The FSDP2 load path in load_and_prepare_fsdp2_model
currently ignores the CLI flag for attention backend; update the function to
accept and propagate args.attn_implementation (or a dedicated
attn_implementation param) when building the model/config so the same attention
implementation used by get_model() is preserved, and ensure callers
(examples/llm_ptq/hf_ptq.py) pass args.attn_implementation into
load_and_prepare_fsdp2_model; also mirror this change in the related FSDP2 code
paths around the block at 146-168 so the attention backend and memory profile
remain consistent.
In `@examples/llm_ptq/hf_ptq.py`:
- Around line 1552-1558: The CLI currently allows combining --use_fsdp2 and
--auto_quantize_bits which later breaks because load_and_prepare_fsdp2_model ->
fsdp2_shard freezes parameters before auto_quantize runs; add an explicit check
after parsing args that if args.use_fsdp2 and args.auto_quantize_bits are both
truthy, call parser.error with a clear message rejecting this combination
(mentioning --use_fsdp2 and --auto_quantize_bits) so the program exits early
instead of failing inside load_and_prepare_fsdp2_model / fsdp2_shard /
auto_quantize.
In `@modelopt/torch/utils/distributed.py`:
- Around line 439-458: The calibration progress bar in fsdp_aware_forward_loop's
inner function calibrate renders on every rank; add an is_main (or rank0)
boolean parameter to fsdp_aware_forward_loop and thread it into calibrate, then
pass disable=not is_main (or only instantiate tqdm when is_main) so tqdm only
shows on the main rank; additionally wrap any shared-side-effect logs with
print_rank_0/warn_rank_0 and ensure any state changes that could race between
ranks are guarded by the same is_main check.
ℹ️ Review info
⚙️ Run configuration
Configuration used: Path: .coderabbit.yaml
Review profile: CHILL
Plan: Enterprise
Run ID: 0c53a76c-6ab3-44c4-92db-4acf83822755
📒 Files selected for processing (3)
examples/llm_ptq/example_utils.pyexamples/llm_ptq/hf_ptq.pymodelopt/torch/utils/distributed.py
sugunav14
marked this pull request as draft
May 29, 2026 20:39
cjluo-nv
reviewed
Jun 2, 2026
cjluo-nv
left a comment
cjluo-nv
left a comment
Collaborator
There was a problem hiding this comment.
Bot review — DM the bot to share feedback.
Substantial new FSDP2 subsystem (~860 LOC across examples/llm_ptq + modelopt/torch/utils/distributed.py) but the PR description is the unfilled template — no problem statement, no testing notes, no Changelog mention, the backward-compat / tests / Claude-approval boxes are unchecked. The deterministic complexity gate fired (1318 lines / 5 dirs / >300-line single non-test file), so design review is required, and the design question is unaddressed in the PR body.
Design protocol — alternatives that are not justified against:
examples/llm_ptq/multinode_ptq.py(deleted in this PR) already did multi-node PTQ via HuggingFace Accelerate's FSDP2 integration. The PR rips that out and replaces it with a hand-rolledfsdp2_wrap/fsdp2_shard/_load_via_parallel_read/broadcast_state_dictstack. The PR body never explains why accelerate's FSDP2 path was insufficient and a parallel custom system is the right answer.modelopt/torch/quantization/plugins/transformers_trainer.py(_quantize_model,_patch_accelerate_for_fsdp2_fix) already runs FSDP2 calibration via accelerate and has the same "use the wrapped model, not the unwrapped arg" forward-loop pattern. The newfsdp_aware_forward_loopcarries aTODO: ... consolidate it onto this helper, i.e. the author knows the duplication exists. Introducing a second FSDP2 calibration entry point before consolidating is exactly the "two systems for one job" anti-pattern.transformers.AutoModelForCausalLM.from_pretrainedalready supports sharded safetensors loading viadevice_map/accelerate. The custom_load_via_parallel_read(per-layer broadcast + FSDP2 set_model_state_dict) is a sizable parallel I/O subsystem that needs justification — at minimum, numbers showing the rank-0-load fallback is unworkable for the target models.
Please address these in the PR body (or extend the existing accelerate path) before this lands.
Concrete issues found:
examples/llm_ptq/hf_ptq.py:948and:1005: preview-generationmax_new_tokens=100→10is changed unconditionally, not just on the FSDP2 path. This silently affects every existing user ofhf_ptq.py(shorter pre/post-PTQ qualitative previews). If it's an FSDP2-perf workaround, gate it onargs.use_fsdp2; otherwise call it out as an intentional global change.examples/llm_ptq/hf_ptq.py:_export_fsdp2_hf_checkpoint: ~30 lines that re-implementexport_hf_checkpointinmodelopt/torch/export/unified_export_hf.py(writehf_quant_config.json,save_pretrainedwith state_dict, patchquantization_configintoconfig.json). The only real differences are rank-0 gating and_original_architectures. Worth extendingexport_hf_checkpointto take a "rank0-only" flag and an architectures override rather than forking a second copy.modelopt/torch/export/unified_export_hf.py:_export_transformers_checkpoint:get_model_state_dict(model, options=StateDictOptions(full_state_dict=True, cpu_offload=True))materializes the full unsharded state dict on every rank (cpu_offload moves it to host, but each rank still gets a full copy). PR comment claims "Rank 0 holds the full model in CPU briefly … other ranks pay ~0 CPU" — that's not whatfull_state_dict=Truegives you. Either passbroadcast_from_rank0-style options to make non-rank-0 cheap, or fix the comment.modelopt/torch/utils/distributed.py:shard_dataloader: silently dropsworker_init_fn,persistent_workers,prefetch_factor,timeout,multiprocessing_context, andgeneratorwhen reconstructing the DataLoader. Probably fine for the current calibration use case but a footgun if reused; either copy them through or document the limitation.modelopt/torch/utils/distributed.py:fsdp2_wrapusesLayerActivationCollector.get_decoder_layers, which returnsnn.ModuleList | None;if not layerstreats both None and an empty ModuleList as the same error, but the messages differ for the two paths and the empty case is currently impossible to hit. Minor — consideris None.
Test coverage:
The only new test (tests/gpu/torch/quantization/test_fsdp2.py:test_writeback_root_unwrapped) is a regression guard for the loosened fsdp2_weight_access_and_writeback_context assert — it does not cover any of the new public API: load_fsdp2_causal_lm, fsdp2_shard, _load_via_parallel_read, broadcast_state_dict, shard_dataloader, fsdp_aware_forward_loop, _export_fsdp2_hf_checkpoint, validate_fsdp2_supported. For ~500 LOC of new framework code in modelopt/torch/utils/distributed.py plus an entirely new CLI flag, that's not enough. Mandatory tests box is unchecked in the PR body too. At minimum: a 2-rank GPU test that loads a tiny causal-LM via load_fsdp2_causal_lm (parallel-read + fallback), runs a calibration forward via fsdp_aware_forward_loop, and exports to HF with rank-0-only writes — checking that the resulting config.json has the original (non-FSDP-prefixed) architectures.
| loader.dataset, | ||
| batch_size=loader.batch_size, | ||
| sampler=sampler, | ||
| collate_fn=loader.collate_fn, |
Collaborator
There was a problem hiding this comment.
Bot comment.
TODO acknowledges that transformers_trainer.py:_quantize_model already inlines this exact "call the wrapped model, ignore the unwrapped arg" pattern. Introducing a second copy of the helper before consolidating is the wrong order — please either land the consolidation in this PR or open a follow-up issue and link it from the TODO so this doesn't quietly become permanent.
| return calibrate | ||
|
|
||
|
|
||
| def broadcast_state_dict( |
Collaborator
There was a problem hiding this comment.
Bot comment.
shard_dataloader only forwards batch_size, collate_fn, num_workers, pin_memory. It silently drops worker_init_fn, persistent_workers, prefetch_factor, timeout, multiprocessing_context, and generator from the input loader. For calibration today this is fine, but the helper is a public-looking entry in __all__ — please either propagate those attributes (getattr(loader, ..., default) is enough) or note the limitation in the docstring.
|
|
||
|
|
||
| def test_writeback_root_unwrapped(dist_workers): | ||
| dist_workers.run(_test_writeback_root_unwrapped) |
Collaborator
There was a problem hiding this comment.
Bot comment.
This is the only new test, and it covers the loosened fsdp2_weight_access_and_writeback_context precondition — none of the new public API (load_fsdp2_causal_lm, fsdp2_shard, _load_via_parallel_read, broadcast_state_dict, shard_dataloader, fsdp_aware_forward_loop, _export_fsdp2_hf_checkpoint, validate_fsdp2_supported) has direct test coverage. Please add at least a 2-rank GPU end-to-end test (tiny causal-LM → load_fsdp2_causal_lm → fsdp_aware_forward_loop calibration → rank-0 export) and verify the exported config.json keeps the pre-FSDP architectures list.
sugunav14
marked this pull request as ready for review
June 4, 2026 18:10
Contributor
There was a problem hiding this comment.
Actionable comments posted: 4
Caution
Some comments are outside the diff and can’t be posted inline due to platform limitations.
⚠️ Outside diff range comments (1)
examples/llm_ptq/hf_ptq.py (1)
1490-1497:⚠️ Potential issue | 🟠 Major | ⚡ Quick winReject
--use_fsdp2 --sparsity_fmt sparsegptwith the other FSDP2 guardrails.The new help text says FSDP2 v1 does not support sparsity, but the parser still accepts that combination and
main()will route it intosparsity_main()after the FSDP2 load path. Please fail fast here next to the existing--use_fsdp2validation.🤖 Prompt for AI Agents
Verify each finding against current code. Fix only still-valid issues, skip the rest with a brief reason, keep changes minimal, and validate. In `@examples/llm_ptq/hf_ptq.py` around lines 1490 - 1497, Add a fast-fail check next to the existing FSDP2 validation to reject using FSDP2 together with the sparsity format that triggers sparsity processing: if args.use_fsdp2 and args.sparsity_fmt == "sparsegpt" then call parser.error(...) with a clear message that FSDP2 v1 does not support sparsity/sparsegpt; this prevents main() from routing into sparsity_main() when --use_fsdp2 is set. Use the existing symbols args.use_fsdp2, args.sparsity_fmt, parser.error and reference sparsity_main() in the message if helpful.
🧹 Nitpick comments (1)
tests/gpu/torch/quantization/test_fsdp2.py (1)
266-313: ⚡ Quick winAdd a regression for the CPU-offload writeback path too.
This new test covers the root-unwrapped case, but the production change also added the
cpu_local/gpu_workingmirror branch infsdp2_weight_access_and_writeback_context(). Because this fixture keeps the gathered weight on CUDA the whole time, that sync-back path is still untested. A small distributed case with FSDP CPU offload would lock down the new behavior much better.As per coding guidelines, "Write focused unit tests during development; before staging or committing, curate production tests by removing redundant lower-level tests and keeping only tests that document expected behavior, protect against regressions, or flag backward-incompatible changes."
🤖 Prompt for AI Agents
Verify each finding against current code. Fix only still-valid issues, skip the rest with a brief reason, keep changes minimal, and validate. In `@tests/gpu/torch/quantization/test_fsdp2.py` around lines 266 - 313, Add a distributed regression that exercises the CPU-offload writeback mirror branch by creating a variant of test_writeback_root_unwrapped which enables FSDP CPU offload (so the gathered weight moves to CPU and then back to GPU_working) and calls the same enable_weight_access_and_writeback(layer[0], model) context; specifically, duplicate or parametrize _test_writeback_root_unwrapped to configure FSDP with cpu_offload=True (or the project’s equivalent flag) and ensure you invoke fsdp2_weight_access_and_writeback_context()/enable_weight_access_and_writeback to validate both the cpu_local/gpu_working sync-back path and that shards are restored and modified values are written back. Ensure the new test runs under dist_workers like test_writeback_root_unwrapped and asserts DTensor sharding before/after and correctness of the mutated weight.
🤖 Prompt for all review comments with AI agents
Verify each finding against current code. Fix only still-valid issues, skip the
rest with a brief reason, keep changes minimal, and validate.
Inline comments:
In `@examples/llm_ptq/example_utils.py`:
- Around line 78-96: The _checkpoint_has_mtp_weights helper currently only
inspects local files or already-cached indexes and thus misses remote HF repo
indexes on a cold cache; change _checkpoint_has_mtp_weights to first attempt to
materialize/download the index (use try_to_load_from_cache(model_path,
"model.safetensors.index.json") when callable and, if that returns None, trigger
a download/materialization of the remote "model.safetensors.index.json" into a
local temp/file cache) before inspecting it, then run the existing
json.load/weight_map check; ensure validate_fsdp2_supported calls this updated
_checkpoint_has_mtp_weights after the checkpoint path has been materialized
locally so MTP checkpoints are detected reliably.
In `@modelopt/torch/export/unified_export_hf.py`:
- Around line 1223-1227: The distributed barrier currently only runs on non-zero
ranks before returning, which can hang if rank 0 raises an exception later; make
the barrier exception-safe by ensuring every rank executes
torch.distributed.barrier() inside a finally block. Concretely, wrap the logic
that calls save_pretrained() and the subsequent config rewrite (the block
guarded by torch.distributed.is_available()/is_initialized() and
torch.distributed.get_rank()) in a try/finally so that
torch.distributed.barrier() is always called by all ranks in the finally, and
only after that should non-zero ranks return; i.e., reference and protect the
calls to save_pretrained() and the config rewrite and ensure
torch.distributed.barrier() is invoked in the finally on both failure and
success.
In `@modelopt/torch/quantization/utils/core_utils.py`:
- Around line 500-513: The temporary local nn.Parameter assignments (created via
local_replicated, working, cpu_mirror and set into module with
_set_parameter(module, name, nn.Parameter(working))) must be restored even if
the context body raises; wrap the code that sets originals[name] and calls
_set_parameter in a try/finally and move the writeback/teardown loop into the
finally so the original parameter, collected shard, original_placements,
original_device_mesh, and any cpu_mirror syncing are always restored/synced;
apply the same try/finally pattern to the similar block around the lines noted
(the block at 523-527) so both success and exception paths perform cleanup.
In `@modelopt/torch/utils/model_load_utils.py`:
- Around line 124-126: The code currently sets model.config.use_cache = False
mutating the live HF config; instead capture the original value (e.g.,
orig_use_cache = model.config.use_cache) before changing it, set use_cache to
False only for the local calibration/export operation (around the calls that
require it), and restore model.config.use_cache = orig_use_cache before
returning or before any save_pretrained/export path; alternatively move the
temporary disable into the specific calibration/export functions so model and
its config are never left mutated (refer to the model.eval() /
model.config.use_cache lines to locate the change).
---
Outside diff comments:
In `@examples/llm_ptq/hf_ptq.py`:
- Around line 1490-1497: Add a fast-fail check next to the existing FSDP2
validation to reject using FSDP2 together with the sparsity format that triggers
sparsity processing: if args.use_fsdp2 and args.sparsity_fmt == "sparsegpt" then
call parser.error(...) with a clear message that FSDP2 v1 does not support
sparsity/sparsegpt; this prevents main() from routing into sparsity_main() when
--use_fsdp2 is set. Use the existing symbols args.use_fsdp2, args.sparsity_fmt,
parser.error and reference sparsity_main() in the message if helpful.
---
Nitpick comments:
In `@tests/gpu/torch/quantization/test_fsdp2.py`:
- Around line 266-313: Add a distributed regression that exercises the
CPU-offload writeback mirror branch by creating a variant of
test_writeback_root_unwrapped which enables FSDP CPU offload (so the gathered
weight moves to CPU and then back to GPU_working) and calls the same
enable_weight_access_and_writeback(layer[0], model) context; specifically,
duplicate or parametrize _test_writeback_root_unwrapped to configure FSDP with
cpu_offload=True (or the project’s equivalent flag) and ensure you invoke
fsdp2_weight_access_and_writeback_context()/enable_weight_access_and_writeback
to validate both the cpu_local/gpu_working sync-back path and that shards are
restored and modified values are written back. Ensure the new test runs under
dist_workers like test_writeback_root_unwrapped and asserts DTensor sharding
before/after and correctness of the mutated weight.
🪄 Autofix (Beta)
Fix all unresolved CodeRabbit comments on this PR:
- Push a commit to this branch (recommended)
- Create a new PR with the fixes
ℹ️ Review info
⚙️ Run configuration
Configuration used: Path: .coderabbit.yaml
Review profile: CHILL
Plan: Enterprise
Run ID: 6556fc96-fe9a-4dd9-b207-1dc70a4eccc3
📒 Files selected for processing (7)
examples/llm_ptq/example_utils.pyexamples/llm_ptq/hf_ptq.pymodelopt/torch/export/unified_export_hf.pymodelopt/torch/quantization/utils/core_utils.pymodelopt/torch/utils/distributed.pymodelopt/torch/utils/model_load_utils.pytests/gpu/torch/quantization/test_fsdp2.py
cjluo-nv
reviewed
Jun 4, 2026
cjluo-nv
left a comment
cjluo-nv
left a comment
Collaborator
There was a problem hiding this comment.
Bot review — DM the bot to share feedback.
Several critical review points from prior rounds addressed, but a few are still open and the design-protocol questions raised on the prior round are still not addressed in the PR body. Recommending human sign-off.
Addressed since prior review:
- 💬
args.device = Noneclobber in non-FSDP path — fixed (setup_distributed_argsno longer overwritesargs.devicein the else branch). - 💬 tqdm progress bar on every rank — fixed via
disable=not is_master()infsdp_aware_forward_loop. - 💬
--attn_implementationignored on FSDP2 load path — now threaded throughload_and_prepare_fsdp2_model→parallel_load_and_prepare_fsdp2→build_meta_causal_lm, andhf_ptq.pypassesargs.attn_implementation. - 💬
--use_fsdp2+--auto_quantize_bits— guarded by an assert at the top ofauto_quantize()(and the help text already declared it unsupported). - 💬
_export_fsdp2_hf_checkpointduplicatingexport_hf_checkpoint— refactored away; the FSDP2 path now goes throughexport_hf_checkpointitself, which gained rank-0 gating + barrier and usesget_model_state_dict(..., broadcast_from_rank0=True)so non-rank-0 no longer pays full host state dict. - 💬
shard_dataloaderdropping DataLoader kwargs — now forwardsworker_init_fn,persistent_workers,prefetch_factor,timeout,multiprocessing_context,generator,pin_memory_device. - 💬
max_new_tokens=100→10global preview change — appears reverted; currenthf_ptq.pyusesmax_new_tokens=100for both pre- and post-PTQ previews.
Still open / worth a human's eyes before approving:
- 💬 Design-protocol questions raised in the prior review (PR body justification for the custom parallel-safetensors loader vs. HF
device_map/accelerate, custom FSDP2 stack vs. the deleted accelerate-basedmultinode_ptq.py, and thefsdp_aware_forward_loop↔transformers_trainer.py:_quantize_modelduplication that still carries aTODO: ... consolidate it onto this helper) — the PR body has not been updated to address them. The code still ships a second FSDP2 calibration helper alongside the existing one intransformers_trainer.py. - Test coverage of the new public surface is still thin:
tests/gpu/torch/quantization/test_fsdp2.pyaddstest_writeback_root_unwrappedandtest_writeback_cpu_offload(good — they cover the loosened assert and the new CPU↔GPU mirror) andtests/unit/torch/utils/test_model_load_utils.pycoversweight_map_for/read_safetensors_subset, butparallel_load_and_prepare_fsdp2,fsdp2_wrap,shard_dataloader,fsdp_aware_forward_loop,broadcast_state_dict, the FSDP2 export path inexport_hf_checkpoint, andvalidate_fsdp2_supportedhave no direct tests. A 2-rank GPU end-to-end (tiny causal-LM → load → calibrate → rank-0 export, asserting the exportedconfig.jsonkeeps the originalarchitectures) is the obvious gap. --use_fsdp2 --sparsity_fmt sparsegptis still accepted byparse_args;main()then routes intosparsity_main()after FSDP2 load. The help text says sparsity is unsupported — please fail-fast inparse_argsnext to the other--use_fsdp2guards.examples/llm_ptq/example_utils.py:_checkpoint_has_mtp_weightsstill only inspects the local checkpoint dir + an already-cachedmodel.safetensors.index.json(try_to_load_from_cachereturnsNoneon a cold cache). For a remote HF repo ID with no cache hit,validate_fsdp2_supportedwill silently accept an MTP checkpoint and the FSDP2 loader will drop the MTP weights at export. Either materialize the index first or move the MTP check after the snapshot download already done insideparallel_load_and_prepare_fsdp2.modelopt/torch/utils/model_load_utils.py:build_meta_causal_lmmutatesmodel.config.use_cache = Falseand never restores it; the same model is later passed tosave_pretrained()so the exported checkpoint inheritsuse_cache=False. Capture/restore around the calibration scope or move the disable into the calibration call site.modelopt/torch/quantization/utils/core_utils.py:fsdp2_weight_access_and_writeback_contextdoes not wrap theyieldintry/finally. If the body raises (calibration error, OOM mid-context, …) the temporary localnn.Parameters set via_set_parameterare never restored and the new CPU mirror never syncs back, leaving the module in a broken state.modelopt/torch/export/unified_export_hf.py:export_hf_checkpointputs the trailingtorch.distributed.barrier()outside thetry/except; if rank 0 raises duringsave_pretrained/config rewrite it re-raises before reaching the barrier, and non-zero ranks have already passed theirs — fine in this direction, but if rank 0 raises before the early non-zero-rank barrier (e.g. inside_export_transformers_checkpoint's rank-0-only work, sincebroadcast_from_rank0=Truemakes rank 0 do extra work after the gather), non-zero ranks wait out NCCL timeout. Putting both barriers in afinallywould close this. The in-code comment acknowledges the gap, but it's still worth a second pair of eyes given how easy the fix is.fsdp_aware_forward_loopstill carriesTODO: consolidate transformers_trainer.py:_quantize_model onto this helper— file a follow-up issue and reference it from the TODO so the duplication doesn't quietly become permanent.
cjluo-nv
reviewed
Jun 4, 2026
cjluo-nv
left a comment
cjluo-nv
left a comment
Collaborator
There was a problem hiding this comment.
Bot review — DM the bot to share feedback.
Major prior-round issues are addressed in the current diff; flagging for human sign-off because this is still a 1500-line architectural change introducing a parallel FSDP2 stack alongside the existing one, and a few lower-priority items remain.
Resolved since prior review:
- 💬 Author addressed:
args.device = Noneclobber removed; tqdm rank-gated viadisable=not is_master();attn_implementationthreaded throughload_and_prepare_fsdp2_model→parallel_load_and_prepare_fsdp2→build_meta_causal_lm. - 💬
--use_fsdp2 + --auto_quantize_bitsand--use_fsdp2 + --sparsity_fmt sparsegptnow fail fast (assert inauto_quantize, parser.error inparse_args). - 💬
_export_fsdp2_hf_checkpointduplication eliminated — FSDP2 path now goes throughexport_hf_checkpoint, which gained rank-0 gating +try/finallybarrier and usesget_model_state_dict(..., broadcast_from_rank0=True)so non-rank-0 ranks don't materialize a full host state dict. - 💬
shard_dataloadernow forwardsworker_init_fn,persistent_workers,prefetch_factor,timeout,multiprocessing_context,generator,pin_memory_device. - 💬
core_utils.pywriteback wrapped intry/finallyso the temporary local params and CPU mirror are always restored on exception. - 💬
_checkpoint_has_mtp_weightsnow useshf_hub_downloadto resolve the index for remote HF repo IDs on a cold cache. - 💬 PR body now has a "Design Notes" section that explicitly addresses the three previously-raised design-protocol questions (custom parallel safetensors loader vs HF/accelerate, custom FSDP2 stack vs accelerate-based multinode_ptq.py, and the
fsdp_aware_forward_loop ↔ transformers_trainer.py:_quantize_modelduplication). - 💬 Test coverage broadened —
tests/gpu/torch/utils/test_model_load_utils.pyadds a 2-rank E2E test (test_parallel_load_and_export) that loads a tiny LlamaForCausalLM viaparallel_load_and_prepare_fsdp2, runs a forward, exports viaexport_hf_checkpoint, and asserts the savedconfig.jsonkeepsarchitectures == ["LlamaForCausalLM"]. Also coversbroadcast_state_dict,shard_dataloader,fsdp_aware_forward_loop, plusweight_map_for/read_safetensors_subsetunit tests.
Still warrants a human's eyes before approval:
- Architectural surface: this lands a second FSDP2 calibration path (
fsdp_aware_forward_loop) alongside the existing one inmodelopt/torch/quantization/plugins/transformers_trainer.py:_quantize_model. The PR body acknowledges the duplication and notes a TODO to consolidate, but the TODO inmodelopt/torch/utils/distributed.py:fsdp_aware_forward_loopis not linked to a tracked follow-up issue, so this risks becoming permanent. Recommend filing a follow-up and referencing it from the TODO before merge. parallel_load_and_prepare_fsdp2only verifies that the FSDP2 export path preservesarchitectures; it does not exercise CPU-offload end-to-end (the new--cpu_offloadcodepath has unit-level coverage only viatest_writeback_cpu_offloadon a synthetic 32-dim model). A 2-rankparallel_load_and_prepare_fsdp2(..., cpu_offload=True)smoke test would lock down the_promote_non_dtensor_to_gpu+ per-layer CPU shard plumbing.parse_argsaccepts but does not validate--use_fsdp2against--vllm_fakequant_exportor--cast_mxfp4_to_nvfp4; both will route intoexport_quantizedpaths that don't go through the rank-0-gated FSDP2 export. v1 scope should either claim or reject these explicitly.unified_export_hf.pycomment correctly notes that if rank 0 raises before the collective gather inside_export_transformers_checkpoint, the other ranks hang on NCCL timeout — out of scope per the comment, but worth a human eyeballing whether that gap is acceptable for v1.
cjluo-nv
reviewed
Jun 4, 2026
cjluo-nv
left a comment
cjluo-nv
left a comment
Collaborator
There was a problem hiding this comment.
Bot review — DM the bot to share feedback.
Most prior-round critical issues are addressed in the current diff; flagging for human sign-off given this is a 1500-line architectural change introducing a parallel FSDP2 stack alongside the existing one.
Resolved since prior reviews:
- 💬
args.device = Noneclobber removed; tqdm rank-gated viadisable=not is_master();attn_implementationthreaded throughload_and_prepare_fsdp2_model→parallel_load_and_prepare_fsdp2→build_meta_causal_lm. - 💬
--use_fsdp2 + --auto_quantize_bits(assert inauto_quantize) and--use_fsdp2 + --sparsity_fmt sparsegpt(parser.error inparse_args) now fail fast. - 💬
_export_fsdp2_hf_checkpointduplication eliminated — FSDP2 path goes throughexport_hf_checkpointwhich gained rank-0 gating +try/finallybarrier and usesget_model_state_dict(..., broadcast_from_rank0=True). - 💬
shard_dataloadernow forwardsworker_init_fn,persistent_workers,prefetch_factor,timeout,multiprocessing_context,generator,pin_memory_device. - 💬
core_utils.pywriteback wrapped intry/finallyso temporary local params and CPU mirror are always restored on exception. - 💬
_checkpoint_has_mtp_weightsuseshf_hub_downloadto resolve the index for remote HF repo IDs on a cold cache. - 💬
model_load_utils.pyno longer mutatesmodel.config.use_cache. - 💬 PR body has a "Design Notes" section explicitly addressing the three previously-raised design-protocol questions (custom parallel safetensors loader vs HF/accelerate, custom FSDP2 stack vs accelerate-based
multinode_ptq.py, and thefsdp_aware_forward_loop ↔ transformers_trainer.py:_quantize_modelduplication). - 💬 Test coverage broadened —
tests/gpu/torch/utils/test_model_load_utils.pyadds a 2-rank E2Etest_parallel_load_and_export(load → forward → export → assertarchitecturespreserved) plus tests forbroadcast_state_dict,shard_dataloader,fsdp_aware_forward_loop,weight_map_for,read_safetensors_subset. Plustest_writeback_cpu_offloadcovers the new CPU↔GPU mirror.
Still warrants a human's eyes:
- Architectural surface — this lands a second FSDP2 calibration helper (
fsdp_aware_forward_loop) alongside the existing inlined one inmodelopt/torch/quantization/plugins/transformers_trainer.py:_quantize_model. The PR body acknowledges the duplication and the helper has aTODO: ... consolidate it onto this helper, but the TODO is not linked to a tracked follow-up issue; recommend filing one and referencing it from the TODO before merge so the second copy doesn't quietly become permanent. - v1 scope guards in
parse_argsreject--use_seq_device_map,--auto_quantize_bits, and--sparsity_fmt sparsegptbut do not validate--use_fsdp2against--vllm_fakequant_exportor--cast_mxfp4_to_nvfp4; both will route intoexport_quantizedpaths that don't go through the rank-0-gated FSDP2 export. v1 should either claim or reject these explicitly. - The 2-rank E2E test exercises the standard FSDP2 path end-to-end, but
parallel_load_and_prepare_fsdp2(..., cpu_offload=True)is only covered at unit level viatest_writeback_cpu_offloadon a synthetic 32-dim model. A 2-rank smoke test withcpu_offload=Truewould lock down the_promote_non_dtensor_to_gpu+ per-layer CPU shard plumbing. unified_export_hf.pycomment correctly notes that if rank 0 raises before the collective gather inside_export_transformers_checkpoint, the other ranks hang on NCCL timeout — out of scope per the comment, but worth a human eyeballing whether that gap is acceptable for v1.
Contributor
There was a problem hiding this comment.
🧹 Nitpick comments (1)
tests/gpu/torch/utils/test_model_load_utils.py (1)
123-127: 💤 Low valueTest leaves temporary directories behind.
The test creates checkpoint and export directories in
tempfile.gettempdir()but never cleans them up. Consider usingtempfile.TemporaryDirectory()with proper synchronization, or add cleanup in afinallyblock.+ import shutil + ckpt_dir = os.path.join(tempfile.gettempdir(), f"_test_parallel_load_{os.getpid()}") + export_dir = os.path.join(tempfile.gettempdir(), f"_test_parallel_export_{os.getpid()}") if rank == 0: os.makedirs(ckpt_dir, exist_ok=True) _build_tiny_llama_checkpoint(ckpt_dir) dist.barrier() ... - export_dir = os.path.join(tempfile.gettempdir(), f"_test_parallel_export_{os.getpid()}") - if rank == 0: - os.makedirs(export_dir, exist_ok=True) - dist.barrier() + try: + if rank == 0: + os.makedirs(export_dir, exist_ok=True) + dist.barrier() ... + finally: + dist.barrier() + if rank == 0: + shutil.rmtree(ckpt_dir, ignore_errors=True) + shutil.rmtree(export_dir, ignore_errors=True)Also applies to: 145-148
🤖 Prompt for AI Agents
Verify each finding against current code. Fix only still-valid issues, skip the rest with a brief reason, keep changes minimal, and validate. In `@tests/gpu/torch/utils/test_model_load_utils.py` around lines 123 - 127, The test creates ckpt_dir (and the export dir around lines 145-148) under tempfile.gettempdir() and never removes them; wrap the setup and teardown in a try/finally so the directories are removed, or use tempfile.TemporaryDirectory() on rank==0 and ensure other ranks wait on dist.barrier() to use the path, then have rank==0 clean up (shutil.rmtree or let TemporaryDirectory go out of scope) after a final dist.barrier(); update the block that calls _build_tiny_llama_checkpoint and the export-dir creation to perform creation on rank==0 and deletion in the finally (or via TemporaryDirectory) with proper synchronization.
🤖 Prompt for all review comments with AI agents
Verify each finding against current code. Fix only still-valid issues, skip the
rest with a brief reason, keep changes minimal, and validate.
Nitpick comments:
In `@tests/gpu/torch/utils/test_model_load_utils.py`:
- Around line 123-127: The test creates ckpt_dir (and the export dir around
lines 145-148) under tempfile.gettempdir() and never removes them; wrap the
setup and teardown in a try/finally so the directories are removed, or use
tempfile.TemporaryDirectory() on rank==0 and ensure other ranks wait on
dist.barrier() to use the path, then have rank==0 clean up (shutil.rmtree or let
TemporaryDirectory go out of scope) after a final dist.barrier(); update the
block that calls _build_tiny_llama_checkpoint and the export-dir creation to
perform creation on rank==0 and deletion in the finally (or via
TemporaryDirectory) with proper synchronization.
ℹ️ Review info
⚙️ Run configuration
Configuration used: Path: .coderabbit.yaml
Review profile: CHILL
Plan: Enterprise
Run ID: 8b1c5e94-fe55-4e7c-9dfc-be7f8cef41cd
📒 Files selected for processing (6)
examples/llm_ptq/example_utils.pyexamples/llm_ptq/hf_ptq.pymodelopt/torch/export/unified_export_hf.pymodelopt/torch/quantization/utils/core_utils.pymodelopt/torch/utils/model_load_utils.pytests/gpu/torch/utils/test_model_load_utils.py
💤 Files with no reviewable changes (1)
- modelopt/torch/utils/model_load_utils.py
🚧 Files skipped from review as they are similar to previous changes (2)
- modelopt/torch/export/unified_export_hf.py
- examples/llm_ptq/example_utils.py
realAsma
reviewed
Jun 4, 2026
Comment on lines
+1224
to
+1227
| # Under torch.distributed: only rank 0 writes; everyone syncs at the finally barrier. | ||
| # If rank 0 raises BEFORE the collective gather inside _export_transformers_checkpoint | ||
| # (e.g. rank-divergent preprocessing), other ranks hang on that collective until NCCL | ||
| # timeout — closing that case would need a broadcast-status pattern; out of scope. |
Contributor
There was a problem hiding this comment.
@sugunav14 do you know how complex it is to write independently from each process? Currently we are bound by the sequential single process disk write speed.
Can we have seperate process write seperate safe-tensor shards?
If this is too much for this PR; can you add a TODO here that we should make the export disk write as well parallel in future?
Contributor
Author
There was a problem hiding this comment.
I still need to look into how much effort it would take to add that support. Maybe we can push it to a follow up PR since this PR is already significantly large to review
Contributor
There was a problem hiding this comment.
Sure; This seems like a require item to me. Current design will cause large models to go OOM. See https://github.com/NVIDIA/Model-Optimizer/pull/1563/changes#r3365625990. Can we have a TODO note atleast?
realAsma
reviewed
Jun 4, 2026
Fridah-nv
reviewed
Jun 4, 2026
Fridah-nv
reviewed
Jun 5, 2026
Signed-off-by: Suguna Velury <178320438+sugunav14@users.noreply.github.com>
Signed-off-by: Suguna Velury <178320438+sugunav14@users.noreply.github.com>
Signed-off-by: Suguna Velury <178320438+sugunav14@users.noreply.github.com>
Signed-off-by: Suguna Velury <178320438+sugunav14@users.noreply.github.com>
Signed-off-by: Suguna Velury <178320438+sugunav14@users.noreply.github.com>
Signed-off-by: Suguna Velury <178320438+sugunav14@users.noreply.github.com>
Signed-off-by: Suguna Velury <178320438+sugunav14@users.noreply.github.com>
Signed-off-by: Suguna Velury <178320438+sugunav14@users.noreply.github.com>
…r param Signed-off-by: Suguna Velury <178320438+sugunav14@users.noreply.github.com>
Signed-off-by: Suguna Velury <178320438+sugunav14@users.noreply.github.com>
Signed-off-by: Suguna Velury <178320438+sugunav14@users.noreply.github.com>
Signed-off-by: Suguna Velury <178320438+sugunav14@users.noreply.github.com>
Signed-off-by: Suguna Velury <178320438+sugunav14@users.noreply.github.com>
Signed-off-by: Suguna Velury <178320438+sugunav14@users.noreply.github.com>
Signed-off-by: Suguna Velury <178320438+sugunav14@users.noreply.github.com>
Signed-off-by: Suguna Velury <178320438+sugunav14@users.noreply.github.com>
Signed-off-by: Suguna Velury <178320438+sugunav14@users.noreply.github.com>
Signed-off-by: Suguna Velury <178320438+sugunav14@users.noreply.github.com>
Signed-off-by: Suguna Velury <178320438+sugunav14@users.noreply.github.com>
Signed-off-by: Suguna Velury <178320438+sugunav14@users.noreply.github.com>
Signed-off-by: Suguna Velury <178320438+sugunav14@users.noreply.github.com>
Replace the hand-rolled _conversion_rules (regex renames + a MergeModulelist/Concatenate whitelist that drifts as HF adds ops) with thin adapters that drive transformers' own conversion APIs: - _conversion_plan splits HF's conversion mapping instead of translating it into our own rule language - _resolve_target uses rename_source_key for target-name resolution (bucketing) before any tensor is read - _convert_keys drives each WeightConverter's ops via op.convert, so any conversion op works; only a many-to-one shape guard remains Guarded behind the conversion engine's presence: pre-v5 transformers keeps legacy _checkpoint_conversion_mapping renames with no fusion. Unit test updated to validate the op-driven path against a real tiny Qwen3-MoE (multi-source gate_up + single-source down_proj fusion). Co-Authored-By: Claude Opus 4.8 (1M context) <noreply@anthropic.com> Signed-off-by: Suguna Velury <178320438+sugunav14@users.noreply.github.com>
Co-authored-by: realAsma <86726418+realAsma@users.noreply.github.com> Signed-off-by: sugunav14 <178320438+sugunav14@users.noreply.github.com>
sugunav14
force-pushed
the
svelury/fsdp2-refactor
branch
from
July 25, 2026 18:19
8311d61 to
0db57e5
Compare
get_max_batch_size now calls the module (model(...)) so FSDP2 hooks fire, not model.forward. Route the Mock's __call__ via side_effect so the OOM-retry regression test records the probed batch sizes again. Co-Authored-By: Claude Opus 4.8 (1M context) <noreply@anthropic.com> Signed-off-by: Suguna Velury <178320438+sugunav14@users.noreply.github.com>
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
6 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.
What does this PR do?
Type of change: New feature (+ refactor/removal of the legacy multi-node path)
Consolidates multi-node FSDP2 post-training quantization into the standard hf_ptq.py entry point behind a --use_fsdp2 flag, and removes the separate
Accelerate-based multinode_ptq.py script and fsdp2.yaml config. FSDP2 PTQ is now launched with torchrun and supports single-node multi-GPU and multi-node out
of the box.
Highlights:
offload of decoder shards between forwards.
broadcasts them so every rank can shard its slice. Non-decoder weights (embed/lm_head/norm) are read on rank 0 and broadcast.
broadcast_from_rank0); rank 0 writes files, other ranks sync on a barrier.
works for CPU-offloaded shards.
v1 scope: standard causal-LM checkpoints only. --use_fsdp2 raises NotImplementedError for VILA, multimodal/VL, pack-quantized/compressed-tensors, speculative decoding, MTP, auto-quantize, and sparsity.
Design Notes
AutoModelForCausalLM.from_pretrained(device_map="auto" | "cpu") and accelerate.load_checkpoint_in_model both load the full checkpoint on every rank from disk (per-rank CPU peak ≈ full model size). For 70B that's ~140 GiB/rank; for 200B+ it OOMs the node before sharding can run. parallel_load_and_prepare_fsdp2 round-robins decoder layers across ranks so each rank reads only ~model_size / world_size from disk, then per-layer broadcasts to peers. Per-rank CPU peak is bounded by the largest single layer + transient broadcast, not the full model. This is what makes 200B+ FSDP2 PTQ feasible on commodity nodes; HF/accelerate's existing entry points don't expose this composition.
The deleted multinode_ptq.py ran on accelerate launch --config_file fsdp2.yaml and duplicated hf_ptq.py's load → calibrate → export path. Two consequences:
The new code path unifies under hf_ptq.py --use_fsdp2. Going direct to FSDP2 primitives (fully_shard, CPUOffloadPolicy, MixedPrecisionPolicy) instead of
routing through accelerate's wrappers buys:
The "custom stack" is intentionally thin: fsdp2_wrap is a 5-line wrapper over fully_shard; the rest is pure torch.distributed. We're not reimplementing FSDP2 —just composing the public PyTorch surface directly rather than via accelerate's adapter.
Both implement the same trick: mtq.quantize unwraps the FSDP module before calling the user's forward_loop, and forwarding through the unwrapped module bypasses FSDP2's pre/post-forward hooks (no all-gather → broken calibration). Both capture the outer wrapped model and forward through it instead. This PR extracts the pattern into fsdp_aware_forward_loop (in distributed.py) as the canonical helper. The QLoRA path (transformers_trainer.py:_quantize_model) keeps its inlined version this PR because the QLoRA forward loop has training-specific quirks (batch shape, loss accumulation, gradient flow) that need a careful pass to share the helper cleanly. The TODO in the helper's docstring marks the consolidation target.
Usage
Testing
round-trip under CPUOffloadPolicy).
Before your PR is "Ready for review"
Make sure you read and follow Contributor guidelines and your commits are signed (
git commit -s -S).Make sure you read and follow the Security Best Practices (e.g. avoiding hardcoded
trust_remote_code=True,torch.load(..., weights_only=False),pickle, etc.).CONTRIBUTING.md: N/AAdditional Information
Summary by CodeRabbit
New Features
Documentation
Removed
Tests