Fridah/single gpu disk offload ptq#2008
Draft
Fridah-nv wants to merge 8 commits into
Draft
Conversation
G1 (offload-aware export): fix _export_transformers_checkpoint to materialize weights per decoder layer via enable_weight_access_and_writeback before serializing, avoiding silent meta-tensor corruption when the model was loaded with accelerate disk/CPU offload. Adds _has_accelerate_offload detection and _process_quantized_modules_offloaded path. Adds meta-tensor guard in _export_quantized_weight to surface silent corruption as loud error. G2 (disk-offload CLI): --offload_folder, --max_gpu_memory_gb, --max_cpu_memory_gb args wired through hf_ptq.py -> example_utils.get_model. Validation guards prevent combining offload_folder with low_memory_mode or use_seq_device_map. G3 (CPU RSS check): psutil RSS verification after model load with disk offload + max_cpu_memory_gb; warns if RSS > 110% of limit. Supporting changes: - NemotronHForCausalLM layer detection: check both backbone.layers and model.layers so the plugin works for both remote-code and native HF variants of the architecture. - New recipe nvfp4_experts_only-kv_fp8_layerwise_offload.yaml: NVFP4 W4A4 on experts, FP8 KV cache, layerwise calibration with calib_mutates_weights: false (disk-offload compatible). - DSR1 BF16 fallback: _FP8BF16Fallback.matmul shim dequants block-scaled FP8 expert weights to BF16 for calibration forward passes when the kernels package is unavailable. - Unit tests (6 CPU-only) and GPU integration tests (2 tiny-LLaMA) for the offload-aware export path. Part of OMNIML-4947. Co-Authored-By: Claude Sonnet 4.6 <noreply@anthropic.com> Signed-off-by: Fridah-nv <201670829+Fridah-nv@users.noreply.github.com>
Adds run scripts and worklog for end-to-end single-GPU disk-offload PTQ validation on two production-scale models (OMNIML-4947): - dsr1_run.sh: DSR1 642GB BF16 on 4 GPUs (fast demo, GPU-only). - dsr1_single_gpu_run.sh: DSR1 single-GPU disk-offload (data collection). - ultra550b_run.sh: Nemotron-Ultra 550B on single GPU with disk offload. - ultra550b_4gpu_run.sh: Nemotron-Ultra 550B on 4 GPUs + CPU offload (demo). - single_gpu_ptq_progress.md: worklog with run history, results, and known limitations. DSR1 4-GPU run (Run 7) passed: 61 layers in 8m25s, 42 export shards in 235s, GPU peak 158-170 GiB, CPU RSS peak 629.82 GiB. Co-Authored-By: Claude Sonnet 4.6 <noreply@anthropic.com> Signed-off-by: Fridah-nv <201670829+Fridah-nv@users.noreply.github.com>
Two fixes for the e2e validation run scripts: 1. Export PYTHONPATH=$REPO before python3 invocations so that the editable install's dev modelopt (0.46.0.dev) is found before system modelopt (0.21.0 at /usr/local/lib/python3.12/dist-packages) in subprocess sys.path. Root cause: pip install -e appends _EditableFinder to sys.meta_path (not prepends), so PathFinder finds /usr/local/modelopt first when the script dir is not the repo root. 2. --gpus now uses CSV format (0,1,2,3) matching the PR #2000 mem_monitor.py interface (--gpus takes a comma-separated string, not nargs="+"). Co-Authored-By: Claude Sonnet 4.6 <noreply@anthropic.com> Signed-off-by: Fridah-nv <201670829+Fridah-nv@users.noreply.github.com>
- Correct branch name (fridah/single-gpu-disk-offload-ptq) - Record 4 commits landed this session - Clarify mem_monitor.py source (cherry-picked from PR #2000, not hand-written) - Document PYTHONPATH editable-install root cause and fix - Note Python environment details (dev modelopt, transformers, torch versions) - Mark Ultra 550B single-GPU run as in-progress (PID 2744499) - Update --gpus format in DSR1 run command to CSV Co-Authored-By: Claude Sonnet 4.6 <noreply@anthropic.com> Signed-off-by: Fridah-nv <201670829+Fridah-nv@users.noreply.github.com>
_process_quantized_modules_offloaded previously only materialized decoder
layers. For models like NemotronH-Ultra-550B where embed_tokens, norm_f,
and lm_head are also disk-offloaded, model.state_dict() returns meta
tensors for those modules. After revert_weight_conversion_quant_aware
renames state_dict keys to hub-original names (backbone.*), transformers'
remove_tied_weights_from_state_dict calls model.get_parameter('backbone.*')
which fails because the model uses in-memory names (model.*).
Fix: after the decoder-layer loop, iterate non-decoder modules with
AlignDevicesHook and use enable_weight_access_and_writeback(writeback=False)
to collect their direct parameters. The returned state dict is now
guaranteed to have no meta tensors.
Adds a unit test that verifies embed.weight is real (not meta) in the
returned state dict when only the embedding is disk-offloaded.
Co-Authored-By: Claude Sonnet 4.6 <noreply@anthropic.com>
Signed-off-by: Fridah-nv <201670829+Fridah-nv@users.noreply.github.com>
Co-Authored-By: Claude Sonnet 4.6 <noreply@anthropic.com> Signed-off-by: Fridah-nv <201670829+Fridah-nv@users.noreply.github.com>
Run completed successfully (PID 2875977, 2026-07-21): - 15 safetensors shards, 33.2 min total wall time - Peak GPU VRAM: 161.94 GB (single GB200) - Peak CPU RSS: ~764 GiB - Calibration: 8m41s (108 layers), Export: 322.6s (15 shards) - Fix: non-decoder disk-offloaded tensors (embed, norm, lm_head) now materialized before state_dict rename to hub-original names, preventing the backbone AttributeError that crashed the first attempt. Co-Authored-By: Claude Sonnet 4.6 <noreply@anthropic.com> Signed-off-by: Fridah-nv <201670829+Fridah-nv@users.noreply.github.com>
Design for replacing the accumulate-all-then-save export pattern with a streaming shard writer. Eliminates the 764 GiB peak CPU RSS that blocks the 80 GB VRAM + 80 GB CPU RAM target from the original OMNIML-4947 ticket. Key pieces: - _StreamingShardWriter: buffers one shard at a time, flushes via safe_save_file - _postprocess_single_tensor: per-tensor extraction from postprocess_state_dict - _export_transformers_checkpoint_streaming: streaming path for offloaded models - E2E test scripts: ultra550b_80gb_run.sh, dsr1_80gb_run.sh with 80 GB budget Estimated peak: ~57 GB (Ultra 550B) and ~55 GB (DSR1), down from 764 GiB. Co-Authored-By: Claude Sonnet 4.6 <noreply@anthropic.com> Signed-off-by: Fridah-nv <201670829+Fridah-nv@users.noreply.github.com>
Contributor
|
Important Review skippedDraft detected. Please check the settings in the CodeRabbit UI or the ⚙️ Run configurationConfiguration used: Path: .coderabbit.yaml Review profile: CHILL Plan: Enterprise Run ID: You can disable this status message by setting the Use the checkbox below for a quick retry:
✨ Finishing Touches🧪 Generate unit tests (beta)
Comment |
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
1 participant
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: ?
Usage
# Add a code snippet demonstrating how to use thisTesting
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