From 31d9d06b22c279f231dee09c1875ce6aa6a6b973 Mon Sep 17 00:00:00 2001 From: chenhany Date: Tue, 7 Apr 2026 14:42:45 -0700 Subject: [PATCH 01/11] fix: handle list-format quant_cfg from ModelOpt PR #1094 ModelOpt PR #1094 changed quant_cfg from a flat dict keyed by quantizer name to an ordered list of QuantizerCfgEntry dicts with an explicit 'quantizer_name' field. quantize.py was directly indexing the dict (e.g. quant_cfg['*medusa_heads**'] = ...) which crashes with a TypeError when quant_cfg is a list. Add _set_quant_cfg_entry() and _get_quant_cfg_entry() helpers that write to and read from either format, and update all call sites in get_modelopt_torch_quantization_config(), get_first_layers_disabled_config(), and get_last_layers_disabled_config() to use them. Signed-off-by: chenhany --- examples/post_training/modelopt/quantize.py | 66 +++++++++++++++------ 1 file changed, 48 insertions(+), 18 deletions(-) diff --git a/examples/post_training/modelopt/quantize.py b/examples/post_training/modelopt/quantize.py index dc4947038e5..e76346f0f7c 100644 --- a/examples/post_training/modelopt/quantize.py +++ b/examples/post_training/modelopt/quantize.py @@ -181,6 +181,35 @@ def _is_last_layers(name: str, num_layers: int = 1, num_layers_to_disable: int = return layer_idx >= num_layers - num_layers_to_disable +def _set_quant_cfg_entry(quant_cfg, quantizer_name, entry_cfg): + """Set a quantizer config entry by name, handling both dict and list formats. + + ModelOpt PR #1094 changed quant_cfg from a flat dict keyed by quantizer name + to an ordered list of QuantizerCfgEntry dicts with an explicit "quantizer_name" + field. This helper writes to either format so callers don't need to branch. + """ + if isinstance(quant_cfg, list): + for entry in quant_cfg: + if entry.get("quantizer_name") == quantizer_name: + entry.update(entry_cfg) + return + new_entry = {"quantizer_name": quantizer_name} + new_entry.update(entry_cfg) + quant_cfg.append(new_entry) + else: + quant_cfg[quantizer_name] = entry_cfg + + +def _get_quant_cfg_entry(quant_cfg, quantizer_name): + """Get a quantizer config entry by name, handling both dict and list formats.""" + if isinstance(quant_cfg, list): + for entry in quant_cfg: + if entry.get("quantizer_name") == quantizer_name: + return entry + return None + return quant_cfg.get(quantizer_name) + + def get_first_layers_disabled_config(config, num_layers: int = 1, num_layers_to_disable: int = 1): """Get a config for `mtq.quantize` with first & last `num_layers_to_disable` layers disabled. @@ -188,13 +217,13 @@ def get_first_layers_disabled_config(config, num_layers: int = 1, num_layers_to_ """ config = copy.deepcopy(config) quant_cfg = config.get("quant_cfg", {}) - quant_cfg.update( - { - functools.partial( - _is_first_layers, num_layers=num_layers, num_layers_to_disable=num_layers_to_disable - ): {"enable": False} - } + predicate = functools.partial( + _is_first_layers, num_layers=num_layers, num_layers_to_disable=num_layers_to_disable ) + if isinstance(quant_cfg, list): + quant_cfg.append({"quantizer_name": predicate, "enable": False}) + else: + quant_cfg.update({predicate: {"enable": False}}) config["quant_cfg"] = quant_cfg return config @@ -206,13 +235,13 @@ def get_last_layers_disabled_config(config, num_layers: int = 1, num_layers_to_d """ config = copy.deepcopy(config) quant_cfg = config.get("quant_cfg", {}) - quant_cfg.update( - { - functools.partial( - _is_last_layers, num_layers=num_layers, num_layers_to_disable=num_layers_to_disable - ): {"enable": False} - } + predicate = functools.partial( + _is_last_layers, num_layers=num_layers, num_layers_to_disable=num_layers_to_disable ) + if isinstance(quant_cfg, list): + quant_cfg.append({"quantizer_name": predicate, "enable": False}) + else: + quant_cfg.update({predicate: {"enable": False}}) config["quant_cfg"] = quant_cfg return config @@ -233,19 +262,20 @@ def get_modelopt_torch_quantization_config(): } if args.export_quant_cfg == "FP8_DEFAULT_CFG": # Enable Medusa heads and kv-cache quantization - mtq_config["quant_cfg"]["*medusa_heads**"] = fp8_config + _set_quant_cfg_entry(mtq_config["quant_cfg"], "*medusa_heads**", fp8_config) if "FP4" in args.export_quant_cfg: # Enable Medusa heads and kv-cache quantization - mtq_config["quant_cfg"]["*medusa_heads**"] = fp4_config + _set_quant_cfg_entry(mtq_config["quant_cfg"], "*medusa_heads**", fp4_config) if "AWQ" in args.export_quant_cfg: - weight_quantizer = mtq_config["quant_cfg"]["*weight_quantizer"] # type: ignore + weight_quantizer = _get_quant_cfg_entry(mtq_config["quant_cfg"], "*weight_quantizer") if isinstance(weight_quantizer, list): weight_quantizer = weight_quantizer[0] - weight_quantizer["block_sizes"][-1] = 128 + if weight_quantizer is not None: + weight_quantizer["block_sizes"][-1] = 128 # Customization if args.disable_qkv_quant: - mtq_config["quant_cfg"]["*self_attention*"] = {"enable": False} + _set_quant_cfg_entry(mtq_config["quant_cfg"], "*self_attention*", {"enable": False}) # KV Cache Quantization enable_quant_kv_cache = args.export_kv_cache_quant != "none" @@ -257,7 +287,7 @@ def get_modelopt_torch_quantization_config(): # Weight Only Quantization if args.weight_only: - mtq_config["quant_cfg"]["*input_quantizer"] = {"enable": False} + _set_quant_cfg_entry(mtq_config["quant_cfg"], "*input_quantizer", {"enable": False}) if args.num_first_layers_to_skip_quant is not None: mtq_config = get_first_layers_disabled_config( mtq_config, From 1bea6cc39aad3ed2a045651edd381699be138eab Mon Sep 17 00:00:00 2001 From: chenhany Date: Wed, 8 Apr 2026 05:37:26 -0700 Subject: [PATCH 02/11] fix: wrap fp8/fp4 quantizer cfgs in 'cfg' key for new ModelOpt format ModelOpt PR #1094 changed QuantizerCfgEntry to require quantizer attributes under a 'cfg' sub-key rather than at the top level. Update fp8_config and fp4_config in quantize.py accordingly. Co-Authored-By: Claude Sonnet 4.6 --- examples/post_training/modelopt/quantize.py | 10 ++++++---- 1 file changed, 6 insertions(+), 4 deletions(-) diff --git a/examples/post_training/modelopt/quantize.py b/examples/post_training/modelopt/quantize.py index e76346f0f7c..b1b81a5c4da 100644 --- a/examples/post_training/modelopt/quantize.py +++ b/examples/post_training/modelopt/quantize.py @@ -253,12 +253,14 @@ def get_modelopt_torch_quantization_config(): raise ValueError(f"Unsupported quantization config {args.export_quant_cfg}.") mtq_config = QUANT_CFG_CHOICES[args.export_quant_cfg] - fp8_config = {"enable": True, "num_bits": (4, 3), "axis": None} + fp8_config = {"enable": True, "cfg": {"num_bits": (4, 3), "axis": None}} fp4_config = { - "num_bits": (2, 1), - "block_sizes": {-1: 16, "type": "dynamic", "scale_bits": (4, 3)}, - "axis": None, "enable": True, + "cfg": { + "num_bits": (2, 1), + "block_sizes": {-1: 16, "type": "dynamic", "scale_bits": (4, 3)}, + "axis": None, + }, } if args.export_quant_cfg == "FP8_DEFAULT_CFG": # Enable Medusa heads and kv-cache quantization From c11adb07feb474167c2421568c6f18da793cb35b Mon Sep 17 00:00:00 2001 From: Chenhan Yu Date: Wed, 15 Apr 2026 08:38:29 -0700 Subject: [PATCH 03/11] fix: update modelopt examples for new args init API initialize_megatron() and pretrain() no longer accept extra_args_provider and args_defaults after commit 97aca2f88. Use parse_and_validate_args() from megatron.training.arguments to parse args before calling initialize_megatron()/pretrain(). Affected files: - quantize.py, validate.py, prune.py, mmlu.py, export.py - convert_model.py, offline_feature_extract.py, generate.py - finetune.py (uses pretrain() instead of initialize_megatron()) Signed-off-by: Chenhan Yu --- examples/post_training/modelopt/convert_model.py | 9 ++++----- examples/post_training/modelopt/export.py | 9 ++++----- examples/post_training/modelopt/finetune.py | 8 ++++++-- examples/post_training/modelopt/generate.py | 9 ++++----- examples/post_training/modelopt/mmlu.py | 9 ++++----- .../post_training/modelopt/offline_feature_extract.py | 8 +++----- examples/post_training/modelopt/prune.py | 9 ++++----- examples/post_training/modelopt/quantize.py | 9 ++++----- examples/post_training/modelopt/validate.py | 9 ++++----- 9 files changed, 37 insertions(+), 42 deletions(-) diff --git a/examples/post_training/modelopt/convert_model.py b/examples/post_training/modelopt/convert_model.py index eaec9789e1e..205e581b596 100644 --- a/examples/post_training/modelopt/convert_model.py +++ b/examples/post_training/modelopt/convert_model.py @@ -27,6 +27,7 @@ from megatron.training import get_args from megatron.training.checkpointing import save_checkpoint from megatron.training.initialize import initialize_megatron +from megatron.training.arguments import parse_and_validate_args from megatron.training.utils import print_rank_0, unwrap_model from model_provider import model_provider @@ -102,14 +103,12 @@ def check_arguments(): if __name__ == "__main__": - initialize_megatron( - extra_args_provider=add_convert_args, - args_defaults={ + parse_and_validate_args(extra_args_provider=add_convert_args, args_defaults={ 'tokenizer_type': 'HuggingFaceTokenizer', 'no_load_rng': True, 'no_load_optim': True, - }, - ) + }) + initialize_megatron() check_arguments() args = get_args() diff --git a/examples/post_training/modelopt/export.py b/examples/post_training/modelopt/export.py index 5e3b2a1716e..9b9b08e8e11 100755 --- a/examples/post_training/modelopt/export.py +++ b/examples/post_training/modelopt/export.py @@ -18,6 +18,7 @@ from megatron.post_training.model_builder import modelopt_gpt_mamba_builder from megatron.training import get_args, get_model from megatron.training.initialize import initialize_megatron +from megatron.training.arguments import parse_and_validate_args from megatron.training.utils import unwrap_model from model_provider import model_provider @@ -49,14 +50,12 @@ def add_modelopt_export_args(parser): if __name__ == "__main__": - initialize_megatron( - extra_args_provider=add_modelopt_export_args, - args_defaults={ + parse_and_validate_args(extra_args_provider=add_modelopt_export_args, args_defaults={ 'tokenizer_type': 'HuggingFaceTokenizer', 'no_load_rng': True, 'no_load_optim': True, - }, - ) + }) + initialize_megatron() args = get_args() diff --git a/examples/post_training/modelopt/finetune.py b/examples/post_training/modelopt/finetune.py index f7f7c24f970..e95e2562686 100755 --- a/examples/post_training/modelopt/finetune.py +++ b/examples/post_training/modelopt/finetune.py @@ -484,12 +484,16 @@ def forward_step(data_iterator, model: GPTModel): if __name__ == "__main__": + from megatron.training.arguments import parse_and_validate_args + + parse_and_validate_args( + extra_args_provider=add_finetune_args, + args_defaults={"tokenizer_type": "HuggingFaceTokenizer"}, + ) pretrain( train_valid_test_sft_datasets_provider, partial(model_provider, modelopt_gpt_mamba_builder), ModelType.encoder_or_decoder, forward_step, - extra_args_provider=add_finetune_args, - args_defaults={"tokenizer_type": "HuggingFaceTokenizer"}, non_loss_data_func=non_loss_data_func, ) diff --git a/examples/post_training/modelopt/generate.py b/examples/post_training/modelopt/generate.py index 3d3f6571b34..2ed430ed7a1 100644 --- a/examples/post_training/modelopt/generate.py +++ b/examples/post_training/modelopt/generate.py @@ -17,6 +17,7 @@ from megatron.post_training.model_builder import modelopt_gpt_mamba_builder from megatron.post_training.utils import report_current_memory_info, to_empty_if_meta from megatron.training import get_args, get_model, initialize_megatron +from megatron.training.arguments import parse_and_validate_args from utils import get_hf_tokenizer from megatron.training.utils import print_rank_0, unwrap_model from model_provider import model_provider @@ -73,14 +74,12 @@ def get_conversations(example): if __name__ == "__main__": - initialize_megatron( - extra_args_provider=add_generate_args, - args_defaults={ + parse_and_validate_args(extra_args_provider=add_generate_args, args_defaults={ 'tokenizer_type': 'HuggingFaceTokenizer', 'no_load_rng': True, 'no_load_optim': True, - }, - ) + }) + initialize_megatron() check_arguments() diff --git a/examples/post_training/modelopt/mmlu.py b/examples/post_training/modelopt/mmlu.py index 5aa5d1c24c7..212e4a42ab1 100644 --- a/examples/post_training/modelopt/mmlu.py +++ b/examples/post_training/modelopt/mmlu.py @@ -20,6 +20,7 @@ from megatron.post_training.model_builder import modelopt_gpt_mamba_builder from megatron.post_training.utils import report_current_memory_info from megatron.training import get_args, get_model, initialize_megatron +from megatron.training.arguments import parse_and_validate_args from utils import get_hf_tokenizer from megatron.training.utils import print_rank_0, unwrap_model import modelopt.torch.quantization as mtq @@ -133,14 +134,12 @@ def generate_prompt(test_example, dev_examples, few_shots=0, no_subject_prompt=F if __name__ == "__main__": - initialize_megatron( - extra_args_provider=add_mmlu_args, - args_defaults={ + parse_and_validate_args(extra_args_provider=add_mmlu_args, args_defaults={ 'tokenizer_type': 'HuggingFaceTokenizer', 'no_load_rng': True, 'no_load_optim': True, - }, - ) + }) + initialize_megatron() args = get_args() cache = Cache(args.cache_dir) diff --git a/examples/post_training/modelopt/offline_feature_extract.py b/examples/post_training/modelopt/offline_feature_extract.py index 80207faf2b2..b5f386fdd23 100644 --- a/examples/post_training/modelopt/offline_feature_extract.py +++ b/examples/post_training/modelopt/offline_feature_extract.py @@ -42,14 +42,12 @@ def extract_feature(dataset, model, output_dir, idx_start, idx_end): torch.distributed.barrier() if __name__ == "__main__": - initialize_megatron( - extra_args_provider=add_extract_args, - args_defaults={ + parse_and_validate_args(extra_args_provider=add_extract_args, args_defaults={ 'tokenizer_type': 'HuggingFaceTokenizer', 'no_load_rng': True, 'no_load_optim': True, - }, - ) + }) + initialize_megatron() args = get_args() tokenizer = get_tokenizer() diff --git a/examples/post_training/modelopt/prune.py b/examples/post_training/modelopt/prune.py index 56bbffa0cd0..0f0ead6610a 100644 --- a/examples/post_training/modelopt/prune.py +++ b/examples/post_training/modelopt/prune.py @@ -33,6 +33,7 @@ report_current_memory_info, ) from megatron.training import get_args, get_model, initialize_megatron +from megatron.training.arguments import parse_and_validate_args from utils import get_hf_tokenizer from megatron.training.checkpointing import save_checkpoint from megatron.training.utils import print_rank_0, unwrap_model @@ -149,14 +150,12 @@ def get_params(model): if __name__ == "__main__": - initialize_megatron( - extra_args_provider=add_prune_args, - args_defaults={ + parse_and_validate_args(extra_args_provider=add_prune_args, args_defaults={ "tokenizer_type": "HuggingFaceTokenizer", "no_load_rng": True, "no_load_optim": True, - }, - ) + }) + initialize_megatron() args = get_args() check_arguments(args) diff --git a/examples/post_training/modelopt/quantize.py b/examples/post_training/modelopt/quantize.py index b1b81a5c4da..8b4310dd6c9 100644 --- a/examples/post_training/modelopt/quantize.py +++ b/examples/post_training/modelopt/quantize.py @@ -45,6 +45,7 @@ report_current_memory_info, ) from megatron.training import get_args, get_model, initialize_megatron +from megatron.training.arguments import parse_and_validate_args from utils import get_hf_tokenizer from megatron.training.checkpointing import save_checkpoint from megatron.training.utils import print_rank_0, unwrap_model @@ -378,14 +379,12 @@ def get_calib_dataloader( if __name__ == "__main__": - initialize_megatron( - extra_args_provider=add_text_generate_ptq_args, - args_defaults={ + parse_and_validate_args(extra_args_provider=add_text_generate_ptq_args, args_defaults={ "tokenizer_type": "HuggingFaceTokenizer", "no_load_rng": True, "no_load_optim": True, - }, - ) + }) + initialize_megatron() check_arguments() diff --git a/examples/post_training/modelopt/validate.py b/examples/post_training/modelopt/validate.py index 8b8f1ffc9dd..10a633dc49c 100644 --- a/examples/post_training/modelopt/validate.py +++ b/examples/post_training/modelopt/validate.py @@ -17,6 +17,7 @@ from megatron.post_training.model_builder import modelopt_gpt_mamba_builder from megatron.post_training.utils import get_mtbench_chat_data from megatron.training import get_args, get_model, initialize_megatron +from megatron.training.arguments import parse_and_validate_args from utils import get_hf_tokenizer from megatron.training.utils import print_rank_0, unwrap_model from model_provider import model_provider @@ -89,14 +90,12 @@ def report_current_memory_info(): if __name__ == "__main__": - initialize_megatron( - extra_args_provider=add_ar_validation_args, - args_defaults={ + parse_and_validate_args(extra_args_provider=add_ar_validation_args, args_defaults={ 'tokenizer_type': 'HuggingFaceTokenizer', 'no_load_rng': True, 'no_load_optim': True, - }, - ) + }) + initialize_megatron() check_arguments() From bda6752f4203085192530c35f9c3a917a9896f79 Mon Sep 17 00:00:00 2001 From: Chenhan Yu Date: Mon, 20 Apr 2026 16:04:36 -0700 Subject: [PATCH 04/11] fix: fallback for privatized get_write_results_queue in nvidia-resiliency-ext Some container versions renamed get_write_results_queue to _get_write_results_queue. Add fallback import to handle both. --- megatron/core/dist_checkpointing/strategies/torch.py | 9 ++++++++- 1 file changed, 8 insertions(+), 1 deletion(-) diff --git a/megatron/core/dist_checkpointing/strategies/torch.py b/megatron/core/dist_checkpointing/strategies/torch.py index 3e467c1f9dd..8098f4f79c4 100644 --- a/megatron/core/dist_checkpointing/strategies/torch.py +++ b/megatron/core/dist_checkpointing/strategies/torch.py @@ -1040,8 +1040,15 @@ def get_async_strategy(async_strategy: str = "nvrx", module: str = None) -> tupl from nvidia_resiliency_ext.checkpointing.async_ckpt.filesystem_async import ( FileSystemWriterAsync, _results_queue, - get_write_results_queue, ) + try: + from nvidia_resiliency_ext.checkpointing.async_ckpt.filesystem_async import ( + get_write_results_queue, + ) + except ImportError: + from nvidia_resiliency_ext.checkpointing.async_ckpt.filesystem_async import ( + _get_write_results_queue as get_write_results_queue, + ) from nvidia_resiliency_ext.checkpointing.async_ckpt.state_dict_saver import ( CheckpointMetadataCache, save_state_dict_async_finalize, From 084cd5ec08b4defa09dd356430db6a1c14f93cf1 Mon Sep 17 00:00:00 2001 From: Chenhan Yu Date: Mon, 20 Apr 2026 20:21:07 -0700 Subject: [PATCH 05/11] fix: add DISABLE_NVRX env var to bypass nvidia-resiliency-ext Set DISABLE_NVRX=1 to force HAVE_NVRX=False and use mcore's own checkpoint implementation. Workaround for containers where nvrx is incompatible with the installed PyTorch version. --- megatron/core/dist_checkpointing/strategies/torch.py | 2 ++ 1 file changed, 2 insertions(+) diff --git a/megatron/core/dist_checkpointing/strategies/torch.py b/megatron/core/dist_checkpointing/strategies/torch.py index 8098f4f79c4..b9a4a6a3691 100644 --- a/megatron/core/dist_checkpointing/strategies/torch.py +++ b/megatron/core/dist_checkpointing/strategies/torch.py @@ -53,6 +53,8 @@ from .checkpointable import CheckpointableShardedTensor, LocalShardsContainer try: + if os.environ.get("DISABLE_NVRX"): + raise ImportError("NVRX disabled by DISABLE_NVRX env var") from nvidia_resiliency_ext.checkpointing.async_ckpt.core import AsyncRequest as NVRxAsyncRequest from nvidia_resiliency_ext.checkpointing.async_ckpt.state_dict_saver import ( CheckpointMetadataCache, From fc6a931680ea1a352bcfbe0cacba4eff9559b1bd Mon Sep 17 00:00:00 2001 From: Chenhan Yu Date: Mon, 20 Apr 2026 20:23:53 -0700 Subject: [PATCH 06/11] fix: set DISABLE_NVRX=1 in all modelopt example scripts nvidia-resiliency-ext's FileSystemWriterAsync is incompatible with some container PyTorch versions. DISABLE_NVRX forces mcore fallback. --- examples/post_training/modelopt/convert.sh | 3 +++ examples/post_training/modelopt/eagle3.sh | 3 +++ examples/post_training/modelopt/export.sh | 3 +++ examples/post_training/modelopt/finetune.sh | 3 +++ examples/post_training/modelopt/generate.sh | 3 +++ examples/post_training/modelopt/generation_server.sh | 3 +++ examples/post_training/modelopt/mmlu.sh | 3 +++ examples/post_training/modelopt/offline_feature_extract.sh | 3 +++ examples/post_training/modelopt/prune.sh | 3 +++ examples/post_training/modelopt/quantize.sh | 3 +++ examples/post_training/modelopt/train.sh | 3 +++ examples/post_training/modelopt/validate.sh | 3 +++ 12 files changed, 36 insertions(+) diff --git a/examples/post_training/modelopt/convert.sh b/examples/post_training/modelopt/convert.sh index 9dbe27da541..8bb51ae6303 100755 --- a/examples/post_training/modelopt/convert.sh +++ b/examples/post_training/modelopt/convert.sh @@ -1,5 +1,8 @@ #!/bin/bash +# Disable nvidia-resiliency-ext (incompatible with some container PyTorch versions) +export DISABLE_NVRX=1 + SCRIPT_DIR="$(dirname "$(readlink -f "$0")")" # Common arguments and base model specific arguments diff --git a/examples/post_training/modelopt/eagle3.sh b/examples/post_training/modelopt/eagle3.sh index c560b1d3a93..e1763fdb845 100755 --- a/examples/post_training/modelopt/eagle3.sh +++ b/examples/post_training/modelopt/eagle3.sh @@ -1,5 +1,8 @@ #!/bin/bash +# Disable nvidia-resiliency-ext (incompatible with some container PyTorch versions) +export DISABLE_NVRX=1 + SCRIPT_DIR="$(dirname "$(readlink -f "$0")")" # Common arguments and base model specific arguments diff --git a/examples/post_training/modelopt/export.sh b/examples/post_training/modelopt/export.sh index d0f0b97f2a0..da3109676c2 100755 --- a/examples/post_training/modelopt/export.sh +++ b/examples/post_training/modelopt/export.sh @@ -1,5 +1,8 @@ #!/bin/bash +# Disable nvidia-resiliency-ext (incompatible with some container PyTorch versions) +export DISABLE_NVRX=1 + SCRIPT_DIR="$(dirname "$(readlink -f "$0")")" # Common arguments and base model specific arguments diff --git a/examples/post_training/modelopt/finetune.sh b/examples/post_training/modelopt/finetune.sh index 1affc3c04da..1bf62fff4e3 100755 --- a/examples/post_training/modelopt/finetune.sh +++ b/examples/post_training/modelopt/finetune.sh @@ -1,5 +1,8 @@ #!/bin/bash +# Disable nvidia-resiliency-ext (incompatible with some container PyTorch versions) +export DISABLE_NVRX=1 + SCRIPT_DIR="$(dirname "$(readlink -f "$0")")" # Common arguments and base model specific arguments diff --git a/examples/post_training/modelopt/generate.sh b/examples/post_training/modelopt/generate.sh index 58578bb3d02..fa0d0630b15 100755 --- a/examples/post_training/modelopt/generate.sh +++ b/examples/post_training/modelopt/generate.sh @@ -1,5 +1,8 @@ #!/bin/bash +# Disable nvidia-resiliency-ext (incompatible with some container PyTorch versions) +export DISABLE_NVRX=1 + SCRIPT_DIR="$(dirname "$(readlink -f "$0")")" # Common arguments and base model specific arguments diff --git a/examples/post_training/modelopt/generation_server.sh b/examples/post_training/modelopt/generation_server.sh index 9acd61f3d04..363cda6caca 100755 --- a/examples/post_training/modelopt/generation_server.sh +++ b/examples/post_training/modelopt/generation_server.sh @@ -1,5 +1,8 @@ #!/bin/bash +# Disable nvidia-resiliency-ext (incompatible with some container PyTorch versions) +export DISABLE_NVRX=1 + SCRIPT_DIR="$(dirname "$(readlink -f "$0")")" # Common arguments and base model specific arguments diff --git a/examples/post_training/modelopt/mmlu.sh b/examples/post_training/modelopt/mmlu.sh index 36918b89df2..275f7256b40 100755 --- a/examples/post_training/modelopt/mmlu.sh +++ b/examples/post_training/modelopt/mmlu.sh @@ -1,5 +1,8 @@ #!/bin/bash +# Disable nvidia-resiliency-ext (incompatible with some container PyTorch versions) +export DISABLE_NVRX=1 + SCRIPT_DIR="$(dirname "$(readlink -f "$0")")" # Common arguments and base model specific arguments diff --git a/examples/post_training/modelopt/offline_feature_extract.sh b/examples/post_training/modelopt/offline_feature_extract.sh index 1ab7a0fcc85..2d458456a70 100755 --- a/examples/post_training/modelopt/offline_feature_extract.sh +++ b/examples/post_training/modelopt/offline_feature_extract.sh @@ -1,5 +1,8 @@ #!/bin/bash +# Disable nvidia-resiliency-ext (incompatible with some container PyTorch versions) +export DISABLE_NVRX=1 + SCRIPT_DIR="$(dirname "$(readlink -f "$0")")" # Common arguments and base model specific arguments diff --git a/examples/post_training/modelopt/prune.sh b/examples/post_training/modelopt/prune.sh index 33f3e615e96..91aaa953bae 100755 --- a/examples/post_training/modelopt/prune.sh +++ b/examples/post_training/modelopt/prune.sh @@ -1,5 +1,8 @@ #!/bin/bash +# Disable nvidia-resiliency-ext (incompatible with some container PyTorch versions) +export DISABLE_NVRX=1 + SCRIPT_DIR="$(dirname "$(readlink -f "$0")")" # Common arguments and base model specific arguments diff --git a/examples/post_training/modelopt/quantize.sh b/examples/post_training/modelopt/quantize.sh index 9119ff4ae76..7261dc422cd 100755 --- a/examples/post_training/modelopt/quantize.sh +++ b/examples/post_training/modelopt/quantize.sh @@ -1,5 +1,8 @@ #!/bin/bash +# Disable nvidia-resiliency-ext (incompatible with some container PyTorch versions) +export DISABLE_NVRX=1 + SCRIPT_DIR="$(dirname "$(readlink -f "$0")")" # Common arguments and base model specific arguments diff --git a/examples/post_training/modelopt/train.sh b/examples/post_training/modelopt/train.sh index 1ebb8bf3d76..730a4408108 100755 --- a/examples/post_training/modelopt/train.sh +++ b/examples/post_training/modelopt/train.sh @@ -1,5 +1,8 @@ #!/bin/bash +# Disable nvidia-resiliency-ext (incompatible with some container PyTorch versions) +export DISABLE_NVRX=1 + SCRIPT_DIR="$(dirname "$(readlink -f "$0")")" # Common arguments and base model specific arguments diff --git a/examples/post_training/modelopt/validate.sh b/examples/post_training/modelopt/validate.sh index 796231e508e..bf3d193ad2f 100755 --- a/examples/post_training/modelopt/validate.sh +++ b/examples/post_training/modelopt/validate.sh @@ -1,5 +1,8 @@ #!/bin/bash +# Disable nvidia-resiliency-ext (incompatible with some container PyTorch versions) +export DISABLE_NVRX=1 + SCRIPT_DIR="$(dirname "$(readlink -f "$0")")" # Common arguments and base model specific arguments From 7e08b1b95f9f0d138ce31edceb0545258e01aa69 Mon Sep 17 00:00:00 2001 From: Jennifer Chen Date: Mon, 27 Apr 2026 11:43:33 -0700 Subject: [PATCH 07/11] revert: remove DISABLE_NVRX=1 from modelopt example scripts Co-Authored-By: Claude Opus 4.7 (1M context) --- examples/post_training/modelopt/convert.sh | 3 --- examples/post_training/modelopt/eagle3.sh | 3 --- examples/post_training/modelopt/export.sh | 3 --- examples/post_training/modelopt/finetune.sh | 3 --- examples/post_training/modelopt/generate.sh | 3 --- examples/post_training/modelopt/generation_server.sh | 3 --- examples/post_training/modelopt/mmlu.sh | 3 --- examples/post_training/modelopt/offline_feature_extract.sh | 3 --- examples/post_training/modelopt/prune.sh | 3 --- examples/post_training/modelopt/quantize.sh | 3 --- examples/post_training/modelopt/train.sh | 3 --- examples/post_training/modelopt/validate.sh | 3 --- 12 files changed, 36 deletions(-) diff --git a/examples/post_training/modelopt/convert.sh b/examples/post_training/modelopt/convert.sh index 8bb51ae6303..9dbe27da541 100755 --- a/examples/post_training/modelopt/convert.sh +++ b/examples/post_training/modelopt/convert.sh @@ -1,8 +1,5 @@ #!/bin/bash -# Disable nvidia-resiliency-ext (incompatible with some container PyTorch versions) -export DISABLE_NVRX=1 - SCRIPT_DIR="$(dirname "$(readlink -f "$0")")" # Common arguments and base model specific arguments diff --git a/examples/post_training/modelopt/eagle3.sh b/examples/post_training/modelopt/eagle3.sh index e1763fdb845..c560b1d3a93 100755 --- a/examples/post_training/modelopt/eagle3.sh +++ b/examples/post_training/modelopt/eagle3.sh @@ -1,8 +1,5 @@ #!/bin/bash -# Disable nvidia-resiliency-ext (incompatible with some container PyTorch versions) -export DISABLE_NVRX=1 - SCRIPT_DIR="$(dirname "$(readlink -f "$0")")" # Common arguments and base model specific arguments diff --git a/examples/post_training/modelopt/export.sh b/examples/post_training/modelopt/export.sh index da3109676c2..d0f0b97f2a0 100755 --- a/examples/post_training/modelopt/export.sh +++ b/examples/post_training/modelopt/export.sh @@ -1,8 +1,5 @@ #!/bin/bash -# Disable nvidia-resiliency-ext (incompatible with some container PyTorch versions) -export DISABLE_NVRX=1 - SCRIPT_DIR="$(dirname "$(readlink -f "$0")")" # Common arguments and base model specific arguments diff --git a/examples/post_training/modelopt/finetune.sh b/examples/post_training/modelopt/finetune.sh index 1bf62fff4e3..1affc3c04da 100755 --- a/examples/post_training/modelopt/finetune.sh +++ b/examples/post_training/modelopt/finetune.sh @@ -1,8 +1,5 @@ #!/bin/bash -# Disable nvidia-resiliency-ext (incompatible with some container PyTorch versions) -export DISABLE_NVRX=1 - SCRIPT_DIR="$(dirname "$(readlink -f "$0")")" # Common arguments and base model specific arguments diff --git a/examples/post_training/modelopt/generate.sh b/examples/post_training/modelopt/generate.sh index fa0d0630b15..58578bb3d02 100755 --- a/examples/post_training/modelopt/generate.sh +++ b/examples/post_training/modelopt/generate.sh @@ -1,8 +1,5 @@ #!/bin/bash -# Disable nvidia-resiliency-ext (incompatible with some container PyTorch versions) -export DISABLE_NVRX=1 - SCRIPT_DIR="$(dirname "$(readlink -f "$0")")" # Common arguments and base model specific arguments diff --git a/examples/post_training/modelopt/generation_server.sh b/examples/post_training/modelopt/generation_server.sh index 363cda6caca..9acd61f3d04 100755 --- a/examples/post_training/modelopt/generation_server.sh +++ b/examples/post_training/modelopt/generation_server.sh @@ -1,8 +1,5 @@ #!/bin/bash -# Disable nvidia-resiliency-ext (incompatible with some container PyTorch versions) -export DISABLE_NVRX=1 - SCRIPT_DIR="$(dirname "$(readlink -f "$0")")" # Common arguments and base model specific arguments diff --git a/examples/post_training/modelopt/mmlu.sh b/examples/post_training/modelopt/mmlu.sh index 275f7256b40..36918b89df2 100755 --- a/examples/post_training/modelopt/mmlu.sh +++ b/examples/post_training/modelopt/mmlu.sh @@ -1,8 +1,5 @@ #!/bin/bash -# Disable nvidia-resiliency-ext (incompatible with some container PyTorch versions) -export DISABLE_NVRX=1 - SCRIPT_DIR="$(dirname "$(readlink -f "$0")")" # Common arguments and base model specific arguments diff --git a/examples/post_training/modelopt/offline_feature_extract.sh b/examples/post_training/modelopt/offline_feature_extract.sh index 2d458456a70..1ab7a0fcc85 100755 --- a/examples/post_training/modelopt/offline_feature_extract.sh +++ b/examples/post_training/modelopt/offline_feature_extract.sh @@ -1,8 +1,5 @@ #!/bin/bash -# Disable nvidia-resiliency-ext (incompatible with some container PyTorch versions) -export DISABLE_NVRX=1 - SCRIPT_DIR="$(dirname "$(readlink -f "$0")")" # Common arguments and base model specific arguments diff --git a/examples/post_training/modelopt/prune.sh b/examples/post_training/modelopt/prune.sh index 91aaa953bae..33f3e615e96 100755 --- a/examples/post_training/modelopt/prune.sh +++ b/examples/post_training/modelopt/prune.sh @@ -1,8 +1,5 @@ #!/bin/bash -# Disable nvidia-resiliency-ext (incompatible with some container PyTorch versions) -export DISABLE_NVRX=1 - SCRIPT_DIR="$(dirname "$(readlink -f "$0")")" # Common arguments and base model specific arguments diff --git a/examples/post_training/modelopt/quantize.sh b/examples/post_training/modelopt/quantize.sh index 7261dc422cd..9119ff4ae76 100755 --- a/examples/post_training/modelopt/quantize.sh +++ b/examples/post_training/modelopt/quantize.sh @@ -1,8 +1,5 @@ #!/bin/bash -# Disable nvidia-resiliency-ext (incompatible with some container PyTorch versions) -export DISABLE_NVRX=1 - SCRIPT_DIR="$(dirname "$(readlink -f "$0")")" # Common arguments and base model specific arguments diff --git a/examples/post_training/modelopt/train.sh b/examples/post_training/modelopt/train.sh index 730a4408108..1ebb8bf3d76 100755 --- a/examples/post_training/modelopt/train.sh +++ b/examples/post_training/modelopt/train.sh @@ -1,8 +1,5 @@ #!/bin/bash -# Disable nvidia-resiliency-ext (incompatible with some container PyTorch versions) -export DISABLE_NVRX=1 - SCRIPT_DIR="$(dirname "$(readlink -f "$0")")" # Common arguments and base model specific arguments diff --git a/examples/post_training/modelopt/validate.sh b/examples/post_training/modelopt/validate.sh index bf3d193ad2f..796231e508e 100755 --- a/examples/post_training/modelopt/validate.sh +++ b/examples/post_training/modelopt/validate.sh @@ -1,8 +1,5 @@ #!/bin/bash -# Disable nvidia-resiliency-ext (incompatible with some container PyTorch versions) -export DISABLE_NVRX=1 - SCRIPT_DIR="$(dirname "$(readlink -f "$0")")" # Common arguments and base model specific arguments From 9056598db03e14ebb06ab525d069be1a420446ec Mon Sep 17 00:00:00 2001 From: Jennifer Chen Date: Mon, 27 Apr 2026 11:55:31 -0700 Subject: [PATCH 08/11] revert: restore torch.py get_write_results_queue fallback to main Co-Authored-By: Claude Opus 4.7 (1M context) --- megatron/core/dist_checkpointing/strategies/torch.py | 11 +---------- 1 file changed, 1 insertion(+), 10 deletions(-) diff --git a/megatron/core/dist_checkpointing/strategies/torch.py b/megatron/core/dist_checkpointing/strategies/torch.py index b9a4a6a3691..3e467c1f9dd 100644 --- a/megatron/core/dist_checkpointing/strategies/torch.py +++ b/megatron/core/dist_checkpointing/strategies/torch.py @@ -53,8 +53,6 @@ from .checkpointable import CheckpointableShardedTensor, LocalShardsContainer try: - if os.environ.get("DISABLE_NVRX"): - raise ImportError("NVRX disabled by DISABLE_NVRX env var") from nvidia_resiliency_ext.checkpointing.async_ckpt.core import AsyncRequest as NVRxAsyncRequest from nvidia_resiliency_ext.checkpointing.async_ckpt.state_dict_saver import ( CheckpointMetadataCache, @@ -1042,15 +1040,8 @@ def get_async_strategy(async_strategy: str = "nvrx", module: str = None) -> tupl from nvidia_resiliency_ext.checkpointing.async_ckpt.filesystem_async import ( FileSystemWriterAsync, _results_queue, + get_write_results_queue, ) - try: - from nvidia_resiliency_ext.checkpointing.async_ckpt.filesystem_async import ( - get_write_results_queue, - ) - except ImportError: - from nvidia_resiliency_ext.checkpointing.async_ckpt.filesystem_async import ( - _get_write_results_queue as get_write_results_queue, - ) from nvidia_resiliency_ext.checkpointing.async_ckpt.state_dict_saver import ( CheckpointMetadataCache, save_state_dict_async_finalize, From 0348040a40f911f3fe6e69da1a6e1e8b0beaccd8 Mon Sep 17 00:00:00 2001 From: Jennifer Chen Date: Mon, 27 Apr 2026 12:30:29 -0700 Subject: [PATCH 09/11] fix quantize.py quant config modification Signed-off-by: Jennifer Chen --- examples/post_training/modelopt/quantize.py | 45 +++++---------------- 1 file changed, 9 insertions(+), 36 deletions(-) diff --git a/examples/post_training/modelopt/quantize.py b/examples/post_training/modelopt/quantize.py index 8b4310dd6c9..df723cdef22 100644 --- a/examples/post_training/modelopt/quantize.py +++ b/examples/post_training/modelopt/quantize.py @@ -182,35 +182,6 @@ def _is_last_layers(name: str, num_layers: int = 1, num_layers_to_disable: int = return layer_idx >= num_layers - num_layers_to_disable -def _set_quant_cfg_entry(quant_cfg, quantizer_name, entry_cfg): - """Set a quantizer config entry by name, handling both dict and list formats. - - ModelOpt PR #1094 changed quant_cfg from a flat dict keyed by quantizer name - to an ordered list of QuantizerCfgEntry dicts with an explicit "quantizer_name" - field. This helper writes to either format so callers don't need to branch. - """ - if isinstance(quant_cfg, list): - for entry in quant_cfg: - if entry.get("quantizer_name") == quantizer_name: - entry.update(entry_cfg) - return - new_entry = {"quantizer_name": quantizer_name} - new_entry.update(entry_cfg) - quant_cfg.append(new_entry) - else: - quant_cfg[quantizer_name] = entry_cfg - - -def _get_quant_cfg_entry(quant_cfg, quantizer_name): - """Get a quantizer config entry by name, handling both dict and list formats.""" - if isinstance(quant_cfg, list): - for entry in quant_cfg: - if entry.get("quantizer_name") == quantizer_name: - return entry - return None - return quant_cfg.get(quantizer_name) - - def get_first_layers_disabled_config(config, num_layers: int = 1, num_layers_to_disable: int = 1): """Get a config for `mtq.quantize` with first & last `num_layers_to_disable` layers disabled. @@ -254,6 +225,10 @@ def get_modelopt_torch_quantization_config(): raise ValueError(f"Unsupported quantization config {args.export_quant_cfg}.") mtq_config = QUANT_CFG_CHOICES[args.export_quant_cfg] + if isinstance(mtq_config["quant_cfg"], dict): + # Normalize old dict format to new list format + mtq_config["quant_cfg"] = mtq.normalize_quant_cfg_list(mtq_config["quant_cfg"]) + fp8_config = {"enable": True, "cfg": {"num_bits": (4, 3), "axis": None}} fp4_config = { "enable": True, @@ -265,20 +240,18 @@ def get_modelopt_torch_quantization_config(): } if args.export_quant_cfg == "FP8_DEFAULT_CFG": # Enable Medusa heads and kv-cache quantization - _set_quant_cfg_entry(mtq_config["quant_cfg"], "*medusa_heads**", fp8_config) + mtq_config["quant_cfg"].append({"quantizer_name": "*medusa_heads**", **fp8_config}) if "FP4" in args.export_quant_cfg: # Enable Medusa heads and kv-cache quantization - _set_quant_cfg_entry(mtq_config["quant_cfg"], "*medusa_heads**", fp4_config) + mtq_config["quant_cfg"].append({"quantizer_name": "*medusa_heads**", **fp4_config}) if "AWQ" in args.export_quant_cfg: - weight_quantizer = _get_quant_cfg_entry(mtq_config["quant_cfg"], "*weight_quantizer") - if isinstance(weight_quantizer, list): - weight_quantizer = weight_quantizer[0] + weight_quantizer = mtq._get_quant_cfg_entry(mtq_config["quant_cfg"], "*weight_quantizer") if weight_quantizer is not None: weight_quantizer["block_sizes"][-1] = 128 # Customization if args.disable_qkv_quant: - _set_quant_cfg_entry(mtq_config["quant_cfg"], "*self_attention*", {"enable": False}) + mtq_config["quant_cfg"].append({"quantizer_name": "*self_attention*", "enable": False}) # KV Cache Quantization enable_quant_kv_cache = args.export_kv_cache_quant != "none" @@ -290,7 +263,7 @@ def get_modelopt_torch_quantization_config(): # Weight Only Quantization if args.weight_only: - _set_quant_cfg_entry(mtq_config["quant_cfg"], "*input_quantizer", {"enable": False}) + mtq_config["quant_cfg"].append({"quantizer_name": "*input_quantizer", "enable": False}) if args.num_first_layers_to_skip_quant is not None: mtq_config = get_first_layers_disabled_config( mtq_config, From 4560fed34086a7dcff1ead106a71ff072bdfa012 Mon Sep 17 00:00:00 2001 From: Jennifer Chen Date: Mon, 27 Apr 2026 12:34:17 -0700 Subject: [PATCH 10/11] only support list-type modelopt config in quantize.py Signed-off-by: Jennifer Chen --- examples/post_training/modelopt/quantize.py | 10 ++-------- 1 file changed, 2 insertions(+), 8 deletions(-) diff --git a/examples/post_training/modelopt/quantize.py b/examples/post_training/modelopt/quantize.py index df723cdef22..09ff272069e 100644 --- a/examples/post_training/modelopt/quantize.py +++ b/examples/post_training/modelopt/quantize.py @@ -192,10 +192,7 @@ def get_first_layers_disabled_config(config, num_layers: int = 1, num_layers_to_ predicate = functools.partial( _is_first_layers, num_layers=num_layers, num_layers_to_disable=num_layers_to_disable ) - if isinstance(quant_cfg, list): - quant_cfg.append({"quantizer_name": predicate, "enable": False}) - else: - quant_cfg.update({predicate: {"enable": False}}) + quant_cfg.append({"quantizer_name": predicate, "enable": False}) config["quant_cfg"] = quant_cfg return config @@ -210,10 +207,7 @@ def get_last_layers_disabled_config(config, num_layers: int = 1, num_layers_to_d predicate = functools.partial( _is_last_layers, num_layers=num_layers, num_layers_to_disable=num_layers_to_disable ) - if isinstance(quant_cfg, list): - quant_cfg.append({"quantizer_name": predicate, "enable": False}) - else: - quant_cfg.update({predicate: {"enable": False}}) + quant_cfg.append({"quantizer_name": predicate, "enable": False}) config["quant_cfg"] = quant_cfg return config From 0fc2c08c82501101d76dd16a98c6641b86459ae8 Mon Sep 17 00:00:00 2001 From: Jennifer Chen Date: Mon, 27 Apr 2026 12:47:02 -0700 Subject: [PATCH 11/11] use existing find_quant_cfg_entry_by_path from mtq Signed-off-by: Jennifer Chen --- examples/post_training/modelopt/quantize.py | 7 ++++--- 1 file changed, 4 insertions(+), 3 deletions(-) diff --git a/examples/post_training/modelopt/quantize.py b/examples/post_training/modelopt/quantize.py index 09ff272069e..1e24d97917f 100644 --- a/examples/post_training/modelopt/quantize.py +++ b/examples/post_training/modelopt/quantize.py @@ -239,10 +239,11 @@ def get_modelopt_torch_quantization_config(): # Enable Medusa heads and kv-cache quantization mtq_config["quant_cfg"].append({"quantizer_name": "*medusa_heads**", **fp4_config}) if "AWQ" in args.export_quant_cfg: - weight_quantizer = mtq._get_quant_cfg_entry(mtq_config["quant_cfg"], "*weight_quantizer") - if weight_quantizer is not None: + try: + weight_quantizer = mtq.find_quant_cfg_entry_by_path(mtq_config["quant_cfg"], "*weight_quantizer") weight_quantizer["block_sizes"][-1] = 128 - + except KeyError: + weight_quantizer = None # Customization if args.disable_qkv_quant: mtq_config["quant_cfg"].append({"quantizer_name": "*self_attention*", "enable": False})