FSDP2 calibration with hf_ptq.py [2/2] by sugunav14 · Pull Request #2013 · NVIDIA/Model-Optimizer

sugunav14 · 2026-07-24T00:00:12Z

What does this PR do?

Type of change: ?

Usage

# Add a code snippet demonstrating how to use this

Testing

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.).

Is this change backward compatible?: ✅ / ❌ / N/A
If you copied code from any other sources or added a new PIP dependency, did you follow guidance in CONTRIBUTING.md: ✅ / ❌ / N/A
Did you write any new necessary tests?: ✅ / ❌ / N/A
Did you update Changelog?: ✅ / ❌ / N/A
Did you get Claude approval on this PR?: ✅ / ❌ / N/A

Additional Information

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>

Distribute quantized weight-packing across FSDP2 ranks: each rank packs its own Shard(0) row-slice in place via a meta-skeleton FSDPParam rebuild (_rebuild_fsdp_param_from_shard) instead of every rank packing the full model. MoE fused experts are packed keep-fused and split into per-expert deployment keys at write time; scale buffers are re-wrapped as Shard(0) DTensors with explicit global shape/stride so unshard round-trips on uneven splits. The FSDP2 export gather is mode-aware (MODELOPT_DISABLE_SHARD_LOCAL escape hatch) with EXPORT_PHASE_TIMING markers for prep/pack/gather. Supersedes the streaming Option A export path on this branch. Validated on Qwen3-235B and NemotronH-120B: bit-exact vs the standard device_map export at world=4 and world=8; packing 148.7s -> 88.4s (1.68x) from 4 -> 8 ranks. Co-Authored-By: Claude Opus 4.8 (1M context) <noreply@anthropic.com> Signed-off-by: Suguna Velury <178320438+sugunav14@users.noreply.github.com>

Each FSDP2 rank writes its owned decoder-layer units concurrently instead of a whole-model full_tensor gather + rank-0 save_pretrained. export_hf_checkpoint forks to _parallel_write_hf_checkpoint (escape hatches MODELOPT_DISABLE_SHARD_LOCAL / MODELOPT_DISABLE_PARALLEL_WRITE); _export_transformers_checkpoint gains a _pack_only early-return so both paths share prep+pack. Per-unit gather (full_tensor) -> HF split_torch_state_dict_into_shards writer -> gather_object index merge. Host memory bounded to ~one chunk of world units; writes parallelize across ranks. Validated: byte-exact vs single-writer reference on tiny Qwen3-MoE + NemotronH (w2) and Qwen3-235B (w8, 146,361 keys); output HF-loadable (get_checkpoint_shard_files). Co-Authored-By: Claude Opus 4.8 (1M context) <noreply@anthropic.com> Signed-off-by: Suguna Velury <178320438+sugunav14@users.noreply.github.com>

Drop all os.environ-gated code from the shard-local + parallel-write path, leaving the validated default behavior as the single path: - EXPORT_PHASE_TIMING phase-timing markers (unified_export_hf) - MODELOPT_DISABLE_SHARD_LOCAL escape hatch + its legacy get_model_state_dict full-pack gather branch (now dead) and _use_shard_local gate - MODELOPT_DISABLE_PARALLEL_WRITE escape hatch on the export fork - SHARD_LOCAL_DEBUG (_dbg) per-rank logging (core_utils) FSDP2 distributed export now unconditionally uses shard-local pack + parallel write. Co-Authored-By: Claude Opus 4.8 (1M context) <noreply@anthropic.com> Signed-off-by: Suguna Velury <178320438+sugunav14@users.noreply.github.com>

copy-pr-bot · 2026-07-24T00:00:16Z

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 · 2026-07-24T00:00:19Z

Important

Review skipped

Draft detected.

Please check the settings in the CodeRabbit UI or the .coderabbit.yaml file in this repository. To trigger a single review, invoke the @coderabbitai review command.

⚙️ Run configuration

Configuration used: Path: .coderabbit.yaml

Review profile: CHILL

Plan: Enterprise

Run ID: 4c288808-db00-4e39-b650-b2a9430d043d

You can disable this status message by setting the reviews.review_status to false in the CodeRabbit configuration file.

Use the checkbox below for a quick retry:

🔍 Trigger review

✨ Finishing Touches

🧪 Generate unit tests (beta)

Create PR with unit tests
Commit unit tests in branch svelury/parallel-export

_{Comment @coderabbitai help to get the list of available commands.}

sugunav14 added 30 commits July 19, 2026 18:24

initial refactor

4343c9d

Signed-off-by: Suguna Velury <178320438+sugunav14@users.noreply.github.com>

cleanup

d78b547

Signed-off-by: Suguna Velury <178320438+sugunav14@users.noreply.github.com>

prototype, untested

a3d1f8b

Signed-off-by: Suguna Velury <178320438+sugunav14@users.noreply.github.com>

tested layerwise + non layerwise

8d39037

Signed-off-by: Suguna Velury <178320438+sugunav14@users.noreply.github.com>

added and tested CPU offloading policy

ce66d36

Signed-off-by: Suguna Velury <178320438+sugunav14@users.noreply.github.com>

added process parallel loading

9b90b4e

Signed-off-by: Suguna Velury <178320438+sugunav14@users.noreply.github.com>

claude self-review comments

a86b744

Signed-off-by: Suguna Velury <178320438+sugunav14@users.noreply.github.com>

coderabbit reviews

d9ad10f

Signed-off-by: Suguna Velury <178320438+sugunav14@users.noreply.github.com>

modelopt bot review

08dcd95

Signed-off-by: Suguna Velury <178320438+sugunav14@users.noreply.github.com>

clean up (1)

b783aff

Signed-off-by: Suguna Velury <178320438+sugunav14@users.noreply.github.com>

refactor

f9b4e45

Signed-off-by: Suguna Velury <178320438+sugunav14@users.noreply.github.com>

claude self review

cec0305

Signed-off-by: Suguna Velury <178320438+sugunav14@users.noreply.github.com>

modelopt bot review

c4f34a1

Signed-off-by: Suguna Velury <178320438+sugunav14@users.noreply.github.com>

minor update

d2e4c1a

Signed-off-by: Suguna Velury <178320438+sugunav14@users.noreply.github.com>

updated README

6f786f9

Signed-off-by: Suguna Velury <178320438+sugunav14@users.noreply.github.com>

removed load_and_prepare_fsdp2 wrapper

5906bf6

Signed-off-by: Suguna Velury <178320438+sugunav14@users.noreply.github.com>

PR review comments

5eb2c25

Signed-off-by: Suguna Velury <178320438+sugunav14@users.noreply.github.com>

modified parallel read to do one broadcast per rank instead of one pe…

1bf9ebb

…r param Signed-off-by: Suguna Velury <178320438+sugunav14@users.noreply.github.com>

PR review addressed

86e05cc

Signed-off-by: Suguna Velury <178320438+sugunav14@users.noreply.github.com>

perf measurements, revert later

b549e62

Signed-off-by: Suguna Velury <178320438+sugunav14@users.noreply.github.com>

refactor + perf measured

88508ba

Signed-off-by: Suguna Velury <178320438+sugunav14@users.noreply.github.com>

cleanup

9bfecd1

Signed-off-by: Suguna Velury <178320438+sugunav14@users.noreply.github.com>

shard_root=True fix

0831244

Signed-off-by: Suguna Velury <178320438+sugunav14@users.noreply.github.com>

added slurm example and updated README.md

4a83d15

Signed-off-by: Suguna Velury <178320438+sugunav14@users.noreply.github.com>

code quality check

ca924ff

Signed-off-by: Suguna Velury <178320438+sugunav14@users.noreply.github.com>

fix unit tests

a9cb58a

Signed-off-by: Suguna Velury <178320438+sugunav14@users.noreply.github.com>

updated unit test for shard_root=True default

88a9145

Signed-off-by: Suguna Velury <178320438+sugunav14@users.noreply.github.com>

export fix

8abf326

Signed-off-by: Suguna Velury <178320438+sugunav14@users.noreply.github.com>

updated agent skills with recipe

b336a64

Signed-off-by: Suguna Velury <178320438+sugunav14@users.noreply.github.com>

updated dedup comment

f449c10

Signed-off-by: Suguna Velury <178320438+sugunav14@users.noreply.github.com>

sugunav14 and others added 3 commits July 23, 2026 23:42

sugunav14 changed the title ~~Svelury/parallel export~~ FSDP2 calibration with hf_ptq.py [2/2] Jul 24, 2026

Provide feedback

Saved searches

Use saved searches to filter your results more quickly

Uh oh!

FSDP2 calibration with hf_ptq.py [2/2]#2013

FSDP2 calibration with hf_ptq.py [2/2]#2013
sugunav14 wants to merge 33 commits into
mainfrom
svelury/parallel-export

sugunav14 commented Jul 24, 2026

Uh oh!

copy-pr-bot Bot commented Jul 24, 2026

Uh oh!

coderabbitai Bot commented Jul 24, 2026

Review skipped

Uh oh!

Reviewers

Assignees

Labels

Projects

Milestone

Development

Uh oh!

1 participant

Uh oh!

Conversation

sugunav14 commented Jul 24, 2026

What does this PR do?

Usage

Testing

Before your PR is "Ready for review"

Additional Information

Uh oh!

copy-pr-bot Bot commented Jul 24, 2026

Uh oh!

coderabbitai Bot commented Jul 24, 2026

Review skipped

Uh oh!

Reviewers

Assignees

Labels

Projects

Milestone

Development

Uh oh!

1 participant