diff --git a/.agents/skills/ptq/references/slurm-setup-ptq.md b/.agents/skills/ptq/references/slurm-setup-ptq.md index 635e5e1a4a8..117494e0074 100644 --- a/.agents/skills/ptq/references/slurm-setup-ptq.md +++ b/.agents/skills/ptq/references/slurm-setup-ptq.md @@ -63,23 +63,26 @@ pip install -U transformers --no-deps Estimate GPU count from model size and available GPU memory. `hf_ptq.py` uses `device_map="auto"` so it fills GPUs automatically — request only as many as needed. -For multi-node PTQ (200B+ params), use `examples/hf_ptq/multinode_ptq.py` with FSDP2 and accelerate: +For multi-node PTQ (200B+ params), use `hf_ptq.py --use_fsdp2` launched via `torchrun`: ```bash -accelerate launch \ - --config_file examples/hf_ptq/fsdp2.yaml \ - --num_machines $NUM_NODES \ - --num_processes $((NUM_NODES * GPUS_PER_NODE)) \ - --main_process_ip $MASTER_ADDR \ - --main_process_port $MASTER_PORT \ - --machine_rank $SLURM_PROCID \ - examples/hf_ptq/multinode_ptq.py \ +torchrun \ + --nnodes=$NUM_NODES \ + --node_rank=$SLURM_PROCID \ + --nproc_per_node=$GPUS_PER_NODE \ + --master_addr=$MASTER_ADDR \ + --master_port=$MASTER_PORT \ + examples/hf_ptq/hf_ptq.py \ --pyt_ckpt_path \ - --qformat \ - --export_path + --recipe \ + --export_path \ + --use_fsdp2 ``` -The `num_machines`, `num_processes`, `main_process_ip`, and `machine_rank` are overridden on the command line — no need to edit `fsdp2.yaml`. Only update `fsdp_transformer_layer_cls_to_wrap` in the YAML if the model uses a non-default decoder layer class. +`--recipe` is the maintained entry point (e.g. `general/ptq/nvfp4_default-kv_fp8_cast`); see the +Recipe-based Quantization section of `examples/hf_ptq/README.md` for the format and built-in names. + +When the per-rank decoder shard approaches GPU capacity (200B+ at low rank count), either add more nodes (more ranks → smaller shard per rank) or add `--cpu_offload`. Layer detection is automatic; no YAML config needed. Use the multi-node template from `skills/common/slurm-setup.md` section 4 as the job script wrapper. diff --git a/examples/hf_ptq/README.md b/examples/hf_ptq/README.md index 3a7380a9f9e..376cdff62e7 100755 --- a/examples/hf_ptq/README.md +++ b/examples/hf_ptq/README.md @@ -428,33 +428,37 @@ mtq.calibrate(model, algorithm="max", forward_loop=calibrate_loop) ## Multi-Node Post-Training Quantization with FSDP2 -ModelOpt enables quantization of LLMs across multiple GPU nodes using various quantization formats. It leverages HuggingFace's Accelerate library and FSDP2 for distributed model sharding and calibration. +ModelOpt enables quantization of LLMs across multiple GPU nodes using FSDP2 for distributed model sharding and calibration, exposed via the `--use_fsdp2` flag on the standard `hf_ptq.py` entry point. ### Usage -For distributed execution across multiple nodes, use the `accelerate` library. A template configuration file (`fsdp2.yaml`) is provided and can be customized for user specific requirements. +#### Slurm (recommended) -On each node run the following command: +Slurm orchestrates launching the job on every node for you, so this is the easiest way to run a multi-node PTQ. A ready-to-run example that quantizes Nemotron-3-Super to NVFP4 is provided in [`slurm/multinode_fsdp2_ptq.slurm`](./slurm/multinode_fsdp2_ptq.slurm). Edit the `CONFIG` block (container image, model path, export path, recipe) and submit: ```bash -accelerate launch --config_file fsdp2.yaml \ - --num_machines= \ - --machine_rank= \ - --main_process_ip= \ - --main_process_port= \ - --fsdp_transformer_layer_cls_to_wrap= - multinode_ptq.py \ +sbatch --nodes=2 slurm/multinode_fsdp2_ptq.slurm +``` + +#### Manual (run on each node) + +Without Slurm, start `torchrun` on every node yourself: + +```bash +torchrun \ + --nnodes= --node_rank= \ + --master_addr= --master_port= \ + --nproc_per_node= \ + hf_ptq.py \ --pyt_ckpt_path \ - --qformat \ - --kv_cache_qformat \ + --recipe general/ptq/nvfp4_default-kv_fp8_cast \ --batch_size \ --calib_size \ - --dataset \ --export_path \ - --trust_remote_code + --use_fsdp2 ``` -The exported checkpoint can be deployed using TensorRT-LLM/ vLLM/ SGLang. For more details refer to the [deployment section](#deployment) of this document. +See [Recipe-based Quantization](#recipe-based-quantization) for the recipe format and built-in recipe names. The exported checkpoint can be deployed using TensorRT-LLM/ vLLM/ SGLang. For more details refer to the [deployment section](#deployment) of this document. > *Performance Note: FSDP2 is designed for training workloads and may result in longer calibration and export times. For faster calibration, maximize the batch size based on available GPU memory and choose the right number of GPUs to avoid unnecessary communication.* diff --git a/examples/hf_ptq/example_utils.py b/examples/hf_ptq/example_utils.py index 83a54849110..8c8ca9f988f 100755 --- a/examples/hf_ptq/example_utils.py +++ b/examples/hf_ptq/example_utils.py @@ -48,11 +48,55 @@ except ImportError: snapshot_download = None +from modelopt.torch.utils import distributed as dist_utils + logger = logging.getLogger(__name__) SPECULATIVE_MODEL_LIST = ["Eagle", "Medusa"] +def setup_distributed_args(args): + """Set ``args.rank``/``world_size``/``device``/``is_main`` (single-process if FSDP2 off).""" + if getattr(args, "use_fsdp2", False): + dist_utils.setup() + args.rank = dist_utils.rank() + args.world_size = dist_utils.size() + args.device = torch.device(f"cuda:{dist_utils.local_rank()}") + args.is_main = args.rank == 0 + else: + args.rank = 0 + args.world_size = 1 + args.is_main = True + + +def cleanup_distributed(args): + """Destroy the process group if ``--use_fsdp2`` set it up.""" + if getattr(args, "use_fsdp2", False): + dist_utils.cleanup() + + +def validate_fsdp2_supported(args, config): + """Raise ``NotImplementedError`` for model/CLI combos the FSDP2 path doesn't support yet.""" + issues = [] + if "vila" in args.pyt_ckpt_path.lower(): + issues.append("VILA (custom builder + non-standard layer layout)") + if is_nemotron_vl(config) or _is_multimodal_config(config): + issues.append("multimodal / VL models (decoder layers not auto-detectable)") + if getattr(config, "quantization_config", None) is not None: + issues.append("pack-quantized / compressed-tensors checkpoints") + if getattr(args, "specdec_offline_dataset", None) is not None: + issues.append("speculative decoding (--specdec_offline_dataset)") + if getattr(args, "low_memory_mode", False): + issues.append("--low_memory_mode (redundant with FSDP2)") + + if issues: + raise NotImplementedError( + "--use_fsdp2 does not support:\n - " + + "\n - ".join(issues) + + "\nRemove --use_fsdp2 or use a standard causal-LM checkpoint." + ) + + def run_nemotron_vl_preview( full_model, tokenizer, @@ -372,6 +416,19 @@ def _apply_to_model_state_dict( return out_state_dict +def mtp_layer_prefixes_from_checkpoint(model_path: str) -> list[str]: + """MTP exclude-prefixes from a checkpoint's safetensors index (``[]`` if none); reads no tensors. + + Local-index-only, matching :func:`load_mtp_weights`, so detection and re-attach stay in sync. + """ + index_file = Path(model_path) / "model.safetensors.index.json" + if not index_file.exists(): + return [] + weight_map = json.load(open(index_file))["weight_map"] + mtp_keys = [k for k, v in weight_map.items() if "mtp" in k or "mtp" in v] + return list(_keys_to_prefixes(mtp_keys)) + + def load_mtp_weights( model: torch.nn.Module, model_path: str ) -> tuple[list[str], dict[str, torch.Tensor]]: diff --git a/examples/hf_ptq/fsdp2.yaml b/examples/hf_ptq/fsdp2.yaml deleted file mode 100644 index 09977835561..00000000000 --- a/examples/hf_ptq/fsdp2.yaml +++ /dev/null @@ -1,30 +0,0 @@ -# ============================================================================= -# FSDP Configuration for running LLM PTQ on multinode setup. This file is consumed by examples/hf_ptq/multinode_ptq.py -# ============================================================================= - -compute_environment: LOCAL_MACHINE -debug: false -distributed_type: FSDP -downcast_bf16: 'no' -enable_cpu_affinity: false -fsdp_config: - fsdp_activation_checkpointing: false - fsdp_auto_wrap_policy: TRANSFORMER_BASED_WRAP - fsdp_cpu_ram_efficient_loading: true - fsdp_offload_params: false - fsdp_reshard_after_forward: true - fsdp_state_dict_type: FULL_STATE_DICT - fsdp_transformer_layer_cls_to_wrap: LlamaDecoderLayer - fsdp_use_orig_params: true - fsdp_version: 2 -machine_rank: 0 -main_training_function: main -mixed_precision: 'no' -num_machines: 2 -num_processes: 16 -rdzv_backend: c10d -same_network: true -tpu_env: [] -tpu_use_cluster: false -tpu_use_sudo: false -use_cpu: false diff --git a/examples/hf_ptq/hf_ptq.py b/examples/hf_ptq/hf_ptq.py index 8dcc78afa27..dae5ebdcfec 100755 --- a/examples/hf_ptq/hf_ptq.py +++ b/examples/hf_ptq/hf_ptq.py @@ -15,6 +15,7 @@ import argparse import copy +import os import random import time import warnings @@ -29,6 +30,7 @@ from example_utils import ( _resolve_model_path, build_quant_cfg, + cleanup_distributed, copy_custom_model_files, create_vlm_calibration_loop, get_model, @@ -37,9 +39,12 @@ is_enc_dec, is_nemotron_vl, load_mtp_weights, + mtp_layer_prefixes_from_checkpoint, needs_checkpoint_path_update, resolve_checkpoint_dir, run_nemotron_vl_preview, + setup_distributed_args, + validate_fsdp2_supported, ) from torch.utils.data import DataLoader from transformers import ( @@ -82,6 +87,7 @@ get_supported_datasets, ) from modelopt.torch.utils.memory_monitor import launch_memory_monitor +from modelopt.torch.utils.plugins.model_load_utils import parallel_load_and_prepare_fsdp2 from modelopt.torch.utils.speech_dataset_utils import get_speech_dataset_dataloader from modelopt.torch.utils.vlm_dataset_utils import get_vlm_dataset_dataloader @@ -251,6 +257,10 @@ def make_calib_dataloader( max_sample_length=args.calib_seq, device=device, include_labels=include_labels, + distributed=args.use_fsdp2, + sampler_kwargs=( + {"num_replicas": args.world_size, "rank": args.rank} if args.use_fsdp2 else None + ), ) return calib_dataloader, first_text_speech_dataset @@ -413,6 +423,11 @@ def auto_quantize( "Auto Quantization is not supported for pipeline parallel size > 1" ) + assert not (args.auto_quantize_bits and args.use_fsdp2), ( + "Auto Quantization is not supported with --use_fsdp2: mtq.auto_quantize " + "is invoked after FSDP2 loading has frozen every parameter." + ) + inputs = _mtq_inputs_from_auto_quantize_config(aq_config, args) # base-model lm_head handling (mirrors the CLI helper) @@ -501,7 +516,27 @@ def _recipe_is_auto_quantize(recipe: str | None) -> bool: def load_model(args: argparse.Namespace): # If low memory mode is enabled, we compress the model while loading the HF checkpoint. calibration_only = False - if args.specdec_offline_dataset is not None or not args.low_memory_mode: + if args.use_fsdp2: + hf_config = AutoConfig.from_pretrained( + args.pyt_ckpt_path, trust_remote_code=args.trust_remote_code + ) + validate_fsdp2_supported(args, hf_config) + full_model = parallel_load_and_prepare_fsdp2( + args.pyt_ckpt_path, + args.device, + args.rank, + args.world_size, + trust_remote_code=args.trust_remote_code, + cpu_offload=args.cpu_offload, + attn_implementation=args.attn_implementation, + hf_config=hf_config, + ) + # The FSDP2 loader drops MTP weights (re-attached BF16 at export); flag their prefixes now + # so the pre-quant exclusion below skips any MTP module from_config did build. + mtp_prefixes = mtp_layer_prefixes_from_checkpoint(args.pyt_ckpt_path) + if mtp_prefixes: + full_model._mtp_layer_prefixes = mtp_prefixes + elif args.specdec_offline_dataset is not None or not args.low_memory_mode: full_model = get_model( args.pyt_ckpt_path, args.device, @@ -536,9 +571,12 @@ def load_model(args: argparse.Namespace): model_type = get_model_type(full_model) - device = full_model.device - if hasattr(full_model, "model"): - device = full_model.model.device + if args.use_fsdp2: + device = args.device + else: + device = full_model.device + if hasattr(full_model, "model"): + device = full_model.model.device processor = None tokenizer = None language_model = full_model @@ -832,7 +870,6 @@ def export_quantized( mtp_layer_prefixes, mtp_state_dict = load_mtp_weights( full_model, args.pyt_ckpt_path ) - if mtp_layer_prefixes: full_model._mtp_layer_prefixes = mtp_layer_prefixes @@ -853,18 +890,21 @@ def export_quantized( tokenizer.padding_side = default_padding_side if default_pad_token is not None: tokenizer.pad_token = default_pad_token - tokenizer.save_pretrained(export_path) + if args.is_main: + tokenizer.save_pretrained(export_path) # Copy custom model files (Python files and JSON configs) if trust_remote_code is used. # This must run AFTER tokenizer.save_pretrained() so original tokenizer files # from the source checkpoint take precedence over regenerated ones (which may # differ in format due to newer transformers versions). - copy_custom_model_files(args.pyt_ckpt_path, export_path, args.trust_remote_code) + if args.is_main: + copy_custom_model_files(args.pyt_ckpt_path, export_path, args.trust_remote_code) end_time = time.time() - print( - f"Quantized model exported to: {export_path}. Total time used {end_time - start_time}s" - ) + if args.is_main: + print( + f"Quantized model exported to: {export_path}. Total time used {end_time - start_time}s" + ) def pre_quantize( @@ -963,7 +1003,7 @@ def post_quantize( ) return - if args.verbose: + if args.verbose and args.is_main: try: mtq.print_quant_summary(full_model, args.export_path) save_expert_token_count_table(full_model, args.export_path) @@ -1418,6 +1458,20 @@ def parse_args() -> argparse.Namespace: default=False, action="store_true", ) + parser.add_argument( + "--use_fsdp2", + action="store_true", + help=( + "Run calibration under PyTorch FSDP2 (requires torchrun); takes precedence over " + "--use_seq_device_map. v1: standard causal-LM only (no VILA / pack-quantized / " + "speculative / auto-quantize / sparsity / VLM / MTP)." + ), + ) + parser.add_argument( + "--cpu_offload", + action="store_true", + help="With --use_fsdp2, keep decoder shards on CPU between forwards (frees GPU memory, adds PCIe traffic).", + ) parser.add_argument( "--verbose", help="Print verbose output (e.g. quantization summary). Disable by --no-verbose.", @@ -1548,6 +1602,19 @@ def parse_args() -> argparse.Namespace: "--low_memory_mode does not support --recipe or AutoQuantize (--auto_quantize_bits); " "the low-memory loader initializes quantizers from --qformat/--kv_cache_qformat." ) + if args.use_fsdp2 and args.use_seq_device_map: + warnings.warn("--use_seq_device_map is ignored when --use_fsdp2 is set.") + args.use_seq_device_map = False + if args.use_fsdp2 and os.environ.get("RANK") is None: + parser.error("--use_fsdp2 requires launching with torchrun") + if args.cpu_offload and not args.use_fsdp2: + parser.error("--cpu_offload requires --use_fsdp2") + if args.use_fsdp2 and args.sparsity_fmt != "dense": + parser.error(f"--use_fsdp2 does not support --sparsity_fmt {args.sparsity_fmt}.") + if args.use_fsdp2 and args.vllm_fakequant_export: + parser.error("--use_fsdp2 does not support --vllm_fakequant_export.") + if args.use_fsdp2 and args.cast_mxfp4_to_nvfp4: + parser.error("--use_fsdp2 does not support --cast_mxfp4_to_nvfp4.") return args @@ -1559,31 +1626,16 @@ def main(args: argparse.Namespace): random.seed(RAND_SEED) np.random.seed(RAND_SEED) - # launch a memory monitor to read the currently used GPU memory. - launch_memory_monitor() + setup_distributed_args(args) - # Force eager execution for all model types. - torch.compiler.set_stance("force_eager") + try: + # launch a memory monitor to read the currently used GPU memory. + launch_memory_monitor() - ( - full_model, - language_model, - model_type, - calibration_only, - processor, - tokenizer, - default_padding_side, - default_pad_token, - device, - ) = load_model(args) + # Force eager execution for all model types. + torch.compiler.set_stance("force_eager") - if args.sparsity_fmt != "dense": - # Sparse - sparsity_main(args, full_model, tokenizer, device) - else: - # Quantize - quantize_main( - args, + ( full_model, language_model, model_type, @@ -1593,7 +1645,27 @@ def main(args: argparse.Namespace): default_padding_side, default_pad_token, device, - ) + ) = load_model(args) + + if args.sparsity_fmt != "dense": + # Sparse + sparsity_main(args, full_model, tokenizer, device) + else: + # Quantize + quantize_main( + args, + full_model, + language_model, + model_type, + calibration_only, + processor, + tokenizer, + default_padding_side, + default_pad_token, + device, + ) + finally: + cleanup_distributed(args) if __name__ == "__main__": diff --git a/examples/hf_ptq/multinode_ptq.py b/examples/hf_ptq/multinode_ptq.py deleted file mode 100644 index 12e6c04e535..00000000000 --- a/examples/hf_ptq/multinode_ptq.py +++ /dev/null @@ -1,369 +0,0 @@ -# SPDX-FileCopyrightText: Copyright (c) 2024 NVIDIA CORPORATION & AFFILIATES. All rights reserved. -# SPDX-License-Identifier: Apache-2.0 -# -# Licensed under the Apache License, Version 2.0 (the "License"); -# you may not use this file except in compliance with the License. -# You may obtain a copy of the License at -# -# http://www.apache.org/licenses/LICENSE-2.0 -# -# Unless required by applicable law or agreed to in writing, software -# distributed under the License is distributed on an "AS IS" BASIS, -# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. -# See the License for the specific language governing permissions and -# limitations under the License. - -"""Multi-node PTQ (Post-Training Quantization) with FSDP2 support.""" - -import argparse -import json -import os -import random -import time -import warnings -from pathlib import Path - -import numpy as np -import torch -import torch.nn as nn -from accelerate import Accelerator -from example_utils import build_quant_cfg, get_tokenizer -from tqdm import tqdm -from transformers import AutoModelForCausalLM, PreTrainedTokenizer, PreTrainedTokenizerFast - -import modelopt.torch.opt as mto -import modelopt.torch.quantization as mtq -from modelopt.recipe.presets import KV_CACHE_NONE, KV_QUANT_CFG_CHOICES, QUANT_CFG_CHOICES -from modelopt.torch.export import get_model_type -from modelopt.torch.export.convert_hf_config import convert_hf_quant_config_format -from modelopt.torch.export.unified_export_hf import _export_transformers_checkpoint -from modelopt.torch.quantization.config import need_calibration -from modelopt.torch.quantization.utils import patch_fsdp_mp_dtypes -from modelopt.torch.utils.dataset_utils import get_dataset_dataloader, get_supported_datasets - -# Constants -RAND_SEED = 1234 - - -# Enable HuggingFace checkpointing -mto.enable_huggingface_checkpointing() - - -def parse_args(): - """Parse command line arguments.""" - parser = argparse.ArgumentParser(description="Multi-node post-training quantization with FSDP2") - - parser.add_argument( - "--pyt_ckpt_path", - required=True, - help="Path to PyTorch checkpoint", - ) - parser.add_argument( - "--qformat", - default="fp8", - choices=list(QUANT_CFG_CHOICES), - help="Quantization format", - ) - parser.add_argument( - "--kv_cache_qformat", - default="fp8", - choices=[KV_CACHE_NONE, *KV_QUANT_CFG_CHOICES], - help="KV cache quantization format", - ) - parser.add_argument( - "--batch_size", - type=int, - default=1, - help="Batch size for calibration", - ) - parser.add_argument( - "--calib_size", - type=str, - default="512", - help="Comma-separated list of calibration sizes per dataset", - ) - parser.add_argument( - "--dataset", - help=( - f"name of a dataset, or a comma separated list of datasets. " - f"dataset choices are {get_supported_datasets()}" - ), - type=str, - default=None, - ) - parser.add_argument( - "--export_path", - default="exported_model", - help="Directory to export the quantized model", - ) - parser.add_argument( - "--trust_remote_code", - action="store_true", - help="Trust remote code for HuggingFace models", - ) - parser.add_argument("--awq_block_size", default=0, type=int) - - args = parser.parse_args() - - # Parse comma-separated lists - args.dataset = args.dataset.split(",") if args.dataset else None - args.calib_size = [int(x) for x in args.calib_size.split(",")] - - return args - - -def load_and_prepare_model( - model_path: str, - calib_dataloader: torch.utils.data.DataLoader, - accelerator: Accelerator, - trust_remote_code: bool = False, -) -> tuple[nn.Module, str, list[str], torch.utils.data.DataLoader]: - """Load model and prepare it for FSDP2 distributed execution. - - Args: - model_path: Path to the HuggingFace model - calibration_dataloader: Calibration dataloader to be sharded for calibration - accelerator: Accelerate's Accelerator instance - trust_remote_code: Whether to trust remote code - - Returns: - Tuple of (prepared_model, model_type, original_architectures, calibration_dataloader) - """ - model = AutoModelForCausalLM.from_pretrained( - model_path, dtype="auto", trust_remote_code=trust_remote_code - ) - model.eval() - model_type = get_model_type(model) - # Need the original architectures for export - # FSDP prefix is added to the architectures for FSDP2 wrapped models - original_architectures = model.config.architectures - - # FSDP2 requires an optimizer to be prepared together with the model - dummy_optimizer = torch.optim.SGD(model.parameters(), lr=0.0) - model, _, calibration_dataloader = accelerator.prepare(model, dummy_optimizer, calib_dataloader) - - return model, model_type, original_architectures, calibration_dataloader - - -def create_calibration_dataloader( - tokenizer: PreTrainedTokenizer | PreTrainedTokenizerFast, - dataset_names: list[str], - calib_sizes: list[int], - batch_size: int, -) -> torch.utils.data.DataLoader: - """Create calibration dataloader from dataset. - - Args: - tokenizer: HuggingFace tokenizer - dataset_names: List of dataset names (defaults to cnn_dailymail) - calib_sizes: Number of samples for each dataset - batch_size: Batch size for calibration - - Returns: - DataLoader for calibration - """ - - return get_dataset_dataloader( - dataset_name=dataset_names, - tokenizer=tokenizer, - batch_size=batch_size, - num_samples=calib_sizes, - device=None, # Keep data on CPU, calibration loop handles device transfer - include_labels=False, - ) - - -def create_fsdp2_calibration_loop( - model: nn.Module, - dataloader: torch.utils.data.DataLoader, - accelerator: Accelerator, -): - """Create calibration loop compatible with FSDP2. - - For FSDP2, we need to use the outer FSDP-wrapped model instead of - the parameter passed by mtq.quantize to properly handle DTensor. - - Args: - model: FSDP2-wrapped model - dataloader: Calibration dataloader - accelerator: Accelerator instance for device management - - Returns: - Calibration function compatible with mtq.quantize - """ - - def calibrate(unwrapped_model): - """Calibration loop that uses the FSDP-wrapped model.""" - for batch in tqdm(dataloader, desc="Calibrating"): - if isinstance(batch, dict): - batch = { - k: v.to(accelerator.device) if isinstance(v, torch.Tensor) else v - for k, v in batch.items() - } - # Use outer model (FSDP-wrapped), not the parameter - # Important: We should forward pass using the unwrapped model - # mtq.quantize will unwrap the model & pass to the forward_loop - model(**batch) - - return calibrate - - -def export_model( - model: nn.Module, - accelerator: Accelerator, - export_path: str | Path, - architectures: list[str], -): - """Export quantized model to HuggingFace format. - - Args: - model: Quantized model - accelerator: Accelerator instance for state dict gathering - export_path: Directory to export model to - """ - export_dir = Path(export_path) - export_dir.mkdir(parents=True, exist_ok=True) - - post_state_dict, hf_quant_config = _export_transformers_checkpoint( - model, torch.bfloat16, accelerator=accelerator - ) - - if accelerator.is_main_process: - # Save hf_quant_config.json for backward compatibility - with open(f"{export_dir}/hf_quant_config.json", "w") as file: - json.dump(hf_quant_config, file, indent=4) - - hf_quant_config = convert_hf_quant_config_format(hf_quant_config) - - # Save model - model.save_pretrained(export_dir, state_dict=post_state_dict, save_modelopt_state=False) - - original_config = f"{export_dir}/config.json" - config_data = {} - - with open(original_config) as file: - config_data = json.load(file) - - config_data["quantization_config"] = hf_quant_config - # Update config architectures to use original architectures that does not have FSDP prefix - config_data["architectures"] = architectures - - with open(original_config, "w") as file: - json.dump(config_data, file, indent=4) - - -def main(args): - """Main quantization workflow.""" - # Validate GPU availability - if not torch.cuda.is_available(): - raise OSError("GPU is required for quantization.") - - # Validate quantization format - if args.qformat not in QUANT_CFG_CHOICES: - raise ValueError( - f"Quantization format {args.qformat} not supported. Choose from: {list(QUANT_CFG_CHOICES)}" - ) - - # Set random seeds - random.seed(RAND_SEED) - np.random.seed(RAND_SEED) - torch.manual_seed(RAND_SEED) - - # Initialize accelerator - accelerator = Accelerator() - - print(f"Rank: {os.environ.get('RANK', 'Not set')}") - print(f"World Size: {os.environ.get('WORLD_SIZE', 'Not set')}") - print(f"Local Rank: {os.environ.get('LOCAL_RANK', 'Not set')}") - - # Load tokenizer - tokenizer = get_tokenizer(args.pyt_ckpt_path, trust_remote_code=args.trust_remote_code) - default_padding_side = tokenizer.padding_side - tokenizer.padding_side = "left" # Left padding for better calibration - - # Set default dataset if not provided - if args.dataset is None: - args.dataset = ["cnn_dailymail", "nemotron-post-training-dataset-v2"] - warnings.warn( - "No dataset specified. Defaulting to cnn_dailymail and nemotron-post-training-dataset-v2." - ) - # Adjust calib_size to match dataset length by extending or truncating as needed - args.calib_size = (args.calib_size + [args.calib_size[-1]] * len(args.dataset))[ - : len(args.dataset) - ] - - # Create calibration dataloader with max batch size - calib_dataloader = create_calibration_dataloader( - tokenizer=tokenizer, - dataset_names=args.dataset, - calib_sizes=args.calib_size, - batch_size=args.batch_size, - ) - - # Load and prepare model - model, model_type, original_architectures, calib_dataloader = load_and_prepare_model( - model_path=args.pyt_ckpt_path, - calib_dataloader=calib_dataloader, - accelerator=accelerator, - trust_remote_code=args.trust_remote_code, - ) - - quant_cfg = QUANT_CFG_CHOICES[args.qformat] - - quant_cfg = build_quant_cfg( - quant_cfg, - args.awq_block_size, - ) - - enable_quant_kv_cache = args.kv_cache_qformat != KV_CACHE_NONE - print(f"{'Enable' if enable_quant_kv_cache else 'Disable'} KV cache quantization") - - # Check if any bmm_quantizer is in the quant_cfg. If so, we need to enable the bmm_quantizer. - if enable_quant_kv_cache: - quant_cfg = mtq.update_quant_cfg_with_kv_cache_quant( - quant_cfg, - KV_QUANT_CFG_CHOICES[args.kv_cache_qformat]["quant_cfg"], - ) - - # Quantize the model - if accelerator.is_main_process: - print("Starting quantization...") - - start_time = time.time() - - if need_calibration(quant_cfg): - calibrate_fn = create_fsdp2_calibration_loop(model, calib_dataloader, accelerator) - else: - calibrate_fn = None - warnings.warn("Dynamic quantization. Calibration skipped.") - - with torch.no_grad(): - model = mtq.quantize(model, quant_cfg, forward_loop=calibrate_fn) - - elapsed = time.time() - start_time - - if accelerator.is_main_process: - print(f"Quantization completed in {elapsed:.2f}s") - mtq.print_quant_summary(model) - - start_time = time.time() - export_model(model, accelerator, args.export_path, original_architectures) - elapsed = time.time() - start_time - - if accelerator.is_main_process: - # Restore default padding and export the tokenizer as well. - if tokenizer is not None: - tokenizer.padding_side = default_padding_side - tokenizer.save_pretrained(args.export_path) - # Export the model - print(f"Export completed in {elapsed:.2f}s") - print(f"Model exported to {args.export_path}") - - print("Unpatching FSDP2 MP dtypes") - - -if __name__ == "__main__": - args = parse_args() - # This context manager can be removed once the update to FSDP2 function is reflected in torch - with patch_fsdp_mp_dtypes(): - main(args) diff --git a/examples/hf_ptq/slurm/multinode_fsdp2_ptq.slurm b/examples/hf_ptq/slurm/multinode_fsdp2_ptq.slurm new file mode 100644 index 00000000000..1baf3e2a79a --- /dev/null +++ b/examples/hf_ptq/slurm/multinode_fsdp2_ptq.slurm @@ -0,0 +1,91 @@ +#!/bin/bash + +# SPDX-FileCopyrightText: Copyright (c) 2025 NVIDIA CORPORATION & AFFILIATES. All rights reserved. +# SPDX-License-Identifier: Apache-2.0 +# +# Licensed under the Apache License, Version 2.0 (the "License"); +# you may not use this file except in compliance with the License. +# You may obtain a copy of the License at +# +# http://www.apache.org/licenses/LICENSE-2.0 +# +# Unless required by applicable law or agreed to in writing, software +# distributed under the License is distributed on an "AS IS" BASIS, +# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. +# See the License for the specific language governing permissions and +# limitations under the License. +# +# Multi-node post-training quantization with FSDP2, ready to run under Slurm. +# +# Slurm allocates the nodes and launches one task per node; each task starts a `torchrun` that +# spawns one process per GPU. `hf_ptq.py --use_fsdp2` then shards the model across all ranks +# (world_size = num_nodes * gpus_per_node) for distributed loading, calibration, and export. +# +# The defaults below quantize nvidia/NVIDIA-Nemotron-3-Super-120B-A12B-BF16 to NVFP4 with an FP8 +# KV cache. Edit the CONFIG block for your cluster/model, then submit with e.g.: +# +# sbatch --nodes=2 multinode_fsdp2_ptq.slurm +# +# For a single-node run, `--nodes=1` works unchanged. + +#SBATCH --job-name=fsdp2-ptq +#SBATCH --account={account} +#SBATCH --partition={partition} +#SBATCH --nodes=2 +#SBATCH --ntasks-per-node=1 # one torchrun launcher per node; it fans out to the GPUs +#SBATCH --gpus-per-node=8 +#SBATCH --exclusive +#SBATCH --time=04:00:00 +#SBATCH --output=%x_%j.log + +set -euo pipefail + +# --------------------------------------------------------------------------- +# CONFIG — edit these for your cluster and model. They are exported so `srun` +# propagates them into the container task below. +# --------------------------------------------------------------------------- +export CONTAINER_IMAGE={container_image} # e.g. nvcr.io#nvidia/pytorch:25.10-py3 +export MODELOPT_PATH={path_to_modelopt_repo} # host clone of TensorRT-Model-Optimizer +export MODEL_PATH=nvidia/NVIDIA-Nemotron-3-Super-120B-A12B-BF16 # HF repo id (auto-downloaded) or a local dir +export EXPORT_PATH={path_to_export_dir} # where the quantized checkpoint is written +export HF_HOME={path_to_hf_cache} # HF cache; persist so a repo id isn't re-downloaded each run +# export HF_TOKEN={hf_token} # required for gated repos (or `huggingface-cli login` on host) +export RECIPE=general/ptq/nvfp4_default-kv_fp8_cast # built-in recipe name or /path/to/recipe.yaml +export CALIB_SIZE=512 +export BATCH_SIZE=4 + +# Rendezvous: node 0 is the master; all ranks meet at MASTER_ADDR:MASTER_PORT. +export MASTER_ADDR=$(scontrol show hostnames "$SLURM_JOB_NODELIST" | head -1) +export MASTER_PORT=29531 + +# Mount the repo, export dir, and HF cache; add the model dir only when MODEL_PATH is a local path +# (a repo id is downloaded into HF_HOME instead). Run from the hf_ptq example dir. +CONTAINER_MOUNTS="${MODELOPT_PATH}:/modelopt,${EXPORT_PATH}:${EXPORT_PATH},${HF_HOME}:${HF_HOME}" +if [ -d "${MODEL_PATH}" ]; then + CONTAINER_MOUNTS="${CONTAINER_MOUNTS},${MODEL_PATH}:${MODEL_PATH}" +fi +srun --container-image="${CONTAINER_IMAGE}" \ + --container-mounts="${CONTAINER_MOUNTS}" \ + --container-workdir=/modelopt/examples/hf_ptq \ + bash -c ' + set -euo pipefail + export PYTHONPATH="/modelopt:${PYTHONPATH:-}" + export PYTHONUNBUFFERED=1 PYTORCH_CUDA_ALLOC_CONF=expandable_segments:True + # ModelOpt + its HF deps (transformers/accelerate/datasets/...) from the mounted source, then example extras. + pip install -q -e "/modelopt[hf]" --no-build-isolation + pip install -q -r requirements.txt + + torchrun \ + --nnodes="${SLURM_NNODES}" \ + --node_rank="${SLURM_NODEID}" \ + --nproc_per_node="$(nvidia-smi -L | wc -l)" \ + --rdzv_backend=c10d \ + --rdzv_endpoint="${MASTER_ADDR}:${MASTER_PORT}" \ + hf_ptq.py \ + --pyt_ckpt_path "${MODEL_PATH}" \ + --recipe "${RECIPE}" \ + --calib_size "${CALIB_SIZE}" \ + --batch_size "${BATCH_SIZE}" \ + --export_path "${EXPORT_PATH}" \ + --use_fsdp2 + ' diff --git a/modelopt/torch/export/hf_export_handlers.py b/modelopt/torch/export/hf_export_handlers.py index 800b51daca9..3449e62ea01 100644 --- a/modelopt/torch/export/hf_export_handlers.py +++ b/modelopt/torch/export/hf_export_handlers.py @@ -20,11 +20,12 @@ import torch.nn as nn -from modelopt.torch.quantization.utils import fsdp2_aware_weight_update +from modelopt.torch.quantization.utils import fsdp2_aware_weight_update, fsdp2_shard_local_pack +from modelopt.torch.utils.distributed import is_fsdp2_model from .layer_utils import get_expert_linear_names, is_quantlinear, set_expert_quantizer_amax from .model_config import QUANTIZATION_NONE -from .moe_utils import _export_fused_experts +from .moe_utils import _export_fused_experts, _export_fused_experts_keep_fused from .quant_utils import get_quantization_format from .registry import ExportContext, ExportModuleRegistry, PrepareMoEInputsRegistry @@ -36,6 +37,11 @@ def _has_fused_experts_quantizers(module: nn.Module) -> bool: return hasattr(module, f"{first_proj_attr}_weight_quantizers") +def _use_shard_local(model: nn.Module) -> bool: + """Whether to use shard-local packing (FSDP2 only).""" + return is_fsdp2_model(model) + + def _export_weight( module: nn.Module, ctx: ExportContext, @@ -128,23 +134,43 @@ def _export_moe_linear(name: str, module: nn.Module, ctx: ExportContext) -> None @ExportModuleRegistry.register(predicate=_has_fused_experts_quantizers) def _export_fused_experts_module(name: str, module: nn.Module, ctx: ExportContext) -> None: - """Split and quantize a fused-experts module with plural weight quantizers.""" - with fsdp2_aware_weight_update(ctx.model, module, reshard=False): - _export_fused_experts( - module, - ctx.dtype, - _moe_tied_cache=ctx.moe_tied_cache, - _tied_cache=ctx.tied_cache, - ) + """Split and quantize a fused-experts module with plural weight quantizers. + + Under FSDP2 the fused weight is ``Shard(0)`` on the expert dim, so the destructive per-rank split + (``_export_fused_experts``) would drop non-owner experts from the gather. Instead pack this rank's + experts IN PLACE and keep them fused (``_export_fused_experts_keep_fused`` inside + ``fsdp2_shard_local_pack``); the gate/up split is deferred to write time + (``_split_fused_experts_state_dict`` in the gather). Non-FSDP is unchanged. + """ + if _use_shard_local(ctx.model): + with fsdp2_shard_local_pack(ctx.model, module): + _export_fused_experts_keep_fused(module, ctx.dtype) + else: + with fsdp2_aware_weight_update(ctx.model, module, reshard=False): + _export_fused_experts( + module, + ctx.dtype, + _moe_tied_cache=ctx.moe_tied_cache, + _tied_cache=ctx.tied_cache, + ) @ExportModuleRegistry.register(predicate=is_quantlinear) def _export_quant_linear(name: str, module: nn.Module, ctx: ExportContext) -> None: - """Export a standard quantized linear layer.""" + """Export a standard quantized linear layer. + + ``fsdp2_shard_local_pack`` packs this rank's ``Shard(0)`` slice in place (no unshard) under FSDP2, + and is a no-op for non-FSDP models -- so the single-process path is unchanged. + """ if get_quantization_format(module) == QUANTIZATION_NONE: return + cm = ( + fsdp2_shard_local_pack(ctx.model, module) + if _use_shard_local(ctx.model) + else fsdp2_aware_weight_update(ctx.model, module, reshard=False) + ) try: - with fsdp2_aware_weight_update(ctx.model, module, reshard=False): + with cm: _export_weight(module, ctx) except AssertionError as e: raise AssertionError( diff --git a/modelopt/torch/export/moe_utils.py b/modelopt/torch/export/moe_utils.py index 787e173959e..a3ea8988236 100644 --- a/modelopt/torch/export/moe_utils.py +++ b/modelopt/torch/export/moe_utils.py @@ -294,6 +294,187 @@ def _export_fused_experts( _moe_tied_cache[_source_key] = module +def _export_fused_experts_keep_fused(module, dtype): + """Shard-local keep-fused NVFP4/FP8 pack for one rank's experts. + + Pack THIS rank's experts IN PLACE and keep the weight FUSED (``Shard(0)`` on E), + instead of splitting into per-expert child modules. + + Self-adapts to FSDP2 vs single-process via ``module._shard_local_start`` (set by + :func:`fsdp2_shard_local_pack`): + + * FSDP2 shard-local: the context has already ``to_local``'d ``gate_up_proj``/``down_proj`` to this + rank's ``[E/world, ...]`` block and recorded ``start`` (this rank's global expert offset), so the + per-expert weight quantizers are indexed ``[start + i]`` (the quantizer ModuleList is global + length-E and replicated -- verified in ``spike_moe_indexing.py``). + * single-process: ``start == 0`` and the tensors are the full ``[E, ...]`` block. + + Each local expert's fused first-projection ``[2I, H]`` is packed WHOLE by reusing the standard + ``_export_quantized_weight`` (so the per-tensor ``weight_scale_2`` is naturally shared by the gate + and up halves -- what vLLM's load-time fusion requires -- and per-block scales are row-local). The + packed per-expert results are stacked back into fused ``[local_n, ...]`` tensors + per-expert + scale buffers; the gate/up split is deferred to write time (``_split_fused_experts_state_dict``). + Non-destructive: the fused params survive, so :func:`fsdp2_shard_local_pack` re-registers them. + """ + from modelopt.torch.export.unified_export_hf import _export_quantized_weight + + first_proj_attr = getattr(module, "_first_proj_attr", "gate_up_proj") + start = getattr(module, "_shard_local_start", {}).get(first_proj_attr, 0) + first_proj_wq = getattr(module, f"{first_proj_attr}_weight_quantizers") + down_wq = module.down_proj_weight_quantizers + first_proj_iq = getattr(module, f"{first_proj_attr}_input_quantizer") + down_iq = module.down_proj_input_quantizer + + first_proj = getattr( + module, first_proj_attr + ).data # local [local_n, 2I, H] (or full [E, 2I, H]) + down = module.down_proj.data + local_n = first_proj.shape[0] + + def _pack_one(weight_2d, w_quant_src, i_quant): + wrapper = nn.Module() + wrapper.weight = nn.Parameter(weight_2d.contiguous(), requires_grad=False) + # deepcopy so packing does not mutate the shared calibrated quantizer state. + wrapper.weight_quantizer = copy.deepcopy(w_quant_src) + wrapper.input_quantizer = i_quant + wq = wrapper.weight_quantizer + if getattr(wq, "is_enabled", False) and ( + not hasattr(wq, "_amax") or wq._amax is None or torch.all(wq._amax == 0) + ): + # Uncalibrated expert (received no tokens): fall back to the weight's own amax. + wq.amax = weight_2d.abs().amax().to(torch.float32) + _export_quantized_weight(wrapper, dtype) + return ( + wrapper.weight.data, + getattr(wrapper, "weight_scale", None), + getattr(wrapper, "weight_scale_2", None), + getattr(wrapper, "input_scale", None), + ) + + fp_w, fp_s, fp_s2 = [], [], [] + dp_w, dp_s, dp_s2 = [], [], [] + fp_input_scale = dp_input_scale = None + for i in range(local_n): + g = start + i + w, s, s2, isc = _pack_one(first_proj[i], first_proj_wq[g], first_proj_iq) + fp_w.append(w) + fp_s.append(s) + fp_s2.append(s2) + fp_input_scale = isc if isc is not None else fp_input_scale + w, s, s2, isc = _pack_one(down[i], down_wq[g], down_iq) + dp_w.append(w) + dp_s.append(s) + dp_s2.append(s2) + dp_input_scale = isc if isc is not None else dp_input_scale + + def _register(attr, weights, scales, scale2s): + setattr(module, attr, nn.Parameter(torch.stack(weights), requires_grad=False)) + if scales[0] is not None: + module.register_buffer(f"{attr}_weight_scale", torch.stack(scales)) + if scale2s[0] is not None: + module.register_buffer( + f"{attr}_weight_scale_2", torch.stack([x.reshape(()) for x in scale2s]) + ) + + _register(first_proj_attr, fp_w, fp_s, fp_s2) + _register("down_proj", dp_w, dp_s, dp_s2) + if fp_input_scale is not None: + module.register_buffer(f"{first_proj_attr}_input_scale", fp_input_scale) + if dp_input_scale is not None: + module.register_buffer("down_proj_input_scale", dp_input_scale) + + +def _split_fused_experts_state_dict(state_dict, model): + """Split keep-fused expert tensors in a gathered ``state_dict`` into per-expert keys. + + Converts the ``_export_fused_experts_keep_fused`` output (fused ``{name}.gate_up_proj [E,2I,H/2]`` + + per-expert scales) into the same per-expert deployment keys ``_export_fused_experts`` emits + (``{name}.{e}.gate_proj.weight`` / ``up_proj`` / ``down_proj`` + scales). Byte-identical: packing + is row-independent, so slicing the packed fused tensor equals packing each half; gate/up share the + per-tensor ``weight_scale_2`` (packed together with one amax). Runs on the gathered (full) dict. + """ + from modelopt.torch.quantization.plugins.huggingface import _get_fused_expert_intermediate_dim + + for name, module in model.named_modules(): + first_proj_attr = getattr(module, "_first_proj_attr", "gate_up_proj") + if not hasattr(module, f"{first_proj_attr}_weight_quantizers"): + continue + prefix = f"{name}." if name else "" + fp_w = state_dict.pop(prefix + first_proj_attr, None) + if fp_w is None: + continue # not keep-fused here (e.g. already split) -> nothing to do + is_gated = getattr(module, "_is_gated", True) + n = module.num_experts + fp_s = state_dict.pop(prefix + f"{first_proj_attr}_weight_scale", None) + fp_s2 = state_dict.pop(prefix + f"{first_proj_attr}_weight_scale_2", None) + fp_in = state_dict.pop(prefix + f"{first_proj_attr}_input_scale", None) + dp_w = state_dict.pop(prefix + "down_proj", None) + dp_s = state_dict.pop(prefix + "down_proj_weight_scale", None) + dp_s2 = state_dict.pop(prefix + "down_proj_weight_scale_2", None) + dp_in = state_dict.pop(prefix + "down_proj_input_scale", None) + + # Drop leftover fused quantizer buffers (keep-fused does not delete them, unlike the + # per-expert path which calls _delete_fused_moe_source_attrs). + for k in [ + k for k in state_dict if k.startswith(prefix) and "_quantizer" in k[len(prefix) :] + ]: + state_dict.pop(k) + + edim = _get_fused_expert_intermediate_dim(module) if is_gated else None + + def _emit(e, proj, w, s, s2, insc): + # clone/contiguous so each per-expert/projection key is a DISTINCT tensor object. The + # shared scales (one input_scale across all experts; one weight_scale_2 across gate|up) + # would otherwise share a data_ptr and get collapsed by postprocess_state_dict's tied- + # weight dedup -- producing fewer keys than the per-expert path, which builds separate + # equal-valued objects. Cloning matches that format exactly (values are identical). + p = f"{prefix}{e}.{proj}." + state_dict[p + "weight"] = w.contiguous() + if s is not None: + state_dict[p + "weight_scale"] = s.contiguous() + if s2 is not None: + state_dict[p + "weight_scale_2"] = s2.clone() + if insc is not None: + state_dict[p + "input_scale"] = insc.clone() + + for e in range(n): + if is_gated: + _emit( + e, + "gate_proj", + fp_w[e, :edim], + fp_s[e, :edim] if fp_s is not None else None, + fp_s2[e] if fp_s2 is not None else None, + fp_in, + ) + _emit( + e, + "up_proj", + fp_w[e, edim:], + fp_s[e, edim:] if fp_s is not None else None, + fp_s2[e] if fp_s2 is not None else None, + fp_in, + ) + else: + _emit( + e, + "up_proj", + fp_w[e], + fp_s[e] if fp_s is not None else None, + fp_s2[e] if fp_s2 is not None else None, + fp_in, + ) + _emit( + e, + "down_proj", + dp_w[e], + dp_s[e] if dp_s is not None else None, + dp_s2[e] if dp_s2 is not None else None, + dp_in, + ) + return state_dict + + def save_expert_token_count_table(model: nn.Module, output_dir: str | Path | None = None): """Collect expert_token_count from all quantized MoE layers and save as an HTML table. diff --git a/modelopt/torch/export/registry.py b/modelopt/torch/export/registry.py index 8e2cda63df9..260cb32eea3 100644 --- a/modelopt/torch/export/registry.py +++ b/modelopt/torch/export/registry.py @@ -32,6 +32,8 @@ import torch import torch.nn as nn +from modelopt.torch.utils.distributed import is_fsdp2_model + __all__ = [ "ExportContext", "ExportHandler", @@ -54,8 +56,17 @@ class ExportContext: model: nn.Module dtype: torch.dtype is_modelopt_qlora: bool = False - tied_cache: dict[int, nn.Module] = field(default_factory=dict) - moe_tied_cache: dict[tuple[int, int], nn.Module] = field(default_factory=dict) + tied_cache: dict[int, nn.Module] | None = field(default_factory=dict) + moe_tied_cache: dict[tuple[int, int], nn.Module] | None = field(default_factory=dict) + + def __post_init__(self) -> None: + # FSDP2 may recycle data_ptr() values as modules are resharded, so pointer-keyed dedup can + # falsely alias distinct weights. Disable it for FSDP2; consequently, legitimately tied + # packed weights and scale buffers are not re-aliased and may be stored as duplicates. + # TODO: replace this with stable, name-based tied-group deduplication. + if is_fsdp2_model(self.model): + self.tied_cache = None + self.moe_tied_cache = None ExportHandler = Callable[[str, nn.Module, ExportContext], None] diff --git a/modelopt/torch/export/unified_export_hf.py b/modelopt/torch/export/unified_export_hf.py index cee64c22c05..f2ed7f9c317 100644 --- a/modelopt/torch/export/unified_export_hf.py +++ b/modelopt/torch/export/unified_export_hf.py @@ -15,6 +15,7 @@ """Code that export quantized Hugging Face models for deployment.""" +import contextlib import json import re import tempfile @@ -57,7 +58,9 @@ from modelopt.torch.quantization.nn import SequentialQuantizer, TensorQuantizer from modelopt.torch.quantization.qtensor import MXFP8QTensor, NVFP4QTensor from modelopt.torch.quantization.utils import fsdp2_aware_weight_update, quantizer_attr_names +from modelopt.torch.utils import distributed as _dist from modelopt.torch.utils.dataset_utils import _disable_use_cache +from modelopt.torch.utils.distributed import is_fsdp2_model try: from modelopt.torch.sparsity.attention_sparsity.conversion import export_sparse_attention_config @@ -840,7 +843,6 @@ def _export_transformers_checkpoint( Args: model: the full torch model to export. The actual quantized model may be a submodule. dtype: the weights data type to export the unquantized layers or the default model data type if None. - accelerator: the accelerator instance in case of distributed export setup. Returns: post_state_dict: Dict containing quantized weights @@ -854,8 +856,6 @@ def _export_transformers_checkpoint( f"({dtype}), which may lead to numerical errors." ) - accelerator = kwargs.get("accelerator") - # Handle input quantizers of experts that are not calibrated. Each MoE block is # dispatched by its experts container to the matching preparation handler. prepare_ctx = ExportContext(model=model, dtype=dtype, is_modelopt_qlora=is_modelopt_qlora) @@ -925,10 +925,29 @@ def _export_transformers_checkpoint( _reconstruct_fused_moe_linear(model) - if accelerator is not None: - # Gather state_dict from all ranks - quantized_state_dict = accelerator.get_state_dict(model) + # The FSDP2 parallel-write path shares this prep+pack, then streams the gather+write itself + # (per-owned-unit) instead of the whole-model gather below. See _parallel_write_hf_checkpoint. + if kwargs.get("_pack_only"): + # Only _parallel_write_hf_checkpoint passes this; it discards the state dict (writes per-unit + # itself), so return an empty dict to keep the tuple[dict, dict] contract for other callers. + return {}, quant_config + + if is_fsdp2_model(model): + # FSDP2 shard-local: each rank packed its own Shard(0) slice (keep-fused for experts, so the + # fused weight stays Shard(0)-on-E and gathers like any other weight). Gather every DTensor to + # full via full_tensor() (an all-gather -> full on every rank), then split the gathered + # keep-fused expert weights into the per-expert deployment keys. + from torch.distributed.tensor import DTensor + + from .moe_utils import _split_fused_experts_state_dict + + quantized_state_dict = { + k: (v.full_tensor().cpu() if isinstance(v, DTensor) else v.detach().cpu()) + for k, v in model.state_dict().items() + } + _split_fused_experts_state_dict(quantized_state_dict, model) else: + # Non-FSDP2: assumes a replicated model (rank 0 has the full state dict). quantized_state_dict = model.state_dict() # We define kv cache scale as amax / 448 for both FP8 and NVFP4 KV cache quantization. @@ -1382,6 +1401,250 @@ def export_speculative_decoding( exporter.export(export_dir, dtype) +# --------------------------------------------------------------------------- +# FSDP2 shard-local parallel checkpoint save (see docs/fsdp2_parallel_write.md). +# Each rank writes its owned decoder-layer units concurrently; host memory is +# bounded to ~one chunk of world units and the disk write parallelizes. +# --------------------------------------------------------------------------- + + +def _materialize_cpu(v: torch.Tensor) -> torch.Tensor: + """Convert a gathered state-dict value to a plain, contiguous CPU tensor. + + A gathered DTensor here is fully replicated (``full_tensor``/replicated buffer), so + ``to_local()`` yields the complete tensor. A ``QTensorWrapper`` (compressed weight) unwraps + to its packed data via ``.data``. + """ + from torch.distributed.tensor import DTensor + + if isinstance(v, DTensor): + v = v.to_local() + v = getattr(v, "data", v) # QTensorWrapper -> packed data; plain tensor -> itself + return v.detach().to("cpu").contiguous() + + +def _enumerate_export_units(model, id_to_name): + """Ordered export units, identical on every rank (this order drives ownership). + + Each decoder layer is one unit; a trailing "root-leaves" unit holds every module that owns + parameters directly and is not under a decoder-layer prefix (embed / lm_head / final norm). + Returns ``[(modules, is_root), ...]``. If decoder layers cannot be discovered, the whole model + becomes a single root unit (correct, but no parallelism). + """ + from modelopt.torch.quantization.utils.layerwise_calib import LayerActivationCollector + + decoder_layers = LayerActivationCollector.get_decoder_layers(model) + units: list[tuple[list[nn.Module], bool]] = [] + layer_prefixes: tuple[str, ...] = () + if decoder_layers is not None: + layer_prefixes = tuple(id_to_name[id(layer)] + "." for layer in decoder_layers) + units = [([layer], False) for layer in decoder_layers] + root_leaves = [ + m + for n, m in model.named_modules() + if next(m.parameters(recurse=False), None) is not None + and not (layer_prefixes and n.startswith(layer_prefixes)) + ] + units.append((root_leaves, True)) + return units + + +def _gather_unit(modules, id_to_name, is_owner): + """Gather one unit's params to the owner via ``full_tensor()`` (all-gather on every rank). + + Only the owner keeps the materialized CPU copy. Non-owners join the collective and drop the + result, so the gate is on the *keep*, never the collective (skipping it on non-owners deadlocks). + """ + from torch.distributed.tensor import DTensor + + unit_sd: dict[str, torch.Tensor] = {} + for m in modules: + base = id_to_name.get(id(m), "") + prefix = (base + ".") if base else "" + for name, v in m.state_dict().items(): + full = v.full_tensor() if isinstance(v, DTensor) else v # COLLECTIVE (all ranks) + if is_owner: + unit_sd[prefix + name] = _materialize_cpu(full) + return unit_sd if is_owner else None + + +def _write_owned_shards(owned, export_dir, rank_id, unit_idx, max_shard_size): + """Write one unit's tensors as safetensors shards via the HF splitter. + + Returns ``(weight_map {key: filename}, total_bytes)`` for the index merge. + Filenames embed rank + unit so ranks never collide; the HF splitter handles ``max_shard_size`` + parsing and never splits an individual tensor (a >cap tensor gets its own over-cap shard). + """ + from huggingface_hub import split_torch_state_dict_into_shards + + pattern = f"model-r{rank_id:02d}-u{unit_idx:04d}{{suffix}}.safetensors" + split = split_torch_state_dict_into_shards( + owned, filename_pattern=pattern, max_shard_size=max_shard_size + ) + for fname, keys in split.filename_to_tensors.items(): + save_file( + {k: owned[k] for k in keys}, str(Path(export_dir) / fname), metadata={"format": "pt"} + ) + return split.tensor_to_filename, split.metadata["total_size"] + + +def _write_unit(sd, model, export_dir, rank, idx, max_shard_size, kv_fmt, is_modelopt_qlora): + """Owner-side per-unit finalize, then write shards. + + Split keep-fused MoE, postprocess (per-key; includes data_ptr tied-weight dedup, matching the + standard path), revert conversion names, write. Returns the unit's ``(weight_map, total_bytes)``. + """ + from .moe_utils import _split_fused_experts_state_dict + + _split_fused_experts_state_dict(sd, model) # only this unit's keep-fused keys are present + sd = postprocess_state_dict(sd, 448, kv_fmt, is_modelopt_qlora) + with contextlib.suppress(Exception): + sd = revert_weight_conversion_quant_aware(model, sd) + return _write_owned_shards(sd, export_dir, rank, idx, max_shard_size) + + +def _finalize_index(local_maps, export_dir, rank, world): + """gather_object the per-rank HF weight_maps to rank 0; write the standard index. No sidecars.""" + from huggingface_hub.constants import SAFETENSORS_INDEX_FILE + + gathered: list[Any] | None + if world > 1: + gathered = [None] * world if rank == 0 else None + torch.distributed.gather_object(local_maps, gathered, dst=0) + else: + gathered = [local_maps] + if rank != 0 or gathered is None: + return + weight_map: dict[str, str] = {} + total = 0 + for rank_maps in gathered: + for t2f, nbytes in rank_maps: + weight_map.update(t2f) + total += nbytes + index = {"metadata": {"total_size": total}, "weight_map": weight_map} + (Path(export_dir) / SAFETENSORS_INDEX_FILE).write_text(json.dumps(index, indent=2)) + + +def _write_hf_metadata( + model, export_dir, quant_config, save_modelopt_state, architectures_override +): + """Rank-0: write the config / quant / generation metadata ``save_pretrained`` would emit. + + Writes config.json (+ quantization_config / sparse_attention), generation_config.json, and + hf_quant_config.json. + """ + export_dir = Path(export_dir) + quantization_details = (quant_config or {}).get("quantization", {}) + is_quantized_export = ( + quantization_details.get("quant_algo") is not None + or quantization_details.get("kv_cache_quant_algo") is not None + ) + hf_quant_config = None + if is_quantized_export: + with contextlib.suppress(Exception): + name_mapper = build_reverse_name_mapper(model) + if name_mapper is not None: + revert_quant_config_names(quant_config.get("quantization", {}), name_mapper) + with open(export_dir / "hf_quant_config.json", "w") as f: + json.dump(quant_config, f, indent=4) + hf_quant_config = convert_hf_quant_config_format(quant_config) + + _sanitize_generation_config_for_save(model) + model.config.save_pretrained(str(export_dir)) + if getattr(model, "generation_config", None) is not None: + model.generation_config.save_pretrained(str(export_dir)) + + original_config = export_dir / "config.json" + with open(original_config) as f: + config_data = json.load(f) + sanitize_hf_config_for_deployment(config_data, model) + if architectures_override: + config_data["architectures"] = architectures_override + if hf_quant_config is not None: + config_data["quantization_config"] = hf_quant_config + if export_sparse_attention_config is not None: + sparse_attn_config = export_sparse_attention_config(model) + if sparse_attn_config is not None: + config_data["sparse_attention_config"] = sparse_attn_config + with open(original_config, "w") as f: + json.dump(config_data, f, indent=4) + + if save_modelopt_state: + warnings.warn( + "save_modelopt_state is not yet supported on the FSDP2 parallel-write export path; " + "modelopt_state was not written. The deployment checkpoint (weights + config) is complete." + ) + + +def _parallel_write_hf_checkpoint( + model, + dtype, + export_dir, + *, + save_modelopt_state=False, + extra_state_dict=None, + max_shard_size="10GB", + architectures_override=None, + is_modelopt_qlora=False, + **kwargs, +): + """FSDP2 shard-local parallel checkpoint save (see docs/fsdp2_parallel_write.md). + + Reuses the standard prep+pack (``_pack_only``) so every rank packs its own ``Shard(0)`` slice, + then streams the write in chunks of ``world`` decoder-layer units: each rank gathers every unit + (collective) but materializes + writes only the one it owns (``i % world``). Host peak is bounded + to ~one chunk; writes parallelize across ranks. Rank 0 merges the index + writes metadata. + """ + export_dir = Path(export_dir) + rank, world = _dist.rank(), _dist.size() + + # Phase A — prep + shard-local pack, shared with the standard path (packs in place, all ranks). + _, quant_config = _export_transformers_checkpoint( + model, dtype, is_modelopt_qlora=is_modelopt_qlora, _pack_only=True, **kwargs + ) + if getattr(model, "hf_quantizer", None) is not None: + model.hf_quantizer = None + + kv_fmt = quant_config["quantization"]["kv_cache_quant_algo"] + id_to_name = {id(m): n for n, m in model.named_modules()} + units = _enumerate_export_units(model, id_to_name) + n = len(units) + + # Phase B — chunk loop: all ranks gather each unit; owner (i % world) writes it. + local_maps = [] + for chunk in range(0, n, world): + mine = None + for j in range(world): + i = chunk + j + if i >= n: + break + modules, _is_root = units[i] + sd = _gather_unit(modules, id_to_name, is_owner=(j == rank)) # COLLECTIVE all ranks + if j == rank: + mine = (sd, i) + if mine is not None and mine[0] is not None: + sd, idx = mine + local_maps.append( + _write_unit( + sd, model, export_dir, rank, idx, max_shard_size, kv_fmt, is_modelopt_qlora + ) + ) + + # extra_state_dict (e.g. MTP) -> final rank-0 shard. + if extra_state_dict and rank == 0: + esd = {k: v.detach().to("cpu") for k, v in extra_state_dict.items()} + local_maps.append(_write_owned_shards(esd, export_dir, rank, n, max_shard_size)) + + # Phase C — index + metadata. + _dist.barrier() + _finalize_index(local_maps, export_dir, rank, world) + if rank == 0: + _write_hf_metadata( + model, export_dir, quant_config, save_modelopt_state, architectures_override + ) + _dist.barrier() + + def export_hf_checkpoint( model: Any, dtype: torch.dtype | None = None, @@ -1397,6 +1660,10 @@ def export_hf_checkpoint( This function automatically detects whether the model is from transformers or diffusers and applies the appropriate export logic. + Under ``torch.distributed`` (e.g. FSDP2), all ranks participate in the + collective state-dict gather inside ``_export_transformers_checkpoint``; + only rank 0 writes files. A final barrier syncs the other ranks. + Args: model: The full torch model to export. The actual quantized model may be a submodule. Supports both transformers models (e.g., LlamaForCausalLM) and diffusers @@ -1430,6 +1697,26 @@ def export_hf_checkpoint( ) return + is_distributed = ( + torch.distributed.is_available() + and torch.distributed.is_initialized() + and is_fsdp2_model(model) + ) + + # FSDP2 shard-local parallel write: each rank writes its owned decoder-layer units concurrently + # (host memory bounded to ~a chunk, no rank-0 full gather). + if is_distributed: + _parallel_write_hf_checkpoint( + model, + dtype, + export_dir, + save_modelopt_state=save_modelopt_state, + extra_state_dict=extra_state_dict, + max_shard_size=max_shard_size, + **kwargs, + ) + return + try: post_state_dict, hf_quant_config = _export_transformers_checkpoint(model, dtype, **kwargs) @@ -1461,6 +1748,10 @@ def export_hf_checkpoint( "names may not match the original HF hub checkpoint." ) + # Under torch.distributed only rank 0 writes; others sync at the finally barrier. + if is_distributed and torch.distributed.get_rank() != 0: + return + # Only treat the export as quantized when at least one quant_algo field is set. # get_quant_config always returns a dict (even for sparsity-only or unmodified models), # so emitting hf_quant_config.json unconditionally produces a file with @@ -1489,6 +1780,7 @@ def export_hf_checkpoint( _sanitize_generation_config_for_save(model) + # TODO: parallelize the disk write across ranks (avoid single-process speed + rank-0 OOM). try: model.save_pretrained( export_dir, @@ -1525,3 +1817,6 @@ def export_hf_checkpoint( " can be saved with torch.save for further inspection." ) raise e + finally: + if is_distributed: + torch.distributed.barrier() diff --git a/modelopt/torch/opt/utils.py b/modelopt/torch/opt/utils.py index 96c0dd555f3..43fc9e6b7a4 100644 --- a/modelopt/torch/opt/utils.py +++ b/modelopt/torch/opt/utils.py @@ -113,6 +113,14 @@ def _lazy_init_retain_mesh_info(self): fsdp_state._lazy_init = types.MethodType(_lazy_init_retain_mesh_info, fsdp_state) fsdp_states.append(fsdp_state) + + # The calibration loop calls model.forward directly, bypassing model.__call__ and the root's + # FSDP2 pre-forward hook, so the root's _lazy_init never runs and decoder layers (called via + # __call__) each self-mark as FSDP root — which later breaks get_model_state_dict at export. + # Run the root's _lazy_init here so it initializes root-first and marks children non-root. + if isinstance(model, FSDPModule): + _get_module_state(model)._lazy_init() + yield for fsdp_state in fsdp_states: if fsdp_state._fsdp_param_group and hasattr(fsdp_state, "_post_forward_mesh_info_after"): diff --git a/modelopt/torch/quantization/model_quant.py b/modelopt/torch/quantization/model_quant.py index 1adda0511a6..843a7024df7 100644 --- a/modelopt/torch/quantization/model_quant.py +++ b/modelopt/torch/quantization/model_quant.py @@ -42,6 +42,7 @@ from .mode import QuantizeModeRegistry, get_modelike_from_algo_cfg from .nn import QuantModule, TensorQuantizer from .utils import is_quantized +from .utils.core_utils import materialize_fsdp2_root __all__ = [ "auto_quantize", @@ -106,7 +107,9 @@ def forward_loop(model): is_training = model.training model.eval() - with forward_with_reshard(model): + # materialize_fsdp2_root: calibration calls model.forward directly (bypasses __call__ and the + # root's FSDP2 pre-forward hook), so unshard the root's embed/lm_head/norm for the forward. + with materialize_fsdp2_root(model), forward_with_reshard(model): apply_mode( model, mode=get_modelike_from_algo_cfg(algorithm), diff --git a/modelopt/torch/quantization/utils/core_utils.py b/modelopt/torch/quantization/utils/core_utils.py index 80352ced9eb..c77aafd96f5 100644 --- a/modelopt/torch/quantization/utils/core_utils.py +++ b/modelopt/torch/quantization/utils/core_utils.py @@ -25,7 +25,7 @@ import torch.nn.functional as F from torch.distributed.fsdp import FSDPModule, MixedPrecisionPolicy, fully_shard from torch.distributed.fsdp._fully_shard._fsdp_param import FSDPParam -from torch.distributed.tensor import Replicate +from torch.distributed.tensor import DTensor, Replicate from modelopt.torch.quantization.config import QuantizerCfgEntry from modelopt.torch.utils import get_unwrapped_name, print_rank_0 @@ -135,7 +135,7 @@ def convert_quantization_axis_to_reduce_axis(input, axis): """ if axis is None: return None - axis = axis if isinstance(axis, (list, tuple)) else [axis] + axis = axis if isinstance(axis, list | tuple) else [axis] # Handle positive and negative axis. reduce_axis = [i for i in range(input.dim()) if i not in axis and (i - input.dim()) not in axis] return reduce_axis @@ -226,13 +226,13 @@ def representative_weight_quantizer(module: nn.Module, weight_name: str = "weigh singular = quantizer_attr_names(weight_name).weight_quantizer q = getattr(module, singular, None) - if isinstance(q, (TensorQuantizer, SequentialQuantizer)): + if isinstance(q, TensorQuantizer | SequentialQuantizer): return q plural = getattr(module, singular + "s", None) if isinstance(plural, nn.ModuleList) and len(plural) > 0: first = plural[0] - if isinstance(first, (TensorQuantizer, SequentialQuantizer)): + if isinstance(first, TensorQuantizer | SequentialQuantizer): return first return None @@ -426,11 +426,14 @@ def _get_fsdp2_mesh(module: nn.Module): return None fsdp_state = _get_module_state(module) - if ( - fsdp_state._fsdp_param_group - and fsdp_state._fsdp_param_group.post_forward_mesh_info is not None - ): - return fsdp_state._fsdp_param_group.post_forward_mesh_info.mesh + pg = fsdp_state._fsdp_param_group + if pg is None: + return None + # A root FSDP module has reshard_after_forward=False by default, so its + # post_forward_mesh_info is None; fall back to the sharding mesh (mesh_info), + # which is the same FSDP shard mesh (post_forward_mesh_info is only the reshard target). + mesh_info = pg.post_forward_mesh_info or pg.mesh_info + return mesh_info.mesh if mesh_info is not None else None def _get_module_name(module: nn.Module, root_model: nn.Module, name_to_module: dict | None = None): @@ -496,8 +499,6 @@ def fsdp2_weight_access_and_writeback_context(module: nn.Module, root_model: nn. If TP is implemented with DTensor, the weight will be a local tensor of the TP DTensor under this context. """ - assert isinstance(root_model, torch.distributed.fsdp.FSDPModule), "We only support FSDP2" - assert not hasattr(module, "_hf_hook"), "We dont support FSDP2 with HF accelerate hooks" fsdp_module = _get_enclosing_fsdp_module(module, root_model) assert fsdp_module is not None, "Module is not wrapped by FSDP" @@ -519,24 +520,42 @@ def fsdp2_weight_access_and_writeback_context(module: nn.Module, root_model: nn. placements=[Replicate()] * fsdp_dim + list(original_placements[fsdp_dim:]), device_mesh=original_device_mesh, ) - originals[name] = (param, collected, original_placements, original_device_mesh) - _set_parameter(module, name, nn.Parameter(collected.to_local())) - - yield - - # Write back and restore original DTensor parameters. - for name, ( - original_param, - collected, - original_placements, - original_device_mesh, - ) in originals.items(): - original_param.to_local().data.copy_( - collected.redistribute( - placements=original_placements, device_mesh=original_device_mesh - ).to_local() + local_replicated = collected.to_local() + # cpu_offload: gathered shard is on CPU; mirror to GPU for forward. + on_cpu = local_replicated.device.type == "cpu" and torch.cuda.is_available() + working = local_replicated.to(torch.cuda.current_device()) if on_cpu else local_replicated + cpu_mirror = local_replicated if on_cpu else None + originals[name] = ( + param, + collected, + original_placements, + original_device_mesh, + cpu_mirror, + working, ) - _set_parameter(module, name, original_param) + _set_parameter(module, name, nn.Parameter(working)) + + try: + yield + finally: + # Write back and restore original DTensor parameters. Runs on both success + # and exception so the module never lingers with the temporary local params. + for name, ( + original_param, + collected, + original_placements, + original_device_mesh, + cpu_local, + gpu_working, + ) in originals.items(): + if cpu_local is not None: + cpu_local.data.copy_(gpu_working.data.to(cpu_local.device)) + original_param.to_local().data.copy_( + collected.redistribute( + placements=original_placements, device_mesh=original_device_mesh + ).to_local() + ) + _set_parameter(module, name, original_param) @contextmanager @@ -937,6 +956,203 @@ def fsdp2_aware_weight_update(root_model, modules_to_update, reshard=True): root_module.reshard() +_ShardInfo = namedtuple("_ShardInfo", ["name", "old", "mesh", "placements"]) + + +def _shard_start(param): + """Global offset along shard dim 0 for a ``Shard(0)`` DTensor on a 1-D FSDP mesh. + + Returns 0 for non-DTensor params. Uses torch's even-split convention: the first ``dim0 % world`` + ranks get an extra row, so rank ``r``'s offset is ``r * (dim0 // world) + min(r, dim0 % world)``. + Correct for both even and uneven sharding. + """ + if not isinstance(param, DTensor): + return 0 + dim0 = param.shape[0] # global size + world = param.device_mesh.size() + r = param.device_mesh.get_local_rank() + base, rem = divmod(dim0, world) + return r * base + min(r, rem) + + +def _rebuild_fsdp_param_from_shard(old_fp, packed_local): + """Re-register a PACKED weight as an FSDPParam from full shape + this rank's packed shard. + + Does so WITHOUT materializing the full weight. + + The obvious route -- ``DTensor.from_local(shard)`` then ``FSDPParam(dtensor, ...)`` -- crashes: + ``_init_sharded_param`` takes a DTensor down its tensor-parallel branch + (``DeviceMesh._concatenate([dp_mesh, tp_mesh])``) which has no ``tp_mesh`` on a plain 1-D FSDP + mesh. Instead build the FSDPParam from a ``meta``-device full-shape tensor (a plain tensor -> + non-DTensor path; ``meta`` is explicitly allowed by ``_init_sharded_param``), which computes every + size/stride/spec field with no real full data and no crash, then overwrite only the two + data-holding fields (``_sharded_param_data``, ``sharded_param``) with this rank's packed shard. + Only dim-0 (the ``Shard(0)`` axis) differs between shard and full; inner dims are identical. + Validated by ``spike_inplace_pack.py`` (world=2, NVFP4). + """ + from modelopt.torch.quantization.qtensor.base_qtensor import QFSDPParam, QTensorWrapper + + full_shape = (old_fp._orig_size[0], *packed_local.shape[1:]) + is_qtw = isinstance(packed_local, QTensorWrapper) + param_class = QFSDPParam if is_qtw else FSDPParam + mp = MixedPrecisionPolicy( + param_dtype=packed_local.dtype, + reduce_dtype=None, + output_dtype=None, + cast_forward_inputs=False, + ) + + # (a) meta full-shape skeleton -> computes sharded_size / padded_sharded_param_size / specs + skeleton = torch.empty(full_shape, dtype=packed_local.dtype, device="meta") + if is_qtw: + skeleton = QTensorWrapper(skeleton, metadata={**packed_local.metadata, "shape": full_shape}) + new_fp = param_class( + nn.Parameter(skeleton, requires_grad=False), + old_fp._module_info, + old_fp.mesh_info, + old_fp.post_forward_mesh_info, + old_fp.device, + None, + mp, + None, + ) + if param_class is FSDPParam: + new_fp.init_dtype_attrs(mp) + + # (b) overwrite the two data fields with the real packed shard (pad the 0th shard as FSDP does) + local = (packed_local.data if is_qtw else packed_local).contiguous().to(old_fp.device) + padded = local.new_zeros(new_fp.padded_sharded_param_size) + padded.narrow(0, 0, new_fp.sharded_size[0]).copy_(local) + new_fp._sharded_param_data = padded.view(-1) + new_fp.sharded_param = nn.Parameter( + new_fp.to_sharded_dtensor(padded.narrow(0, 0, new_fp.sharded_size[0])), requires_grad=False + ) + new_fp._setattr_on_modules(new_fp.sharded_param) + return new_fp + + +def _rewrap_scale_buffers_shard0(module, captured, local_dim0): + """Re-wrap packed per-shard scale buffers as ``Shard(0)`` DTensors. + + A later ``full_tensor()`` gather then reconstructs the full scale. + + Scale buffers are ordinary registered buffers (not FSDP params), so they can use + ``DTensor.from_local`` directly (the FSDPParam-constructor crash does not apply). A buffer is + treated as sharded iff its dim-0 matches the local shard row/expert count (``local_dim0``): + ``weight_scale [rows/world, ...]`` and per-expert ``weight_scale_2 [E/world]`` -> ``Shard(0)``; + per-tensor scalars (``weight_scale_2`` for a plain linear, ``input_scale``) stay replicated. + Validated by ``spike_scale_gather.py`` (world=2, NVFP4). + """ + info = next(iter(captured.values())) + mesh, placements = info.mesh, info.placements + # The scale's dim-0 matches the weight's global dim-0 (per-row / per-expert). Pass the true global + # shape+stride to from_local: without it, from_local infers global = local * world (even-shard + # assumption), so under UNEVEN sharding ranks disagree on the global shape and full_tensor() + # deadlocks. sharded_param DTensors (weights) already carry the right shape via their spec. + global_dim0 = info.old._orig_size[0] + for bname, buf in list(module._buffers.items()): + if buf is None or isinstance(buf, DTensor) or buf.dim() == 0: + continue + if buf.shape[0] == local_dim0: + local = buf.contiguous() + global_shape = torch.Size((global_dim0, *local.shape[1:])) + global_stride = torch.empty(global_shape, device="meta").stride() + module._buffers[bname] = DTensor.from_local( + local, mesh, placements, shape=global_shape, stride=global_stride + ) + + +@contextmanager +def fsdp2_shard_local_pack(root_model, module): + """Pack a module's ``Shard(0)`` weights on the LOCAL shard, in place, keeping them sharded. + + Works for plain quant-linears (``Shard(0)`` on ``out``) and keep-fused experts (``Shard(0)`` on + ``E``). No unshard, no reshard, no collective -> leaves a valid packed FSDP module. No-ops for + non-FSDP models (so the same handler code covers the single-process/standard path). + + Enter: ``to_local`` each weight param (recording its old FSDPParam + mesh/placements and this + rank's global shard offset in ``module._shard_local_start``); the wrapped handler packs the plain + local block in place. Exit: re-register each packed weight via :func:`_rebuild_fsdp_param_from_shard` + and re-wrap the packed scale buffers via :func:`_rewrap_scale_buffers_shard0`; no reshard. + """ + if not isinstance(root_model, FSDPModule): + yield + return + + root_module = _get_enclosing_fsdp_module(module, root_model) + group = fully_shard.state(root_module)._fsdp_param_group + mapping = create_fsdp_param_mapping(group.fsdp_params, root_model) + + captured = {} + module._shard_local_start = {} + for pname, param in list(module.named_parameters(recurse=False)): + name = f"{_get_module_name(module, root_model)}.{pname}" + if name not in mapping: + continue + # A non-DTensor param is replicated / not row-sharded (e.g. shard_root=False root params): + # shard-local packing does not apply -> leave it for the handler to pack in place (it stays + # full and passes through the gather unchanged). Only Shard(0) DTensors get the shard-local path. + if not isinstance(param, DTensor): + continue + # Fail fast + clear on uneven sharding rather than deadlocking the gather later. Symmetric + # (global shape + world are identical on every rank) so this raise can't itself hang. + world = param.device_mesh.size() + if param.shape[0] % world != 0: + raise NotImplementedError( + f"fsdp2_shard_local_pack does not support uneven sharding: '{name}' has dim0=" + f"{param.shape[0]}, not divisible by world size {world}. Use a world size that " + f"divides every sharded dim-0." + ) + captured[pname] = _ShardInfo(name, mapping[name], param.device_mesh, param.placements) + module._shard_local_start[pname] = _shard_start(param) + module._parameters[pname] = nn.Parameter(param.to_local(), requires_grad=False) + try: + yield + finally: + local_dim0 = None + with no_requires_grad(), enable_fake_quant(module): + for pname, info in captured.items(): + packed_local = getattr(module, pname) + if local_dim0 is None: + local_dim0 = packed_local.shape[0] + mapping[info.name] = _rebuild_fsdp_param_from_shard(info.old, packed_local) + info.old._post_load_hook_handle.remove() + if captured: + _rewrap_scale_buffers_shard0(module, captured, local_dim0) + group.fsdp_params = list(mapping.values()) + module.__dict__.pop("_shard_local_start", None) + + +@contextmanager +def materialize_fsdp2_root(model: nn.Module): + """Unshard a sharded FSDP2 root's own params (embed/lm_head/norm) for a calibration forward. + + The calibration loop calls ``model.forward(**batch)`` directly (``dataset_utils._forward_loop``) + rather than ``model(**batch)``, so ``nn.Module.__call__`` is bypassed and the root's FSDP2 + forward pre-hook never fires. Its own params (embed/lm_head/norm) stay sharded DTensors and the + forward hits ``aten.embedding: mixed Tensor and DTensor``. Decoder layers are unaffected: they are + called as ``layer(...)`` inside ``forward``, so their pre-hooks fire and they unshard normally. + + Unshard the root up front so its params are full tensors, then reshard on exit. The bypass also + skips the root's post-forward hook, so nothing reshards it mid-calibration and a single unshard + holds across all batches. Cheap: only the root's own param group is gathered, not the decoder + layers. No-op for non-FSDP2 or already-replicated roots. + """ + root_sharded = False + if isinstance(model, FSDPModule): + pg = fully_shard.state(model)._fsdp_param_group + root_sharded = pg is not None and pg.is_sharded + if root_sharded: + with enable_fake_quant(model): + model.unshard() + try: + yield + finally: + if root_sharded: + with enable_fake_quant(model): + model.reshard() + + def update_quant_cfg_with_kv_cache_quant( quant_cfg: dict[str, Any], kv_cache_quant_cfg: list[QuantizerCfgEntry] ) -> dict[str, Any]: diff --git a/modelopt/torch/utils/distributed.py b/modelopt/torch/utils/distributed.py index 7922b688052..12287865b7f 100644 --- a/modelopt/torch/utils/distributed.py +++ b/modelopt/torch/utils/distributed.py @@ -27,6 +27,7 @@ import torch import torch.distributed +from torch.distributed.fsdp import CPUOffloadPolicy, FSDPModule, fully_shard from torch.distributed.tensor import DTensor __all__ = [ @@ -34,7 +35,9 @@ "ParallelState", "backend", "barrier", + "fsdp2_wrap", "is_available", + "is_fsdp2_model", "is_initialized", "is_master", "rank", @@ -216,6 +219,82 @@ def cleanup(): torch.distributed.destroy_process_group() +def is_fsdp2_model(model) -> bool: + """Return True if any submodule of ``model`` has been wrapped with FSDP2 ``fully_shard``.""" + return any(isinstance(m, FSDPModule) for m in model.modules()) + + +def fsdp2_wrap(model, shard_root=True, mp_policy=None, cpu_offload: bool = False): + """Auto-detect a HF causal-LM's decoder layers and FSDP2 ``fully_shard`` each one. + + By default (``shard_root=True``) the root module is wrapped too, so embed/lm_head/norm are + sharded instead of replicated per rank; pass ``shard_root=False`` to leave the root replicated + (only decoder layers sharded). Returns the detected decoder layers so callers can reuse the + detection result. + """ + # Lazy import: layerwise_calib imports this module at top level (circular). + from modelopt.torch.quantization.utils.layerwise_calib import LayerActivationCollector + + decoder_layers = LayerActivationCollector.get_decoder_layers(model) + if decoder_layers is None: + raise RuntimeError( + "Could not auto-detect decoder layers; FSDP2 wrap requires a standard HF causal-LM layout." + ) + + fsdp_kwargs: dict[str, Any] = {"reshard_after_forward": True} + if mp_policy is not None: + fsdp_kwargs["mp_policy"] = mp_policy + if cpu_offload: + fsdp_kwargs["offload_policy"] = CPUOffloadPolicy() + + # Snapshot/restore config.architectures: some HF builders mutate it during fully_shard. + config = getattr(model, "config", None) + architectures = list(getattr(config, "architectures", []) or []) + for layer in decoder_layers: + fully_shard(layer, **fsdp_kwargs) + if shard_root: + fully_shard(model, **fsdp_kwargs) + if config is not None and architectures: + config.architectures = architectures + + return decoder_layers + + +def broadcast_state_dict( + state_dict_or_none: dict | None, + src: int, + device: torch.device, + pg=None, +) -> dict: + """Broadcast a dict of CPU tensors from rank ``src`` to all ranks. + + Two phases: (1) broadcast metadata (key list + shape/dtype) via + ``broadcast_object_list``, (2) broadcast each tensor via ``dist.broadcast``. + Source rank passes the populated dict; non-source ranks pass ``None``. + Returns a dict of tensors on ``device`` on every rank. + """ + is_src = torch.distributed.get_rank() == src + meta: list[Any] = ( + [{name: (tuple(t.shape), t.dtype) for name, t in state_dict_or_none.items()}] + if is_src and state_dict_or_none is not None + else [None] + ) + torch.distributed.broadcast_object_list(meta, src=src, group=pg) + meta_dict = meta[0] + assert meta_dict is not None, f"src rank {src} passed no state dict to broadcast" + + src_state_dict = state_dict_or_none or {} + out: dict[str, torch.Tensor] = {} + for name, (shape, dtype) in meta_dict.items(): + if is_src: + t = src_state_dict[name].to(device) + else: + t = torch.empty(shape, dtype=dtype, device=device) + torch.distributed.broadcast(t, src=src, group=pg) + out[name] = t + return out + + class DistributedProcessGroup: """A convenient wrapper around torch.distributed.ProcessGroup objects.""" diff --git a/modelopt/torch/utils/plugins/model_load_utils.py b/modelopt/torch/utils/plugins/model_load_utils.py new file mode 100644 index 00000000000..bab24034b46 --- /dev/null +++ b/modelopt/torch/utils/plugins/model_load_utils.py @@ -0,0 +1,428 @@ +# SPDX-FileCopyrightText: Copyright (c) 2024 NVIDIA CORPORATION & AFFILIATES. All rights reserved. +# SPDX-License-Identifier: Apache-2.0 +# +# Licensed under the Apache License, Version 2.0 (the "License"); +# you may not use this file except in compliance with the License. +# You may obtain a copy of the License at +# +# http://www.apache.org/licenses/LICENSE-2.0 +# +# Unless required by applicable law or agreed to in writing, software +# distributed under the License is distributed on an "AS IS" BASIS, +# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. +# See the License for the specific language governing permissions and +# limitations under the License. + +"""HuggingFace-coupled FSDP2 model loading helpers.""" + +import json +import logging +import os +import re +from collections.abc import Callable +from itertools import chain +from typing import Any + +import torch +import torch.nn as nn +from accelerate import init_empty_weights +from huggingface_hub import snapshot_download +from safetensors import safe_open +from torch.distributed.checkpoint.state_dict import StateDictOptions, set_model_state_dict +from torch.distributed.tensor import DTensor +from transformers import AutoConfig, AutoModelForCausalLM + +try: + from transformers.conversion_mapping import get_model_conversion_mapping +except ImportError: # transformers<5 has no fused-MoE weight-conversion engine + get_model_conversion_mapping = None + +from modelopt.torch.utils.distributed import ( + barrier, + broadcast_state_dict, + fsdp2_wrap, + is_initialized, +) + +logger = logging.getLogger(__name__) + + +def read_safetensors_subset( + ckpt_path: str, + weight_map: dict, + select: Callable[[str], bool], +) -> dict: + """Read tensors whose name satisfies ``select`` from safetensors files. + + Groups param names by file to avoid re-opening. Returns CPU tensors. + Uses ``safe_open`` so only the requested tensors' bytes are read. + + ``get_tensor`` returns a zero-copy view into the mmap'd file; the bytes are + not actually read from disk until first touched. We ``clone()`` here to force + the read eagerly, while this function runs (each rank reading its own layers + in parallel). Without it the read is deferred to the later per-source + broadcast (``.to(device)``), which is serialized across ranks and silently + destroys the read parallelism this loader exists to provide. + """ + by_file: dict[str, list[str]] = {} + for name, file in weight_map.items(): + if select(name): + by_file.setdefault(file, []).append(name) + + state: dict[str, torch.Tensor] = {} + for file, names in by_file.items(): + with safe_open(os.path.join(ckpt_path, file), framework="pt", device="cpu") as f: + for name in names: + state[name] = f.get_tensor(name).clone() + return state + + +def weight_map_for(ckpt_path: str) -> dict[str, str]: + """Return the ``param_name → safetensors_file`` map for a local checkpoint directory.""" + index_path = os.path.join(ckpt_path, "model.safetensors.index.json") + single_file = os.path.join(ckpt_path, "model.safetensors") + if os.path.exists(index_path): + with open(index_path) as f: + return json.load(f)["weight_map"] + if os.path.exists(single_file): + with safe_open(single_file, framework="pt", device="cpu") as f: + return dict.fromkeys(f.keys(), "model.safetensors") + raise RuntimeError( + f"No safetensors checkpoint at {ckpt_path} " + "(expected model.safetensors or model.safetensors.index.json)." + ) + + +def _resolve_checkpoint_dir(ckpt_path: str, rank: int) -> str: + """Local dir for ``ckpt_path``; resolves an HF Hub ID (rank 0 downloads, others wait).""" + if os.path.isdir(ckpt_path): + return ckpt_path + if rank == 0: + snapshot_download(ckpt_path) + if is_initialized(): + barrier() + return snapshot_download(ckpt_path) + + +def _materialize_meta_model(model: nn.Module, device: torch.device) -> None: + """Replace meta params/buffers with empty real ones on ``device``; move real buffers there. + + Goes through ``model._apply`` so FSDP2's override refreshes its internal + ``_sharded_param_data`` pointers via ``reset_sharded_param``. + """ + model._apply(lambda t: torch.empty_like(t, device=device) if t.is_meta else t.to(device)) + + +def _promote_non_dtensor_to_gpu(model: nn.Module, device: torch.device) -> None: + """Move all non-DTensor params + buffers in ``model`` to ``device`` in-place. + + Used after CPU-offload loading: decoder DTensor shards stay on CPU (FSDP2 + streams them to GPU per layer), while root-level plain params and buffers + need to live on GPU so forwards work. + """ + for module in model.modules(): + for name, param in list(module._parameters.items()): + if param is None or isinstance(param, DTensor): + continue + module._parameters[name] = nn.Parameter( + param.data.to(device), requires_grad=param.requires_grad + ) + for name, buf in list(module._buffers.items()): + if buf is None or isinstance(buf, DTensor): + continue + module._buffers[name] = buf.to(device) + + +def _conversion_rules(model: nn.Module) -> dict | None: + """Rename + fuse rules for ``model``, read from transformers' own conversion table. + + Returns ``{"renames": {pattern: repl}, "fuses": [{"src_res", "target", "stack", "cat"}]}``, or + ``None`` when nothing needs converting (transformers<5 / already-matching checkpoint). A fuse + stacks its per-expert sources on ``stack`` and, when multi-source, concats them on ``cat``. + """ + renames = dict(getattr(model, "_checkpoint_conversion_mapping", None) or {}) + fuses = [] + for rule in get_model_conversion_mapping(model) if get_model_conversion_mapping else []: + sources, targets = list(rule.source_patterns), list(rule.target_patterns) + ops = getattr(rule, "operations", None) or [] # rename-only rules carry none + if not ops: + renames.update(dict.fromkeys(sources, targets[0])) + continue + dims = {type(op).__name__: op.dim for op in ops} + if len(targets) != 1 or set(dims) - {"MergeModulelist", "Concatenate"}: + raise NotImplementedError( + f"Unsupported conversion rule {sources} -> {targets} ({set(dims)})" + ) + fuses.append( + { + "src_res": [re.compile(s.replace("*", "([^.]+)")) for s in sources], + "target": targets[0], + "stack": dims["MergeModulelist"], + "cat": dims.get("Concatenate"), + } + ) + return {"renames": renames, "fuses": fuses} if (renames or fuses) else None + + +def _rename_key(key: str, renames: dict) -> str: + for old, new in renames.items(): + key = re.sub(old, new, key) + return key + + +def _resolve_fuse_source(key: str, fuses: list): + """Where ``key`` lands in a fused param, or ``None`` if it isn't a fuse source. + + On a hit returns ``(target_name, fuse, source_index, expert_idx)``: the fused param name, the + matched rule, which source slot matched (e.g. 0=gate, 1=up), and the expert index. + """ + for fuse in fuses: + for i, rx in enumerate(fuse["src_res"]): + m = rx.search(key) + if m: + return key[: m.start()] + fuse["target"] + key[m.end() :], fuse, i, m.group(1) + return None + + +def _target_name(rules: dict, key: str) -> str: + key = _rename_key(key, rules["renames"]) + hit = _resolve_fuse_source(key, rules["fuses"]) + return hit[0] if hit else key + + +def _convert_keys(rules: dict, state: dict) -> dict: + """Rename 1:1 keys and fuse per-expert keys into the model's fused params.""" + result, groups = {}, {} # groups[target] = (fuse, {source_index: {expert_idx: tensor}}) + for key, tensor in state.items(): + key = _rename_key(key, rules["renames"]) + hit = _resolve_fuse_source(key, rules["fuses"]) + if hit is None: + result[key] = tensor + continue + target, fuse, i, expert = hit + groups.setdefault(target, (fuse, {}))[1].setdefault(i, {})[expert] = tensor + for target, (fuse, by_src) in groups.items(): + stacks = [ + torch.stack([by_src[i][e] for e in sorted(by_src[i], key=int)], fuse["stack"]) + for i in range(len(fuse["src_res"])) + ] + result[target] = stacks[0] if fuse["cat"] is None else torch.cat(stacks, fuse["cat"]) + return result + + +def build_meta_causal_lm( + ckpt_path: str, + trust_remote_code: bool, + attn_implementation: str | None, + hf_config=None, +): + """Build a meta-init causal LM (no real storage allocated).""" + if hf_config is None: + config_kwargs: dict[str, Any] = {"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) + elif attn_implementation is not None: + # Honor the override even when the caller passed in a pre-fetched config. + hf_config._attn_implementation = attn_implementation + dtype = getattr(hf_config, "torch_dtype", None) or torch.bfloat16 + with init_empty_weights(include_buffers=False): + model = AutoModelForCausalLM.from_config( + hf_config, torch_dtype=dtype, trust_remote_code=trust_remote_code + ) + model.eval() + return model + + +def _layers_for_rank(n_layers: int, world_size: int, r: int) -> list[int]: + return [i for i in range(n_layers) if i % world_size == r] + + +def _read_and_convert( + resolved_path: str, weight_map: dict, keyset: set[str], rules: dict | None +) -> dict: + raw = read_safetensors_subset(resolved_path, weight_map, lambda k: k in keyset) + return _convert_keys(rules, raw) if rules else raw + + +def _read_owned_layers( + resolved_path: str, + weight_map: dict, + layer_sources: dict, + owned: list[int], + rules: dict | None, +) -> dict: + """Read + convert this rank's owned decoder layers from disk (ranks read in parallel).""" + return { + layer_idx: _read_and_convert( + resolved_path, weight_map, set(layer_sources[layer_idx]), rules + ) + for layer_idx in owned + } + + +def _broadcast_load_group( + group: list[int], + src: int, + rank: int, + owned: dict, + decoder_layers: list, + layer_prefixes: list, + device: torch.device, + cpu_offload: bool, +) -> None: + """Broadcast layers ``group`` from rank ``src`` to all ranks and reshard into their FSDP2 shards. + + The owner assembles the group's full tensors; every rank receives them, reshards its local slice, + then frees the full copy (capping the transient GPU peak). The owner drops its read copy after. + """ + big_dict: dict | None = None + if rank == src: + big_dict = {} + for i in group: + big_dict.update(owned[i]) + full = broadcast_state_dict(big_dict, src=src, device=device) + for layer_idx in group: + prefix = layer_prefixes[layer_idx] + stripped = {k[len(prefix) :]: v for k, v in full.items() if k.startswith(prefix)} + if cpu_offload: + stripped = {k: v.cpu() for k, v in stripped.items()} + set_model_state_dict( + decoder_layers[layer_idx], + stripped, + options=StateDictOptions(full_state_dict=True, broadcast_from_rank0=False), + ) + del stripped + del full + if rank == src: + for i in group: + del owned[i] + + +def _group_sources_by_layer( + weight_map: dict, rules: dict | None, model_param_names: set[str], layer_prefixes: list[str] +) -> tuple[dict[int, list[str]], list[str], int]: + """Bucket checkpoint keys by the decoder layer their converted target lives in. + + Returns ``(layer_sources, non_layer_sources, skipped)``: ``layer_sources[i]`` holds the keys + targeting decoder layer ``i``, ``non_layer_sources`` holds root (embed/lm_head/norm) keys, and + ``skipped`` counts keys whose target isn't in the model (aux weights, e.g. an MTP head). + """ + layer_sources: dict[int, list[str]] = {i: [] for i in range(len(layer_prefixes))} + non_layer_sources: list[str] = [] + skipped = 0 + for ckpt_key in weight_map: + target = _target_name(rules, ckpt_key) if rules else ckpt_key + if target not in model_param_names: + skipped += 1 + continue + for i, prefix in enumerate(layer_prefixes): + if target.startswith(prefix): + layer_sources[i].append(ckpt_key) + break + else: + non_layer_sources.append(ckpt_key) + return layer_sources, non_layer_sources, skipped + + +def parallel_load_and_prepare_fsdp2( + ckpt_path: str, + device: torch.device, + rank: int, + world_size: int, + trust_remote_code: bool = False, + mp_policy=None, + cpu_offload: bool = False, + attn_implementation: str | None = None, + hf_config=None, + broadcast_chunk_size: int | None = 8, +) -> nn.Module: + """Load and FSDP2-shard a HuggingFace causal LM via parallel safetensors reads. + + Round-robin assigns decoder layers to ranks; each rank reads only its owned + layers' weights from disk in parallel, then broadcasts to the others. Non-decoder + weights (embed, lm_head, norm) are read on rank 0 and broadcast. + + Requires an initialized ``torch.distributed`` process group (FSDP2's ``fully_shard`` + and the per-layer broadcasts both need it). A 1-rank PG (e.g. ``torchrun + --nproc_per_node=1``) is allowed; bare single-process is not. + + Pass ``hf_config`` if the caller has already fetched it (skips a redundant fetch). + + ``broadcast_chunk_size`` sets how many of a source's owned layers are broadcast per collective: + a smaller value lowers the peak transient GPU memory at the cost of more collectives (default 8; + pass ``None`` to broadcast all of a source's layers at once). + """ + resolved_path = _resolve_checkpoint_dir(ckpt_path, rank) + weight_map = weight_map_for(resolved_path) + + model = build_meta_causal_lm(resolved_path, trust_remote_code, attn_implementation, hf_config) + + # fsdp2_wrap shards each decoder layer + the root (embed/lm_head/norm sharded, not replicated). + decoder_layers = fsdp2_wrap(model, mp_policy=mp_policy, cpu_offload=cpu_offload) + module_to_name = {m: n for n, m in model.named_modules()} + layer_prefixes = [module_to_name[layer] + "." for layer in decoder_layers] + + # transformers>=5 fuses/renames checkpoint keys so they no longer match param names 1:1 + # (None => the pre-5.x identity path). + rules = _conversion_rules(model) + + # Valid targets; keys converting to anything else are aux weights (e.g. an MTP head) we skip. + model_param_names = {n for n, _ in chain(model.named_parameters(), model.named_buffers())} + + # Bucket each checkpoint key by its target's decoder layer (root params go to non_layer_sources). + layer_sources, non_layer_sources, skipped = _group_sources_by_layer( + weight_map, rules, model_param_names, layer_prefixes + ) + if skipped: + logger.debug( + "skipping %d checkpoint keys not present in the model (e.g. MTP head)", skipped + ) + + _materialize_meta_model(model, torch.device("cpu") if cpu_offload else device) + + owned_indices = _layers_for_rank(len(decoder_layers), world_size, rank) + owned = _read_owned_layers(resolved_path, weight_map, layer_sources, owned_indices, rules) + + # Smaller broadcast_chunk_size lowers the transient GPU peak (more, smaller collectives). + for src in range(world_size): + src_layers = _layers_for_rank(len(decoder_layers), world_size, src) + if not src_layers: + continue + chunk = broadcast_chunk_size or len(src_layers) + for start in range(0, len(src_layers), chunk): + _broadcast_load_group( + src_layers[start : start + chunk], + src, + rank, + owned, + decoder_layers, + layer_prefixes, + device, + cpu_offload, + ) + + # Non-decoder params: rank 0 reads + broadcasts; resharded into the root below. + # TODO: layerwise support. + non_layer = None + if rank == 0: + non_layer = _read_and_convert(resolved_path, weight_map, set(non_layer_sources), rules) + non_layer = broadcast_state_dict(non_layer, src=0, device=device) + if cpu_offload: + non_layer = {k: v.cpu() for k, v in non_layer.items()} + # shard_root=True makes the root params sharded DTensors, so reshard the full tensors via + # set_model_state_dict. strict=False: decoder keys are absent here (loaded above). + set_model_state_dict( + model, + non_layer, + options=StateDictOptions(full_state_dict=True, broadcast_from_rank0=False, strict=False), + ) + + if cpu_offload: + # Loaded on CPU for set_model_state_dict; FSDP2 streams decoder shards per forward, but + # the unwrapped root must live on GPU, so promote it. + _promote_non_dtensor_to_gpu(model, device) + if hasattr(model, "tie_weights"): + model.tie_weights() + model.requires_grad_(False) + return model diff --git a/tests/gpu/torch/quantization/test_fsdp2.py b/tests/gpu/torch/quantization/test_fsdp2.py index 6d47e1620ab..ea0444c9b10 100644 --- a/tests/gpu/torch/quantization/test_fsdp2.py +++ b/tests/gpu/torch/quantization/test_fsdp2.py @@ -303,3 +303,88 @@ def _test_persistent_materialization(rank, size): def test_persistent_materialization(dist_workers): dist_workers.run(_test_persistent_materialization) + + +def _test_writeback_root_unwrapped(rank, size): + """Writeback works when only the decoder layers are wrapped and the root is left + unsharded -- the layout ``fsdp2_wrap`` produces (root deliberately not wrapped) and + the one ``layerwise_calib`` save()/full_restore() rely on via + ``enable_weight_access_and_writeback(layer, model)``. + + Regression guard for the stale ``isinstance(root_model, FSDPModule)`` assert that + previously required the root itself to be FSDP-wrapped. + """ + from modelopt.torch.quantization.utils import enable_weight_access_and_writeback + + dim = 32 + torch.manual_seed(1) + # Root is a plain container; model[0] stands in for a decoder layer. + model = nn.Sequential(nn.Sequential(nn.Linear(dim, dim), nn.Linear(dim, dim))).cuda(rank) + synchronize_state_dict(model) + + # Wrap ONLY the "decoder layer" -- intentionally NO ``fully_shard(model)`` on the root, + # mirroring fsdp2_wrap. ``root_model`` (model) is therefore not an FSDPModule. + fully_shard(model[0]) + layer = model[0] + inputs = torch.randn(2, dim).cuda(rank) + + # Warmup forward to trigger FSDP2's lazy_init (mirrors layerwise calibration). + model(inputs) + + # This is the exact call save()/full_restore() make. Before the fix it tripped the + # ``assert isinstance(root_model, FSDPModule)`` because the root is unwrapped — that's + # the regression we guard. The DTensor-shape checks are not portable across torch + # versions when the root is not FSDP-wrapped, so we just verify the writeback path + # runs and mutations persist. + with enable_weight_access_and_writeback(layer[0], model): + ref_weight = layer[0].weight.clone() + layer[0].weight.data.add_(1.0) # mutate -> exercises the writeback path + + # Modification was written back into the shards. + with enable_weight_access_and_writeback(layer[0], model): + assert torch.allclose(layer[0].weight, ref_weight + 1.0) + + +def test_writeback_root_unwrapped(dist_workers): + dist_workers.run(_test_writeback_root_unwrapped) + + +def _test_writeback_cpu_offload(rank, size): + """Writeback round-trip when the FSDP2 shard is CPU-resident (``CPUOffloadPolicy``). + + Regression guard for the CPU↔GPU mirror added to + ``fsdp2_weight_access_and_writeback_context``: the gathered shard is on CPU, + so the helper mirrors it to GPU for in-context mutation and must copy + modifications back to the CPU shard on exit. + """ + from torch.distributed.fsdp import CPUOffloadPolicy + + from modelopt.torch.quantization.utils import enable_weight_access_and_writeback + + dim = 32 + torch.manual_seed(1) + model = nn.Sequential(nn.Sequential(nn.Linear(dim, dim), nn.Linear(dim, dim))).cuda(rank) + synchronize_state_dict(model) + + # Wrap the "decoder layer" with cpu_offload; root stays unwrapped. + fully_shard(model[0], offload_policy=CPUOffloadPolicy()) + layer = model[0] + + # Warmup forward triggers FSDP2's lazy_init. + model(torch.randn(2, dim).cuda(rank)) + + # Regression guard for the CPU→GPU mirror in fsdp2_weight_access_and_writeback_context: + # if the helper handed back a CPU tensor under cpu_offload, calibration ops would crash + # on the in-context mutation below (GPU activations vs CPU weight). The fact that this + # block runs and the mutation persists is the evidence the mirror trip worked. + with enable_weight_access_and_writeback(layer[0], model): + ref_weight = layer[0].weight.clone() + layer[0].weight.data.add_(1.0) + + # Mutation written back to the CPU shard. + with enable_weight_access_and_writeback(layer[0], model): + assert torch.allclose(layer[0].weight, ref_weight + 1.0) + + +def test_writeback_cpu_offload(dist_workers): + dist_workers.run(_test_writeback_cpu_offload) diff --git a/tests/gpu/torch/utils/test_model_load_utils.py b/tests/gpu/torch/utils/test_model_load_utils.py new file mode 100644 index 00000000000..8371f40d2d7 --- /dev/null +++ b/tests/gpu/torch/utils/test_model_load_utils.py @@ -0,0 +1,129 @@ +# SPDX-FileCopyrightText: Copyright (c) 2024 NVIDIA CORPORATION & AFFILIATES. All rights reserved. +# SPDX-License-Identifier: Apache-2.0 +# +# Licensed under the Apache License, Version 2.0 (the "License"); +# you may not use this file except in compliance with the License. +# You may obtain a copy of the License at +# +# http://www.apache.org/licenses/LICENSE-2.0 +# +# Unless required by applicable law or agreed to in writing, software +# distributed under the License is distributed on an "AS IS" BASIS, +# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. +# See the License for the specific language governing permissions and +# limitations under the License. + +"""GPU/distributed tests for the FSDP2 load path and its helpers.""" + +import json +import os +import tempfile +from functools import partial + +import pytest +import torch +import torch.distributed as dist +from torch.distributed.tensor import DTensor + + +def _test_broadcast_state_dict_roundtrip(rank, size): + """Round-trip from every rank as source (matches the per-layer rotation in the loader).""" + from modelopt.torch.utils.distributed import broadcast_state_dict + + device = torch.device(f"cuda:{rank}") + # Distinct payload per source rank so a wrong-src result would fail content checks. + for source in range(size): + src_dict = { + "w": torch.full((2, 4), float(source)), + "b": torch.tensor([float(source), float(source) + 1.0]), + } + out = broadcast_state_dict(src_dict if rank == source else None, src=source, device=device) + assert set(out.keys()) == {"w", "b"} + assert out["w"].device == device + assert torch.equal(out["w"].cpu(), src_dict["w"]) + assert torch.equal(out["b"].cpu(), src_dict["b"]) + + +def test_broadcast_state_dict_roundtrip(dist_workers): + dist_workers.run(_test_broadcast_state_dict_roundtrip) + + +def _build_tiny_llama_checkpoint(path: str) -> None: + """Write a tiny LlamaForCausalLM checkpoint (config + safetensors) to ``path``.""" + from transformers import LlamaConfig, LlamaForCausalLM + + config = LlamaConfig( + vocab_size=64, + hidden_size=32, + intermediate_size=64, + num_hidden_layers=2, + num_attention_heads=4, + max_position_embeddings=32, + torch_dtype="bfloat16", + ) + model = LlamaForCausalLM(config).to(torch.bfloat16) + model.save_pretrained(path) + + +def _test_parallel_load_and_export(rank, size, cpu_offload): + """Load a tiny Llama via the FSDP2 loader, forward, then export — config.architectures preserved. + + Parametrized over ``cpu_offload`` to cover both shard placements: + - off: decoder DTensor shards on GPU, plain root on GPU. + - on: decoder DTensor shards on CPU (streamed per layer), root promoted to GPU + via ``_promote_non_dtensor_to_gpu``. + """ + from modelopt.torch.export.unified_export_hf import export_hf_checkpoint + from modelopt.torch.utils.plugins.model_load_utils import parallel_load_and_prepare_fsdp2 + + suffix = "offload" if cpu_offload else "noffload" + ckpt_dir = os.path.join(tempfile.gettempdir(), f"_test_parallel_load_{suffix}_{os.getpid()}") + if rank == 0: + os.makedirs(ckpt_dir, exist_ok=True) + _build_tiny_llama_checkpoint(ckpt_dir) + dist.barrier() + + device = torch.device(f"cuda:{rank}") + model = parallel_load_and_prepare_fsdp2( + ckpt_dir, + device, + rank, + size, + cpu_offload=cpu_offload, + ) + + # Decoder layers AND root params (embed/lm_head) are sharded DTensors under shard_root=True. + decoder_params = list(model.model.layers[0].parameters()) + assert any(isinstance(p, DTensor) for p in decoder_params) + assert isinstance(model.model.embed_tokens.weight, DTensor) + if not cpu_offload: + # Non-offload: the root's local shard lives on GPU. + assert model.model.embed_tokens.weight.to_local().device.type == "cuda" + if cpu_offload: + # Under cpu_offload the decoder shards live on CPU between forwards. + decoder_dtensors = [p for p in decoder_params if isinstance(p, DTensor)] + assert all(p.to_local().device.type == "cpu" for p in decoder_dtensors) + + # Forward exercises FSDP2 hooks + (under cpu_offload) the per-layer CPU↔GPU stream. + input_ids = torch.randint(0, 64, (1, 8), device=device) + out = model(input_ids=input_ids).logits + assert out.shape == (1, 8, 64) + + # Export and verify the saved config.json retains the original architectures. + export_dir = os.path.join( + tempfile.gettempdir(), f"_test_parallel_export_{suffix}_{os.getpid()}" + ) + if rank == 0: + os.makedirs(export_dir, exist_ok=True) + dist.barrier() + export_hf_checkpoint(model, export_dir=export_dir, dtype=torch.bfloat16) + + if rank == 0: + with open(os.path.join(export_dir, "config.json")) as f: + cfg = json.load(f) + assert cfg["architectures"] == ["LlamaForCausalLM"] + + +@pytest.mark.parametrize("cpu_offload", [False, True]) +def test_parallel_load_and_export(dist_workers, cpu_offload): + dist_workers.run(partial(_test_parallel_load_and_export, cpu_offload=cpu_offload)) diff --git a/tests/unit/torch/utils/test_model_load_utils.py b/tests/unit/torch/utils/test_model_load_utils.py new file mode 100644 index 00000000000..c5cc7a6d324 --- /dev/null +++ b/tests/unit/torch/utils/test_model_load_utils.py @@ -0,0 +1,156 @@ +# SPDX-FileCopyrightText: Copyright (c) 2024 NVIDIA CORPORATION & AFFILIATES. All rights reserved. +# SPDX-License-Identifier: Apache-2.0 +# +# Licensed under the Apache License, Version 2.0 (the "License"); +# you may not use this file except in compliance with the License. +# You may obtain a copy of the License at +# +# http://www.apache.org/licenses/LICENSE-2.0 +# +# Unless required by applicable law or agreed to in writing, software +# distributed under the License is distributed on an "AS IS" BASIS, +# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. +# See the License for the specific language governing permissions and +# limitations under the License. + +"""Pure-function tests for ``modelopt.torch.utils.plugins.model_load_utils``.""" + +import json + +import pytest +import torch +from packaging.version import Version +from safetensors.torch import save_file + +pytest.importorskip("accelerate") + +from modelopt.torch.utils.plugins.model_load_utils import ( + _conversion_rules, + _convert_keys, + _target_name, + read_safetensors_subset, + weight_map_for, +) + + +def test_weight_map_for_sharded(tmp_path): + save_file({"a.weight": torch.zeros(2)}, str(tmp_path / "shard1.safetensors")) + save_file({"b.weight": torch.zeros(2)}, str(tmp_path / "shard2.safetensors")) + (tmp_path / "model.safetensors.index.json").write_text( + json.dumps( + {"weight_map": {"a.weight": "shard1.safetensors", "b.weight": "shard2.safetensors"}} + ) + ) + + assert weight_map_for(str(tmp_path)) == { + "a.weight": "shard1.safetensors", + "b.weight": "shard2.safetensors", + } + + +def test_weight_map_for_single_file(tmp_path): + save_file( + {"a.weight": torch.zeros(2), "b.weight": torch.zeros(2)}, + str(tmp_path / "model.safetensors"), + ) + + assert weight_map_for(str(tmp_path)) == { + "a.weight": "model.safetensors", + "b.weight": "model.safetensors", + } + + +def test_weight_map_for_missing(tmp_path): + with pytest.raises(RuntimeError, match="No safetensors checkpoint"): + weight_map_for(str(tmp_path)) + + +def test_read_safetensors_subset(tmp_path): + save_file( + {"a.weight": torch.tensor([1.0, 2.0]), "a.bias": torch.tensor([3.0])}, + str(tmp_path / "shard1.safetensors"), + ) + save_file({"b.weight": torch.tensor([4.0])}, str(tmp_path / "shard2.safetensors")) + weight_map = { + "a.weight": "shard1.safetensors", + "a.bias": "shard1.safetensors", + "b.weight": "shard2.safetensors", + } + + result = read_safetensors_subset(str(tmp_path), weight_map, lambda n: n.startswith("a.")) + + assert set(result.keys()) == {"a.weight", "a.bias"} + assert torch.equal(result["a.weight"], torch.tensor([1.0, 2.0])) + assert torch.equal(result["a.bias"], torch.tensor([3.0])) + + +def _build_tiny_qwen3_moe(): + """A tiny meta-init Qwen3-MoE (fused ``gate_up_proj`` + ``down_proj`` experts) for converter tests.""" + transformers = pytest.importorskip("transformers") + if Version(transformers.__version__) < Version("5.0"): + pytest.skip("multi-source fused-MoE conversion needs transformers>=5") + from accelerate import init_empty_weights + from transformers import AutoConfig, AutoModelForCausalLM + + cfg = AutoConfig.for_model( + "qwen3_moe", + hidden_size=8, + intermediate_size=16, + moe_intermediate_size=6, + num_hidden_layers=1, + num_attention_heads=2, + num_key_value_heads=2, + head_dim=4, + num_experts=4, + num_experts_per_tok=2, + vocab_size=32, + max_position_embeddings=16, + ) + try: + with init_empty_weights(): + model = AutoModelForCausalLM.from_config(cfg) + except Exception as e: # modeling class unavailable in this transformers build + pytest.skip(f"qwen3_moe modeling unavailable: {e}") + return model, cfg + + +def test_checkpoint_key_converter_multisource_expert_fusion(): + """A multi-source fused MoE (Qwen3: gate_proj+up_proj -> gate_up_proj) converts correctly. + + Exercises the multi-source path (gate/up concatenated after expert stacking) AND the + single-source path (down_proj is a plain expert stack) in one model. + """ + model, cfg = _build_tiny_qwen3_moe() + rules = _conversion_rules(model) + assert rules is not None + + names = dict(model.named_parameters()) + gname = next(n for n in names if n.endswith("mlp.experts.gate_up_proj")) + prefix = gname[: -len("mlp.experts.gate_up_proj")] + n_exp, inter, hidden = cfg.num_experts, cfg.moe_intermediate_size, cfg.hidden_size + + gate = [torch.randn(inter, hidden) for _ in range(n_exp)] + up = [torch.randn(inter, hidden) for _ in range(n_exp)] + down = [torch.randn(hidden, inter) for _ in range(n_exp)] + state = {} + for e in range(n_exp): + state[f"{prefix}mlp.experts.{e}.gate_proj.weight"] = gate[e] + state[f"{prefix}mlp.experts.{e}.up_proj.weight"] = up[e] + state[f"{prefix}mlp.experts.{e}.down_proj.weight"] = down[e] + + out = _convert_keys(rules, state) + + # Multi-source: experts stacked (dim 0) then gate|up concatenated (dim 1), gate first. + gate_up = out[f"{prefix}mlp.experts.gate_up_proj"] + assert gate_up.shape == (n_exp, 2 * inter, hidden) == tuple(names[gname].shape) + assert torch.equal(gate_up, torch.cat([torch.stack(gate), torch.stack(up)], dim=1)) + + # Single-source (regression): down_proj is a plain expert stack. + down_proj = out[f"{prefix}mlp.experts.down_proj"] + assert down_proj.shape == (n_exp, hidden, inter) + assert torch.equal(down_proj, torch.stack(down)) + + # Name-only mapping routes every source key to the fused target. + for e in range(n_exp): + assert _target_name(rules, f"{prefix}mlp.experts.{e}.gate_proj.weight") == gname + assert _target_name(rules, f"{prefix}mlp.experts.{e}.up_proj.weight") == gname