From 445e10728d44a103ebe807c2edfb55424edf8ec2 Mon Sep 17 00:00:00 2001 From: dimapihtar Date: Tue, 5 May 2026 15:13:24 -0700 Subject: [PATCH 1/7] fix docs Signed-off-by: dimapihtar --- docs/llama_mistral.md | 130 ++++++++---------------------------------- 1 file changed, 23 insertions(+), 107 deletions(-) diff --git a/docs/llama_mistral.md b/docs/llama_mistral.md index 95568adce78..2754405610c 100644 --- a/docs/llama_mistral.md +++ b/docs/llama_mistral.md @@ -22,13 +22,11 @@ Architecturally Llama-2, Llama-3 and Mistral-7b are very similar. As such Megatr - [Llama, Mistral and other Llama-like model support in Megatron-LM](#llama-mistral-and-other-llama-like-model-support-in-megatron-lm) - [Contents](#contents) - [Llama-2](#llama-2) - - [Download Meta or Huggingface checkpoints](#download-meta-or-huggingface-checkpoints) + - [Download Huggingface checkpoints](#download-huggingface-checkpoints) - [Convert checkpoint format](#convert-checkpoint-format) - - [Meta format](#meta-format) - [Huggingface format](#huggingface-format) - [Launch model](#launch-model) - [Launch Megatron](#launch-megatron) - - [Launch Meta](#launch-meta) - [Launch Huggingface](#launch-huggingface) - [Benchmark results](#benchmark-results) - [Big Bench](#big-bench) @@ -48,72 +46,35 @@ Architecturally Llama-2, Llama-3 and Mistral-7b are very similar. As such Megatr - [Launch model](#launch-model) - [Other Llama-like model support](#other-llama-like-model-support) - [Known numerical differences](#known-numerical-differences) -- [Using legacy model format](#using-legacy-model-format) # Llama-2 Llama-2 checkpoints can be loaded into Megatron for inference and for finetuning. Loading these checkpoints consists of three steps: 1. Get access to download the checkpoints. -2. Convert the checkpoints from Meta/Huggingface format to Megatron format. +2. Convert the checkpoints from Huggingface format to Megatron format. 3. Setup arguments for launching the model. The following sections detail these steps. The final section lists benchmark result comparisons between: 1) Llama-2 inference code running the Meta-format checkpoints, and 2) Megatron inference code running the converted checkpoints. -## Download Meta or Huggingface checkpoints +## Download Huggingface checkpoints -Users must first apply for access to download the Llama-2 checkpoints either directly [Huggingface](https://huggingface.co/docs/transformers/main/model_doc/llama2) (HF). The checkpoints are available in two formats, Meta's native format (available from both the Meta and HF links), and HF's format (available only from HF). Either format can be converted to Megatron, as detailed next. +Users must first apply for access to download the Llama-2 checkpoints either directly [Huggingface](https://huggingface.co/docs/transformers/main/model_doc/llama2) (HF). The checkpoints are available in HF's format (available only from HF). HF format can be converted to Megatron, as detailed next. ## Convert checkpoint format We recommend passing `--dtype bf16` for training or finetuning. Inference can be done in bfloat16 or float16. -### Meta format - -The Meta format checkpoints are converted to HF format as an intermediate step before converting to Megatron format. The `transformers` package is required, and must have version >=4.31.0 (e.g., `pip install transformers>=4.31.0`). (**Note**: we have specifically tested with versions `4.31.0` and `4.32.0`; your experience may vary with newer versions.) Assuming the downloaded checkpoints are in `$CHECKPOINT_DIR` (with separate sub-directories for 7B, 13B, 70B, etc.), the following example command can be used to convert from Llama-2 format to HF format in bfloat16: - -``` -python tools/checkpoint/convert.py \ -> --model-type GPT \ -> --loader llama_mistral \ -> --load-dir ${META_FORMAT_DIR} \ -> --model-size ${MODEL_SIZE} \ -> --checkpoint-type meta \ -> --tokenizer-model ${TOKENIZER_MODEL} \ -> --saver core \ -> --save-dir ${MEGATRON_FORMAT_DIR} \ -> --target-tensor-parallel-size ${TP} \ -> --target-pipeline-parallel-size ${PP} \ -> --bf16 -``` - -Valid values for `--model-size` are `llama2-7B`, `llama2-13B`, and `llama2-70B` (for pretrained-only models), and `llama2-7Bf`, `llama2-13Bf`, and `llama2-70Bf` (for chat-finetuned models). - ### Huggingface format -The HF checkpoints can be converted to Megatron format by using Megatron's own Llama-2 checkpoint converter for HF format (see script `tools/checkpoint/loader_llama_mistral.py`). One important argument that must be set correctly is the tensor parallel size (`TP`) for each model. The following table shows these values: - -| Model size | Tensor parallel size (`TP`) | -| ---------- | --------------------------- | -| 7B | 1 | -| 13B | 2 | -| 70B | 8 | - -Using these values for `TP`, along with the path to the Llama-2 tokenizer model (automatically downloaded with original checkpoint download; see `${TOKENIZER_MODEL}` below), run the following command from the root of your Megatron source code to convert from HF format to Megatron format: +The HF checkpoints can be converted to Megatron format by using Megatron-Bridge's checkpoint converter for HF format [see script](https://github.com/NVIDIA-NeMo/Megatron-Bridge/blob/main/examples/conversion/convert_checkpoints.py). ``` -python tools/checkpoint/convert.py \ -> --model-type GPT \ -> --loader llama_mistral \ -> --load-dir ${HF_FORMAT_DIR} \ -> --model-size ${MODEL_SIZE} \ -> --checkpoint-type hf \ -> --tokenizer-model ${TOKENIZER_MODEL} \ -> --saver core \ -> --save-dir ${MEGATRON_FORMAT_DIR} \ -> --target-tensor-parallel-size ${TP} \ -> --target-pipeline-parallel-size ${PP} \ -> --bf16 +python Megatron-Bridge/examples/conversion/convert_checkpoints.py import \ + --hf-model meta-llama/Llama-2-7B \ + --megatron-path ./checkpoints/llama2_7b \ + --torch-dtype bfloat16 \ + --device-map auto ``` After this conversion, we are ready to load the checkpoints into a Megatron GPT model. @@ -144,12 +105,6 @@ If loading for either inference or finetuning, use the following arguments: --attention-softmax-in-fp32 ``` -**Note:** If you converted to the legacy model format (i.e., `--saver legacy`), please see [here](#using-legacy-model-format). - -### Launch Meta - -Meta checkpoints can be launched with: - ### Launch Huggingface Huggingface checkpoints can be launched with: @@ -243,29 +198,14 @@ We recommend passing `--dtype bf16` for training or finetuning. Inference can be ### Huggingface format -The HF checkpoints can be converted to Megatron format by using Megatron's own Llama-3.x checkpoint converter for HF format (see script `tools/checkpoint/loader_llama_mistral.py`). One important argument that must be set correctly is the tensor parallel size (`TP`) for each model. The following table shows these values: - -| Model size | Tensor parallel size (`TP`) | -| ---------- | --------------------------- | -| 1B | 1 | -| 3B | 1 | -| 8B | 1 | -| 70B | 8 | - -Using these values for `TP`, along with the path to the Llama-3.x tokenizer model (automatically downloaded with original checkpoint download; see `${TOKENIZER_MODEL}` below), run the following command from the root of your Megatron source code to convert from HF format to Megatron format: +The HF checkpoints can be converted to Megatron format by using Megatron-Bridge's checkpoint converter for HF format [see script](https://github.com/NVIDIA-NeMo/Megatron-Bridge/blob/main/examples/conversion/convert_checkpoints.py). ``` -$>: python tools/checkpoint/convert.py \ - > --bf16 \ - > --model-type GPT \ - > --loader llama_mistral \ - > --saver core \ - > --target-tensor-parallel-size ${TP} \ - > --checkpoint-type hf \ - > --load-dir ${HF_FORMAT_DIR} \ - > --save-dir ${MEGATRON_FORMAT_DIR} \ - > --tokenizer-model ${TOKENIZER_MODEL} \ - > --model-size llama3 \ +python Megatron-Bridge/examples/conversion/convert_checkpoints.py import \ + --hf-model meta-llama/Llama-3.2-1B \ + --megatron-path ./checkpoints/llama3_2_1b \ + --torch-dtype bfloat16 \ + --device-map auto ``` After this conversion, we are ready to load the checkpoints into a Megatron GPT model. @@ -345,8 +285,6 @@ For Llama3.1 please use the following arguments: --bf16 \ ``` -**Note:** If you converted to the legacy model format (i.e., `--saver legacy`), please see [here](#using-legacy-model-format). - # Mistral-7b Megatron currently supports loading the v0.3 release of Mistral-7b (which does not use sliding window attention and offers a larger 32768 vocabulary) for inference and finetuning. Loading these checkpoints consists of several steps: @@ -364,25 +302,17 @@ Users must first apply for access to download the Mistral-7b checkpoints through ## Convert checkpoint format -The HF checkpoints can be converted to Megatron format by using Megatron's own Mistral checkpoint converter for HF format (see script `tools/checkpoint/loader_llama_mistral.py`). - -Using the path to the Mistral tokenizer model (downloaded alongside the HF checkpoint), run the following command from the root of your Megatron source code to convert from HF format to the Megatron core format: +The HF checkpoints can be converted to Megatron format by using Megatron-Bridge's checkpoint converter for HF format [see script](https://github.com/NVIDIA-NeMo/Megatron-Bridge/blob/main/examples/conversion/convert_checkpoints.py). ``` -$>: python tools/checkpoint/convert.py \ - > --bf16 \ - > --model-type GPT \ - > --loader llama_mistral \ - > --saver core \ - > --target-tensor-parallel-size ${TP} \ - > --checkpoint-type hf \ - > --load-dir ${HF_FORMAT_DIR} \ - > --save-dir ${MEGATRON_FORMAT_DIR} \ - > --tokenizer-model ${TOKENIZER_MODEL} \ - > --model-size mistral \ +python Megatron-Bridge/examples/conversion/convert_checkpoints.py import \ + --hf-model mistralai/Mistral-7B-Instruct-v0.3 \ + --megatron-path ./checkpoints/mistral_7b \ + --torch-dtype bfloat16 \ + --device-map auto ``` -After this conversion, we are ready to load the checkpoints into a Megatron core GPT model. +After this conversion, we are ready to load the checkpoints into a Megatron GPT model. ## (Optional) Validate checkpoints @@ -424,8 +354,6 @@ If loading for either inference or finetuning, use the following arguments: --num-attention-heads 32 ``` -**Note:** If you converted to the legacy model format (i.e., `--saver legacy`), please see [here](#using-legacy-model-format). - # Other Llama-like model support *Note: Experimental* @@ -438,15 +366,3 @@ It is not expected that the megatron and Huggingface implementations of llama3.x 1. TransformerEngine (TE) uses the model params_dtype inside RMSNorm whereas the Huggingface implementation uses fp32. See for details: 2. Huggingface `transformers` implements the q, k and v projections in self-attention as separate GEMMs whereas Megatron core combines them into a single GEMM for efficiency. This leads to small numerical differences. - -# Using legacy model format - -In all the checkpoint conversion examples used in this document, the saver format `--saver core` is used, signifying that the newer (and recommended) Megatron GPT model class will be used. I.e.: - -- old class: `megatron.legacy.model.gpt_model.GPTModel` -- new class: `megatron.core.models.gpt.gpt_model.GPTModel` - -Using this new format is the recommended approach. However, if your use case requires using the older class (i.e., convert using `--saver legacy`), then when launching training or finetuning, the following args must be added: - -- `--use-legacy-models`: use the older model class -- `--ckpt-format torch`: use the `torch` checkpoint format, which is the only checkpoint format that is compatible with the legacy model format From 75b1e682e33306bfb08bd926d1618082837844e2 Mon Sep 17 00:00:00 2001 From: dimapihtar Date: Fri, 3 Jul 2026 06:58:54 -0700 Subject: [PATCH 2/7] move modelopt dist chcekpointing plugins to mcore Signed-off-by: dimapihtar --- megatron/post_training/checkpointing.py | 100 +++++++++++++++++++++++- megatron/training/checkpointing.py | 59 +++++++++++++- 2 files changed, 154 insertions(+), 5 deletions(-) diff --git a/megatron/post_training/checkpointing.py b/megatron/post_training/checkpointing.py index 4449a36221b..970792256d0 100644 --- a/megatron/post_training/checkpointing.py +++ b/megatron/post_training/checkpointing.py @@ -1,15 +1,20 @@ # Copyright (c) 2024, NVIDIA CORPORATION. All rights reserved. import logging +import os from pathlib import Path -from typing import Optional, Tuple, Union +from typing import Any, Optional, Tuple, Union +import modelopt import modelopt.torch.opt as mto +from modelopt.torch.utils import safe_load +import torch import torch.nn as nn -from modelopt.torch.opt.plugins import restore_sharded_modelopt_state from megatron.core import dist_checkpointing -from megatron.core.utils import get_torch_version, is_torch_min_version, unwrap_model +from megatron.core.dist_checkpointing.strategies.torch import TorchDistLoadShardedStrategy +from megatron.core.dist_checkpointing.validation import StrictHandling +from megatron.core.utils import unwrap_model from megatron.training import get_args from megatron.training.checkpointing import _load_base_checkpoint, load_checkpoint from megatron.training.utils import print_rank_0 @@ -18,6 +23,7 @@ logger = logging.getLogger(__name__) NEMO_WEIGHT_DIR_NAMES = {"model_weights": "model.", "weights": "module."} +COMMON_STATE_FNAME = "common.pt" def has_modelopt_state(checkpoint_path: str) -> bool: @@ -194,4 +200,90 @@ def _remove_prefix_state_dict_pre_hook( unwrapped_model[0].load_state_dict(model_state_dict, strict=False) print_distributed_quant_summary(unwrapped_model[0]) else: - _ = load_checkpoint(model, optimizer, opt_param_scheduler, strict=strict, load_arg=load_arg) \ No newline at end of file + _ = load_checkpoint(model, optimizer, opt_param_scheduler, strict=strict, load_arg=load_arg) + + +def restore_sharded_modelopt_state( + model: list[torch.nn.Module], + checkpoint_name: str | Path, + prefix: str = "", + metadata: dict[str, Any] | None = None, +) -> None: + """Restore modelopt_state from the sharded state_dict format. + + Args: + model: the model to restore the modelopt optimization + checkpoint_name: the checkpoint folder path + prefix: the prefix to add to the modelopt_state keys ("model." for NeMo) + metadata: the metadata for distributed checkpointing + + Note: + The metadata includes several breaking changes. For example, `singleton_local_shards` + is set to `True` (was not set before) in megatron-core-0.15.0. This flag affects the + sharded state_dict format and must be consistent between saving and loading. + """ + if len(model) > 1: + raise ValueError("sharded_modelopt_state does not support virtual pipeline parallel!") + + modelopt_checkpoint_name = f"{checkpoint_name}/modelopt_state" + + # Early return if the model already has a modelopt_state or the checkpoint does not exist. + if not os.path.exists(modelopt_checkpoint_name) or mto.ModeloptStateManager.is_converted( + model[0] + ): + return + + # Loading the common modelopt_state (replicated on all ranks). + # Detect format: legacy checkpoints store common state in a standalone common.pt file; + # newer sharded checkpoints store it as a ShardedObject inside the torch_dist checkpoint. + legacy_common_path = os.path.join(modelopt_checkpoint_name, COMMON_STATE_FNAME) + if os.path.exists(legacy_common_path): + common_modelopt_state = safe_load(legacy_common_path) + else: + common_modelopt_state = dist_checkpointing.load_common_state_dict(modelopt_checkpoint_name) + + modelopt_load_version = common_modelopt_state["modelopt_version"] + + print_rank_0(f"nvidia-modelopt ckpt/inst version: {modelopt_load_version}/{modelopt.__version__}") + + model[0] = mto.restore_from_modelopt_state(model[0], common_modelopt_state) + + _load_extra_state_from_sharded_checkpoint(model[0], checkpoint_name, prefix, metadata=metadata) + + +def _load_extra_state_from_sharded_checkpoint( + model: torch.nn.Module, + checkpoint_name: str | Path, + prefix: str, + metadata: dict[str, Any] | None = None, +) -> None: + """Load extra state from sharded checkpoint. + + Note: since extra_state is a subset of full the sharded_state_dict, we use + strict=StrictHandling.LOG_UNEXPECTED instead of LOG_ALL. + + Args: + model: the model to load extra state into + checkpoint_name: the checkpoint folder path + prefix: the prefix to add to the modelopt_state keys + metadata: the metadata for distributed checkpointing + + Note: + The metadata includes several breaking changes. For example, `singleton_local_shards` + is set to `True` (was not set before) in megatron-core-0.15.0. This flag affects the + sharded state_dict format and must be consistent between saving and loading. + """ + sharded_state_dict = model.sharded_state_dict(prefix=prefix) + extra_sharded_state_dict = {k: v for k, v in sharded_state_dict.items() if "_extra_state" in k} + extra_state_dict = dist_checkpointing.load( + extra_sharded_state_dict, + checkpoint_name, + TorchDistLoadShardedStrategy(), + strict=StrictHandling.LOG_UNEXPECTED, + ) + extra_state_dict_no_prefix = {} + + for k, v in extra_state_dict.items(): + if k.startswith(prefix): + extra_state_dict_no_prefix[k[len(prefix) :]] = v + model.load_state_dict(extra_state_dict_no_prefix, strict=False) diff --git a/megatron/training/checkpointing.py b/megatron/training/checkpointing.py index ad9de5b15b2..f01df7b97e2 100644 --- a/megatron/training/checkpointing.py +++ b/megatron/training/checkpointing.py @@ -3,6 +3,7 @@ """Input/output checkpointing.""" import contextlib +import copy import inspect import multiprocessing import os @@ -65,9 +66,12 @@ # [ModelOpt]: Import try: - from modelopt.torch.opt.plugins import save_modelopt_state, save_sharded_modelopt_state + import modelopt.torch.utils.distributed as dist + import modelopt.torch.opt as mto + from modelopt.torch.opt.plugins import save_modelopt_state from megatron.post_training.utils import print_distributed_quant_summary + has_nvidia_modelopt = True except Exception: has_nvidia_modelopt = False @@ -2286,3 +2290,56 @@ def load_biencoder_checkpoint(model, only_query_model=False, print(' successfully loaded {}'.format(checkpoint_name)) return model + + +def remove_per_module_state( + modelopt_state: dict[str, Any], +) -> None: + """Remove metadata from the modelopt_state. + + The metadata of the modelopt_state contains keys which may change with different pipeline + and expert parallelism. As a result, the metadata must be stored as several ShardedObject with + global and local layer offset mapping. + + Args: + modelopt_state: the state_dict that contains all algorithms that have been applied + to the given model. + """ + if "modelopt_state_dict" not in modelopt_state: + return + + for mode, config in modelopt_state["modelopt_state_dict"]: + metadata = config.get("metadata", None) + if metadata is not None: + _ = metadata.pop("quantizer_state", None) + _ = metadata.pop("subnet_config", None) + _ = metadata.pop("real_quantizer_state", None) + _ = metadata.pop("q_tensor_state", None) + else: + config["metadata"] = {} + + +def save_sharded_modelopt_state( + model: list[torch.nn.Module], + checkpoint_name: str | Path, + sharded_strategy: tuple[str, int] | None = None, + prefix: str = "", +) -> None: + """Save modelopt_state in the sharded state_dict format. + + Args: + model: the model to restore the modelopt optimization + checkpoint_name: the checkpoint folder path + sharded_strategy: configures sharded tensors saving behavior and backend + prefix: the prefix to add to the modelopt_state keys ("model." for NeMo) + """ + if not mto.ModeloptStateManager.is_converted(model[0]): + return + if len(model) > 1: + raise ValueError("sharded_modelopt_state does not support virtual pipeline parallel!") + modelopt_checkpoint_name = f"{checkpoint_name}/modelopt_state" + if dist.is_master(): + os.makedirs(modelopt_checkpoint_name, exist_ok=True) + modelopt_state = copy.deepcopy(mto.modelopt_state(model[0])) + remove_per_module_state(modelopt_state) + dist_checkpointing.save(modelopt_state, modelopt_checkpoint_name, sharded_strategy) From c2e5206e3fd15ab486b1caf12280b2af9c77add1 Mon Sep 17 00:00:00 2001 From: dimapihtar Date: Mon, 6 Jul 2026 03:01:16 -0700 Subject: [PATCH 3/7] move modelopt modules to dist_checkpointing.strategies Signed-off-by: dimapihtar --- .../dist_checkpointing/strategies/modelopt.py | 187 ++++++++++++++++++ megatron/post_training/checkpointing.py | 100 +--------- megatron/training/checkpointing.py | 59 +----- 3 files changed, 192 insertions(+), 154 deletions(-) create mode 100644 megatron/core/dist_checkpointing/strategies/modelopt.py diff --git a/megatron/core/dist_checkpointing/strategies/modelopt.py b/megatron/core/dist_checkpointing/strategies/modelopt.py new file mode 100644 index 00000000000..da0be0d3468 --- /dev/null +++ b/megatron/core/dist_checkpointing/strategies/modelopt.py @@ -0,0 +1,187 @@ +# Copyright (c) 2026, NVIDIA CORPORATION. All rights reserved. + +"""Dist checkpointing modules needed for ModelOpt.""" + +import copy +import logging +import os +from pathlib import Path +from typing import Any + +import torch + +from .torch import TorchDistLoadShardedStrategy +from ..serialization import load, load_common_state_dict, save +from ..validation import StrictHandling + +from megatron.core import mpu +from megatron.core.safe_globals import safe_load_from_bytes + +try: + import modelopt + import modelopt.torch.utils.distributed as dist + import modelopt.torch.opt as mto + + has_nvidia_modelopt = True +except (ModuleNotFoundError, ImportError): + has_nvidia_modelopt = False + +logger = logging.getLogger(__name__) + + +def remove_per_module_state( + modelopt_state: dict[str, Any], +) -> None: + """Remove metadata from the modelopt_state. + + The metadata of the modelopt_state contains keys which may change with different pipeline + and expert parallelism. As a result, the metadata must be stored as several ShardedObject with + global and local layer offset mapping. + + Args: + modelopt_state: the state_dict that contains all algorithms that have been applied + to the given model. + """ + if "modelopt_state_dict" not in modelopt_state: + return + + for mode, config in modelopt_state["modelopt_state_dict"]: + metadata = config.get("metadata", None) + if metadata is not None: + _ = metadata.pop("quantizer_state", None) + _ = metadata.pop("subnet_config", None) + _ = metadata.pop("real_quantizer_state", None) + _ = metadata.pop("q_tensor_state", None) + else: + config["metadata"] = {} + + +def save_modelopt_state(model: list[torch.nn.Module], state_dict: dict[str, Any]) -> None: + """Save modelopt_state as a part of the per rank state_dict. + + NOTE: Only used for Megatron-LM. + + Args: + model: the modelopt optimized model + state_dict: the current modelopt optimized model state_dict to store + """ + if not mto.ModeloptStateManager.is_converted(model[0]): + return + if len(model) == 1: + state_dict["modelopt_state"] = mto.modelopt_state(model[0]) + else: + for i in range(len(model)): + mpu.set_virtual_pipeline_model_parallel_rank(i) + state_dict[f"modelopt_state_{i}"] = mto.modelopt_state(model[i]) + + +def save_sharded_modelopt_state( + model: list[torch.nn.Module], + checkpoint_name: str | Path, + sharded_strategy: tuple[str, int] | None = None, + prefix: str = "", +) -> None: + """Save modelopt_state in the sharded state_dict format. + + Args: + model: the model to restore the modelopt optimization + checkpoint_name: the checkpoint folder path + sharded_strategy: configures sharded tensors saving behavior and backend + prefix: the prefix to add to the modelopt_state keys ("model." for NeMo) + """ + if not mto.ModeloptStateManager.is_converted(model[0]): + return + if len(model) > 1: + raise ValueError("sharded_modelopt_state does not support virtual pipeline parallel!") + modelopt_checkpoint_name = f"{checkpoint_name}/modelopt_state" + if dist.is_master(): + os.makedirs(modelopt_checkpoint_name, exist_ok=True) + modelopt_state = copy.deepcopy(mto.modelopt_state(model[0])) + remove_per_module_state(modelopt_state) + save(modelopt_state, modelopt_checkpoint_name, sharded_strategy) + + +def _load_extra_state_from_sharded_checkpoint( + model: torch.nn.Module, + checkpoint_name: str | Path, + prefix: str, + metadata: dict[str, Any] | None = None, +) -> None: + """Load extra state from sharded checkpoint. + + Note: since extra_state is a subset of full the sharded_state_dict, we use + strict=StrictHandling.LOG_UNEXPECTED instead of LOG_ALL. + + Args: + model: the model to load extra state into + checkpoint_name: the checkpoint folder path + prefix: the prefix to add to the modelopt_state keys + metadata: the metadata for distributed checkpointing + + Note: + The metadata includes several breaking changes. For example, `singleton_local_shards` + is set to `True` (was not set before) in megatron-core-0.15.0. This flag affects the + sharded state_dict format and must be consistent between saving and loading. + """ + sharded_state_dict = model.sharded_state_dict(prefix=prefix) + extra_sharded_state_dict = {k: v for k, v in sharded_state_dict.items() if "_extra_state" in k} + extra_state_dict = load( + extra_sharded_state_dict, + checkpoint_name, + TorchDistLoadShardedStrategy(), + strict=StrictHandling.LOG_UNEXPECTED, + ) + extra_state_dict_no_prefix = {} + + for k, v in extra_state_dict.items(): + if k.startswith(prefix): + extra_state_dict_no_prefix[k[len(prefix) :]] = v + model.load_state_dict(extra_state_dict_no_prefix, strict=False) + + +def restore_sharded_modelopt_state( + model: list[torch.nn.Module], + checkpoint_name: str | Path, + prefix: str = "", + metadata: dict[str, Any] | None = None, +) -> None: + """Restore modelopt_state from the sharded state_dict format. + + Args: + model: the model to restore the modelopt optimization + checkpoint_name: the checkpoint folder path + prefix: the prefix to add to the modelopt_state keys ("model." for NeMo) + metadata: the metadata for distributed checkpointing + + Note: + The metadata includes several breaking changes. For example, `singleton_local_shards` + is set to `True` (was not set before) in megatron-core-0.15.0. This flag affects the + sharded state_dict format and must be consistent between saving and loading. + """ + if len(model) > 1: + raise ValueError("sharded_modelopt_state does not support virtual pipeline parallel!") + + modelopt_checkpoint_name = f"{checkpoint_name}/modelopt_state" + + # Early return if the model already has a modelopt_state or the checkpoint does not exist. + if not os.path.exists(modelopt_checkpoint_name) or mto.ModeloptStateManager.is_converted( + model[0] + ): + return + + # Loading the common modelopt_state (replicated on all ranks). + # Detect format: legacy checkpoints store common state in a standalone common.pt file; + # newer sharded checkpoints store it as a ShardedObject inside the torch_dist checkpoint. + legacy_common_path = os.path.join(modelopt_checkpoint_name, "common.py") + if os.path.exists(legacy_common_path): + common_modelopt_state = safe_load_from_bytes(legacy_common_path) + else: + common_modelopt_state = load_common_state_dict(modelopt_checkpoint_name) + + modelopt_load_version = common_modelopt_state["modelopt_version"] + + logger.info(f"nvidia-modelopt ckpt/inst version: {modelopt_load_version}/{modelopt.__version__}") + + model[0] = mto.restore_from_modelopt_state(model[0], common_modelopt_state) + + _load_extra_state_from_sharded_checkpoint(model[0], checkpoint_name, prefix, metadata=metadata) diff --git a/megatron/post_training/checkpointing.py b/megatron/post_training/checkpointing.py index 970792256d0..4449a36221b 100644 --- a/megatron/post_training/checkpointing.py +++ b/megatron/post_training/checkpointing.py @@ -1,20 +1,15 @@ # Copyright (c) 2024, NVIDIA CORPORATION. All rights reserved. import logging -import os from pathlib import Path -from typing import Any, Optional, Tuple, Union +from typing import Optional, Tuple, Union -import modelopt import modelopt.torch.opt as mto -from modelopt.torch.utils import safe_load -import torch import torch.nn as nn +from modelopt.torch.opt.plugins import restore_sharded_modelopt_state from megatron.core import dist_checkpointing -from megatron.core.dist_checkpointing.strategies.torch import TorchDistLoadShardedStrategy -from megatron.core.dist_checkpointing.validation import StrictHandling -from megatron.core.utils import unwrap_model +from megatron.core.utils import get_torch_version, is_torch_min_version, unwrap_model from megatron.training import get_args from megatron.training.checkpointing import _load_base_checkpoint, load_checkpoint from megatron.training.utils import print_rank_0 @@ -23,7 +18,6 @@ logger = logging.getLogger(__name__) NEMO_WEIGHT_DIR_NAMES = {"model_weights": "model.", "weights": "module."} -COMMON_STATE_FNAME = "common.pt" def has_modelopt_state(checkpoint_path: str) -> bool: @@ -200,90 +194,4 @@ def _remove_prefix_state_dict_pre_hook( unwrapped_model[0].load_state_dict(model_state_dict, strict=False) print_distributed_quant_summary(unwrapped_model[0]) else: - _ = load_checkpoint(model, optimizer, opt_param_scheduler, strict=strict, load_arg=load_arg) - - -def restore_sharded_modelopt_state( - model: list[torch.nn.Module], - checkpoint_name: str | Path, - prefix: str = "", - metadata: dict[str, Any] | None = None, -) -> None: - """Restore modelopt_state from the sharded state_dict format. - - Args: - model: the model to restore the modelopt optimization - checkpoint_name: the checkpoint folder path - prefix: the prefix to add to the modelopt_state keys ("model." for NeMo) - metadata: the metadata for distributed checkpointing - - Note: - The metadata includes several breaking changes. For example, `singleton_local_shards` - is set to `True` (was not set before) in megatron-core-0.15.0. This flag affects the - sharded state_dict format and must be consistent between saving and loading. - """ - if len(model) > 1: - raise ValueError("sharded_modelopt_state does not support virtual pipeline parallel!") - - modelopt_checkpoint_name = f"{checkpoint_name}/modelopt_state" - - # Early return if the model already has a modelopt_state or the checkpoint does not exist. - if not os.path.exists(modelopt_checkpoint_name) or mto.ModeloptStateManager.is_converted( - model[0] - ): - return - - # Loading the common modelopt_state (replicated on all ranks). - # Detect format: legacy checkpoints store common state in a standalone common.pt file; - # newer sharded checkpoints store it as a ShardedObject inside the torch_dist checkpoint. - legacy_common_path = os.path.join(modelopt_checkpoint_name, COMMON_STATE_FNAME) - if os.path.exists(legacy_common_path): - common_modelopt_state = safe_load(legacy_common_path) - else: - common_modelopt_state = dist_checkpointing.load_common_state_dict(modelopt_checkpoint_name) - - modelopt_load_version = common_modelopt_state["modelopt_version"] - - print_rank_0(f"nvidia-modelopt ckpt/inst version: {modelopt_load_version}/{modelopt.__version__}") - - model[0] = mto.restore_from_modelopt_state(model[0], common_modelopt_state) - - _load_extra_state_from_sharded_checkpoint(model[0], checkpoint_name, prefix, metadata=metadata) - - -def _load_extra_state_from_sharded_checkpoint( - model: torch.nn.Module, - checkpoint_name: str | Path, - prefix: str, - metadata: dict[str, Any] | None = None, -) -> None: - """Load extra state from sharded checkpoint. - - Note: since extra_state is a subset of full the sharded_state_dict, we use - strict=StrictHandling.LOG_UNEXPECTED instead of LOG_ALL. - - Args: - model: the model to load extra state into - checkpoint_name: the checkpoint folder path - prefix: the prefix to add to the modelopt_state keys - metadata: the metadata for distributed checkpointing - - Note: - The metadata includes several breaking changes. For example, `singleton_local_shards` - is set to `True` (was not set before) in megatron-core-0.15.0. This flag affects the - sharded state_dict format and must be consistent between saving and loading. - """ - sharded_state_dict = model.sharded_state_dict(prefix=prefix) - extra_sharded_state_dict = {k: v for k, v in sharded_state_dict.items() if "_extra_state" in k} - extra_state_dict = dist_checkpointing.load( - extra_sharded_state_dict, - checkpoint_name, - TorchDistLoadShardedStrategy(), - strict=StrictHandling.LOG_UNEXPECTED, - ) - extra_state_dict_no_prefix = {} - - for k, v in extra_state_dict.items(): - if k.startswith(prefix): - extra_state_dict_no_prefix[k[len(prefix) :]] = v - model.load_state_dict(extra_state_dict_no_prefix, strict=False) + _ = load_checkpoint(model, optimizer, opt_param_scheduler, strict=strict, load_arg=load_arg) \ No newline at end of file diff --git a/megatron/training/checkpointing.py b/megatron/training/checkpointing.py index f01df7b97e2..ad9de5b15b2 100644 --- a/megatron/training/checkpointing.py +++ b/megatron/training/checkpointing.py @@ -3,7 +3,6 @@ """Input/output checkpointing.""" import contextlib -import copy import inspect import multiprocessing import os @@ -66,12 +65,9 @@ # [ModelOpt]: Import try: - import modelopt.torch.utils.distributed as dist - import modelopt.torch.opt as mto + from modelopt.torch.opt.plugins import save_modelopt_state, save_sharded_modelopt_state - from modelopt.torch.opt.plugins import save_modelopt_state from megatron.post_training.utils import print_distributed_quant_summary - has_nvidia_modelopt = True except Exception: has_nvidia_modelopt = False @@ -2290,56 +2286,3 @@ def load_biencoder_checkpoint(model, only_query_model=False, print(' successfully loaded {}'.format(checkpoint_name)) return model - - -def remove_per_module_state( - modelopt_state: dict[str, Any], -) -> None: - """Remove metadata from the modelopt_state. - - The metadata of the modelopt_state contains keys which may change with different pipeline - and expert parallelism. As a result, the metadata must be stored as several ShardedObject with - global and local layer offset mapping. - - Args: - modelopt_state: the state_dict that contains all algorithms that have been applied - to the given model. - """ - if "modelopt_state_dict" not in modelopt_state: - return - - for mode, config in modelopt_state["modelopt_state_dict"]: - metadata = config.get("metadata", None) - if metadata is not None: - _ = metadata.pop("quantizer_state", None) - _ = metadata.pop("subnet_config", None) - _ = metadata.pop("real_quantizer_state", None) - _ = metadata.pop("q_tensor_state", None) - else: - config["metadata"] = {} - - -def save_sharded_modelopt_state( - model: list[torch.nn.Module], - checkpoint_name: str | Path, - sharded_strategy: tuple[str, int] | None = None, - prefix: str = "", -) -> None: - """Save modelopt_state in the sharded state_dict format. - - Args: - model: the model to restore the modelopt optimization - checkpoint_name: the checkpoint folder path - sharded_strategy: configures sharded tensors saving behavior and backend - prefix: the prefix to add to the modelopt_state keys ("model." for NeMo) - """ - if not mto.ModeloptStateManager.is_converted(model[0]): - return - if len(model) > 1: - raise ValueError("sharded_modelopt_state does not support virtual pipeline parallel!") - modelopt_checkpoint_name = f"{checkpoint_name}/modelopt_state" - if dist.is_master(): - os.makedirs(modelopt_checkpoint_name, exist_ok=True) - modelopt_state = copy.deepcopy(mto.modelopt_state(model[0])) - remove_per_module_state(modelopt_state) - dist_checkpointing.save(modelopt_state, modelopt_checkpoint_name, sharded_strategy) From f52c7bb90ac9bd3589c6be945c661b9ffd3dfc7e Mon Sep 17 00:00:00 2001 From: dimapihtar Date: Mon, 6 Jul 2026 03:10:02 -0700 Subject: [PATCH 4/7] fix imports Signed-off-by: dimapihtar --- megatron/post_training/checkpointing.py | 2 +- megatron/training/checkpointing.py | 6 ++++-- 2 files changed, 5 insertions(+), 3 deletions(-) diff --git a/megatron/post_training/checkpointing.py b/megatron/post_training/checkpointing.py index 4449a36221b..c3326cc24aa 100644 --- a/megatron/post_training/checkpointing.py +++ b/megatron/post_training/checkpointing.py @@ -6,9 +6,9 @@ import modelopt.torch.opt as mto import torch.nn as nn -from modelopt.torch.opt.plugins import restore_sharded_modelopt_state from megatron.core import dist_checkpointing +from megatron.core.dist_checkpointing.strategies.modelopt import restore_sharded_modelopt_state from megatron.core.utils import get_torch_version, is_torch_min_version, unwrap_model from megatron.training import get_args from megatron.training.checkpointing import _load_base_checkpoint, load_checkpoint diff --git a/megatron/training/checkpointing.py b/megatron/training/checkpointing.py index ad9de5b15b2..e9a7e2c799d 100644 --- a/megatron/training/checkpointing.py +++ b/megatron/training/checkpointing.py @@ -30,6 +30,10 @@ FullyParallelLoadStrategyWrapper, FullyParallelSaveStrategyWrapper, ) +from megatron.core.dist_checkpointing.strategies.modelopt import ( + save_modelopt_state, + save_sharded_modelopt_state, +) from megatron.core.dist_checkpointing.strategies.torch import ( TorchDistLoadShardedStrategy, TorchDistSaveShardedStrategy, @@ -65,8 +69,6 @@ # [ModelOpt]: Import try: - from modelopt.torch.opt.plugins import save_modelopt_state, save_sharded_modelopt_state - from megatron.post_training.utils import print_distributed_quant_summary has_nvidia_modelopt = True except Exception: From f9a7a8edd512ee4e795a8858b3d709821bb03bbc Mon Sep 17 00:00:00 2001 From: dimapihtar Date: Mon, 6 Jul 2026 03:12:38 -0700 Subject: [PATCH 5/7] fix code style Signed-off-by: dimapihtar --- megatron/core/dist_checkpointing/strategies/modelopt.py | 8 ++++---- 1 file changed, 4 insertions(+), 4 deletions(-) diff --git a/megatron/core/dist_checkpointing/strategies/modelopt.py b/megatron/core/dist_checkpointing/strategies/modelopt.py index da0be0d3468..2644fafb37a 100644 --- a/megatron/core/dist_checkpointing/strategies/modelopt.py +++ b/megatron/core/dist_checkpointing/strategies/modelopt.py @@ -29,9 +29,7 @@ logger = logging.getLogger(__name__) -def remove_per_module_state( - modelopt_state: dict[str, Any], -) -> None: +def remove_per_module_state(modelopt_state: dict[str, Any]) -> None: """Remove metadata from the modelopt_state. The metadata of the modelopt_state contains keys which may change with different pipeline @@ -180,7 +178,9 @@ def restore_sharded_modelopt_state( modelopt_load_version = common_modelopt_state["modelopt_version"] - logger.info(f"nvidia-modelopt ckpt/inst version: {modelopt_load_version}/{modelopt.__version__}") + logger.info( + f"nvidia-modelopt ckpt/inst version: {modelopt_load_version}/{modelopt.__version__}" + ) model[0] = mto.restore_from_modelopt_state(model[0], common_modelopt_state) From de51811e5a4b29b7206cfcedeb070d4beee57de4 Mon Sep 17 00:00:00 2001 From: dimapihtar Date: Mon, 6 Jul 2026 03:23:27 -0700 Subject: [PATCH 6/7] fix imports Signed-off-by: dimapihtar --- .../core/dist_checkpointing/strategies/modelopt.py | 10 +++++----- 1 file changed, 5 insertions(+), 5 deletions(-) diff --git a/megatron/core/dist_checkpointing/strategies/modelopt.py b/megatron/core/dist_checkpointing/strategies/modelopt.py index 2644fafb37a..a76c844a6e3 100644 --- a/megatron/core/dist_checkpointing/strategies/modelopt.py +++ b/megatron/core/dist_checkpointing/strategies/modelopt.py @@ -10,17 +10,17 @@ import torch -from .torch import TorchDistLoadShardedStrategy -from ..serialization import load, load_common_state_dict, save -from ..validation import StrictHandling - from megatron.core import mpu from megatron.core.safe_globals import safe_load_from_bytes +from ..serialization import load, load_common_state_dict, save +from ..validation import StrictHandling +from .torch import TorchDistLoadShardedStrategy + try: import modelopt - import modelopt.torch.utils.distributed as dist import modelopt.torch.opt as mto + import modelopt.torch.utils.distributed as dist has_nvidia_modelopt = True except (ModuleNotFoundError, ImportError): From d720d2e53c26aa2f2565418c1ccc6763db36edfd Mon Sep 17 00:00:00 2001 From: Dmytro Pykhtar <37850217+dimapihtar@users.noreply.github.com> Date: Mon, 6 Jul 2026 19:11:28 +0300 Subject: [PATCH 7/7] minor change Co-authored-by: Keval Morabia <28916987+kevalmorabia97@users.noreply.github.com> --- megatron/core/dist_checkpointing/strategies/modelopt.py | 2 +- 1 file changed, 1 insertion(+), 1 deletion(-) diff --git a/megatron/core/dist_checkpointing/strategies/modelopt.py b/megatron/core/dist_checkpointing/strategies/modelopt.py index a76c844a6e3..c50bc22cfd0 100644 --- a/megatron/core/dist_checkpointing/strategies/modelopt.py +++ b/megatron/core/dist_checkpointing/strategies/modelopt.py @@ -23,7 +23,7 @@ import modelopt.torch.utils.distributed as dist has_nvidia_modelopt = True -except (ModuleNotFoundError, ImportError): +except ImportError: has_nvidia_modelopt = False logger = logging.getLogger(__name__)