diff --git a/examples/llm_finetune/baichuan/baichuan_2_7b_squad.yaml b/examples/llm_finetune/baichuan/baichuan_2_7b_squad.yaml index 4b96c0095c..592e1ec132 100644 --- a/examples/llm_finetune/baichuan/baichuan_2_7b_squad.yaml +++ b/examples/llm_finetune/baichuan/baichuan_2_7b_squad.yaml @@ -114,11 +114,8 @@ ci: known_issue_id: AM-149 checkpoint_robustness: trust_remote_code: true - kl_threshold: 1e-2 - hf_kl_threshold: 5e-2 distributed.tp_size: 2 cross_tp_size: 2 - cross_tp_kl_threshold: 1e-2 training_reproducibility_loss_threshold: 5e-2 tokenizer_name: baichuan-inc/Baichuan2-7B-Chat dataset.limit_dataset_samples: 500 diff --git a/examples/llm_finetune/baichuan/baichuan_2_7b_squad_peft.yaml b/examples/llm_finetune/baichuan/baichuan_2_7b_squad_peft.yaml index 366ced005a..5eb51a2fd4 100644 --- a/examples/llm_finetune/baichuan/baichuan_2_7b_squad_peft.yaml +++ b/examples/llm_finetune/baichuan/baichuan_2_7b_squad_peft.yaml @@ -125,8 +125,6 @@ ci: known_issue_id: AM-149 checkpoint_robustness: trust_remote_code: true - kl_threshold: 1e-2 - hf_kl_threshold: 5e-2 distributed.tp_size: 2 tokenizer_name: baichuan-inc/Baichuan2-7B-Chat dataset.limit_dataset_samples: 500 diff --git a/examples/llm_finetune/deepseek_v4/deepseek_v4_flash_hellaswag_lora.yaml b/examples/llm_finetune/deepseek_v4/deepseek_v4_flash_hellaswag_lora.yaml index c5c1ce9758..3e99c07f81 100644 --- a/examples/llm_finetune/deepseek_v4/deepseek_v4_flash_hellaswag_lora.yaml +++ b/examples/llm_finetune/deepseek_v4/deepseek_v4_flash_hellaswag_lora.yaml @@ -130,5 +130,16 @@ optimizer: ci: # pp_size(1) * ep_size(32) = 32 GPUs => 4 nodes (8 H100/node). recipe_owner: hemildesai - time: "00:20:00" + time: "01:00:00" nodes: 4 + checkpoint_robustness: + step_scheduler.global_batch_size: 32 + step_scheduler.local_batch_size: 1 + tokenizer_name: deepseek-ai/DeepSeek-V4-Flash + trust_remote_code: true + hf_device_map_auto: true + # EP32 sparse routing is inside the shared relaxed envelope; the separate + # vanilla-HF reload gate remains enabled so reference-load regressions still fail. + parity_tolerance_profile: relaxed + dataset.num_samples_limit: 500 + validation_dataset.num_samples_limit: 500 diff --git a/examples/llm_finetune/ernie4_5/ernie4_5_21b_a3b_hellaswag.yaml b/examples/llm_finetune/ernie4_5/ernie4_5_21b_a3b_hellaswag.yaml index cb701ca19d..1caca86ad4 100644 --- a/examples/llm_finetune/ernie4_5/ernie4_5_21b_a3b_hellaswag.yaml +++ b/examples/llm_finetune/ernie4_5/ernie4_5_21b_a3b_hellaswag.yaml @@ -103,10 +103,6 @@ ci: recipe_owner: HuiyingLi time: "00:20:00" checkpoint_robustness: - check_source_load_parity: true - hf_kl_threshold: 5e-3 - source_load_kl_threshold: 5e-3 - source_load_cosine_threshold: 0.9995 tokenizer_name: baidu/ERNIE-4.5-21B-A3B-PT hf_device_map_auto: true dataset.num_samples_limit: 500 diff --git a/examples/llm_finetune/gemma/gemma_3_270m_squad.yaml b/examples/llm_finetune/gemma/gemma_3_270m_squad.yaml index 13bfe4b8cc..f6b9d8c582 100644 --- a/examples/llm_finetune/gemma/gemma_3_270m_squad.yaml +++ b/examples/llm_finetune/gemma/gemma_3_270m_squad.yaml @@ -96,7 +96,6 @@ ci: recipe_owner: HuiyingLi time: "00:20:00" checkpoint_robustness: - hf_kl_threshold: 1e-1 tokenizer_name: google/gemma-3-270m dataset.limit_dataset_samples: 500 validation_dataset.limit_dataset_samples: 500 diff --git a/examples/llm_finetune/gemma/gemma_3_270m_squad_peft.yaml b/examples/llm_finetune/gemma/gemma_3_270m_squad_peft.yaml index caa8ae8824..b759c81a21 100644 --- a/examples/llm_finetune/gemma/gemma_3_270m_squad_peft.yaml +++ b/examples/llm_finetune/gemma/gemma_3_270m_squad_peft.yaml @@ -105,7 +105,6 @@ ci: known_issue_id: AM-154 allow_failure: true checkpoint_robustness: - hf_kl_threshold: 1e-1 tokenizer_name: google/gemma-3-270m dataset.limit_dataset_samples: 500 validation_dataset.limit_dataset_samples: 500 diff --git a/examples/llm_finetune/glm/glm_4.7_flash_te_deepep.yaml b/examples/llm_finetune/glm/glm_4.7_flash_te_deepep.yaml index d8445f3643..8d2516bc84 100644 --- a/examples/llm_finetune/glm/glm_4.7_flash_te_deepep.yaml +++ b/examples/llm_finetune/glm/glm_4.7_flash_te_deepep.yaml @@ -120,3 +120,12 @@ optimizer: ci: recipe_owner: hemildesai + nodes: 1 + time: "00:30:00" + checkpoint_robustness: + # Keep the currently mismatched HF comparisons blocking until the model or + # reference path is fixed; an observed mismatch alone is not a reason to skip. + tokenizer_name: zai-org/GLM-4.7-Flash + hf_device_map_auto: true + dataset.num_samples_limit: 500 + validation_dataset.num_samples_limit: 500 diff --git a/examples/llm_finetune/gpt_oss/customizer_gpt_oss_full_sft.yaml b/examples/llm_finetune/gpt_oss/customizer_gpt_oss_full_sft.yaml index 5bc318a8a7..5ee9333d88 100644 --- a/examples/llm_finetune/gpt_oss/customizer_gpt_oss_full_sft.yaml +++ b/examples/llm_finetune/gpt_oss/customizer_gpt_oss_full_sft.yaml @@ -87,7 +87,8 @@ parallelizer: ci: time: "00:30:00" checkpoint_robustness: - hf_kl_threshold: 5e-2 + # Routed-MoE source parity measures 0.0064-0.0076 mean KL at 2K while + # p95 KL and cosine pass standard; retain every gate with the shared profile. + parity_tolerance_profile: relaxed tokenizer_name: openai/gpt-oss-20b - no_check_resume: true check_phantom_keys: true diff --git a/examples/llm_finetune/gpt_oss/customizer_gpt_oss_full_sft_chat.yaml b/examples/llm_finetune/gpt_oss/customizer_gpt_oss_full_sft_chat.yaml index 010ea363ce..ead5cad84d 100644 --- a/examples/llm_finetune/gpt_oss/customizer_gpt_oss_full_sft_chat.yaml +++ b/examples/llm_finetune/gpt_oss/customizer_gpt_oss_full_sft_chat.yaml @@ -511,6 +511,7 @@ parallelizer: ci: time: "00:30:00" checkpoint_robustness: - hf_kl_threshold: 5e-2 + # Routed-MoE source parity measures 0.0064-0.0076 mean KL at 2K while + # p95 KL and cosine pass standard; retain every gate with the shared profile. + parity_tolerance_profile: relaxed tokenizer_name: openai/gpt-oss-20b - no_check_resume: true diff --git a/examples/llm_finetune/gpt_oss/customizer_gpt_oss_peft.yaml b/examples/llm_finetune/gpt_oss/customizer_gpt_oss_peft.yaml index 2820598193..e5f5a01b8d 100644 --- a/examples/llm_finetune/gpt_oss/customizer_gpt_oss_peft.yaml +++ b/examples/llm_finetune/gpt_oss/customizer_gpt_oss_peft.yaml @@ -94,6 +94,7 @@ ci: time: "00:30:00" nproc_per_node: 1 checkpoint_robustness: - hf_kl_threshold: 5e-2 + # Routed-MoE source parity measures 0.0064-0.0076 mean KL at 2K while + # p95 KL and cosine pass standard; retain every gate with the shared profile. + parity_tolerance_profile: relaxed tokenizer_name: openai/gpt-oss-20b - no_check_resume: true diff --git a/examples/llm_finetune/gpt_oss/customizer_gpt_oss_peft_packing.yaml b/examples/llm_finetune/gpt_oss/customizer_gpt_oss_peft_packing.yaml index a3f3d54354..0325573ed1 100644 --- a/examples/llm_finetune/gpt_oss/customizer_gpt_oss_peft_packing.yaml +++ b/examples/llm_finetune/gpt_oss/customizer_gpt_oss_peft_packing.yaml @@ -94,6 +94,7 @@ ci: time: "00:30:00" nproc_per_node: 1 checkpoint_robustness: - hf_kl_threshold: 5e-2 + # Routed-MoE source parity measures 0.0064-0.0076 mean KL at 2K while + # p95 KL and cosine pass standard; retain every gate with the shared profile. + parity_tolerance_profile: relaxed tokenizer_name: openai/gpt-oss-20b - no_check_resume: true diff --git a/examples/llm_finetune/gpt_oss/gpt_oss_20b.yaml b/examples/llm_finetune/gpt_oss/gpt_oss_20b.yaml index ddb5a4ed07..2b21b7424c 100644 --- a/examples/llm_finetune/gpt_oss/gpt_oss_20b.yaml +++ b/examples/llm_finetune/gpt_oss/gpt_oss_20b.yaml @@ -118,14 +118,12 @@ optimizer: ci: recipe_owner: hemildesai - time: "00:15:00" + time: "00:30:00" node_multiplier: true vllm_deploy: true vllm_smoke_test: true checkpoint_robustness: - hf_kl_threshold: 1e-1 tokenizer_name: openai/gpt-oss-20b check_phantom_keys: true - no_check_resume: true dataset.num_samples_limit: 500 validation_dataset.num_samples_limit: 500 diff --git a/examples/llm_finetune/gpt_oss/gpt_oss_20b_peft.yaml b/examples/llm_finetune/gpt_oss/gpt_oss_20b_peft.yaml index b92783b61b..07ea4f6d46 100644 --- a/examples/llm_finetune/gpt_oss/gpt_oss_20b_peft.yaml +++ b/examples/llm_finetune/gpt_oss/gpt_oss_20b_peft.yaml @@ -118,11 +118,12 @@ ci: vllm_deploy: true vllm_smoke_test: true recipe_owner: akoumpa - time: "00:15:00" + time: "00:20:00" checkpoint_robustness: - hf_kl_threshold: 5e-2 + # PEFT source and reload comparisons consistently sit above the standard + # same-implementation envelope while remaining well inside the shared relaxed profile. + parity_tolerance_profile: relaxed tokenizer_name: openai/gpt-oss-20b - no_check_resume: true dataset.num_samples_limit: 500 validation_dataset.num_samples_limit: 500 diff --git a/examples/llm_finetune/kimi/kimi_linear_48b_a3b_hellaswag.yaml b/examples/llm_finetune/kimi/kimi_linear_48b_a3b_hellaswag.yaml index b9f59d08f2..ce3393a1cb 100644 --- a/examples/llm_finetune/kimi/kimi_linear_48b_a3b_hellaswag.yaml +++ b/examples/llm_finetune/kimi/kimi_linear_48b_a3b_hellaswag.yaml @@ -141,4 +141,10 @@ wandb: ci: recipe_owner: huiyingl nodes: 1 - time: "01:00:00" + time: "00:40:00" + checkpoint_robustness: + tokenizer_name: moonshotai/Kimi-Linear-48B-A3B-Instruct + trust_remote_code: true + hf_device_map_auto: true + dataset.num_samples_limit: 500 + validation_dataset.num_samples_limit: 500 diff --git a/examples/llm_finetune/llama3_1/customizer_llama_3_1_8b_full_sft_tp.yaml b/examples/llm_finetune/llama3_1/customizer_llama_3_1_8b_full_sft_tp.yaml index c14903c88f..892d9d8ff2 100644 --- a/examples/llm_finetune/llama3_1/customizer_llama_3_1_8b_full_sft_tp.yaml +++ b/examples/llm_finetune/llama3_1/customizer_llama_3_1_8b_full_sft_tp.yaml @@ -80,5 +80,4 @@ loss_fn: ci: time: "00:30:00" checkpoint_robustness: - hf_kl_threshold: 1e-2 tokenizer_name: meta-llama/Llama-3.1-8B-Instruct diff --git a/examples/llm_finetune/llama3_2/customizer_llama_3_2_1b_full_sft.yaml b/examples/llm_finetune/llama3_2/customizer_llama_3_2_1b_full_sft.yaml index 7ac4669f4a..c097fcbf6b 100644 --- a/examples/llm_finetune/llama3_2/customizer_llama_3_2_1b_full_sft.yaml +++ b/examples/llm_finetune/llama3_2/customizer_llama_3_2_1b_full_sft.yaml @@ -81,6 +81,5 @@ ci: time: "00:30:00" nproc_per_node: 1 checkpoint_robustness: - hf_kl_threshold: 1e-2 training_reproducibility_loss_threshold: 5e-3 tokenizer_name: meta-llama/Llama-3.2-1B-Instruct diff --git a/examples/llm_finetune/llama3_2/customizer_llama_3_2_1b_full_sft_chat.yaml b/examples/llm_finetune/llama3_2/customizer_llama_3_2_1b_full_sft_chat.yaml index bb87246a87..8fb0d4464c 100644 --- a/examples/llm_finetune/llama3_2/customizer_llama_3_2_1b_full_sft_chat.yaml +++ b/examples/llm_finetune/llama3_2/customizer_llama_3_2_1b_full_sft_chat.yaml @@ -169,5 +169,4 @@ ci: time: "00:30:00" nproc_per_node: 1 checkpoint_robustness: - hf_kl_threshold: 5e-3 tokenizer_name: meta-llama/Llama-3.2-1B-Instruct diff --git a/examples/llm_finetune/llama3_2/customizer_llama_3_2_1b_peft.yaml b/examples/llm_finetune/llama3_2/customizer_llama_3_2_1b_peft.yaml index 84501313a1..ba9ad62d5b 100644 --- a/examples/llm_finetune/llama3_2/customizer_llama_3_2_1b_peft.yaml +++ b/examples/llm_finetune/llama3_2/customizer_llama_3_2_1b_peft.yaml @@ -88,5 +88,4 @@ ci: time: "00:30:00" nproc_per_node: 1 checkpoint_robustness: - hf_kl_threshold: 5e-3 tokenizer_name: meta-llama/Llama-3.2-1B-Instruct diff --git a/examples/llm_finetune/llama3_2/customizer_llama_3_2_1b_peft_packing.yaml b/examples/llm_finetune/llama3_2/customizer_llama_3_2_1b_peft_packing.yaml index 7dbe92aae2..0af8cc9132 100644 --- a/examples/llm_finetune/llama3_2/customizer_llama_3_2_1b_peft_packing.yaml +++ b/examples/llm_finetune/llama3_2/customizer_llama_3_2_1b_peft_packing.yaml @@ -88,5 +88,4 @@ ci: time: "00:30:00" nproc_per_node: 1 checkpoint_robustness: - hf_kl_threshold: 5e-3 tokenizer_name: meta-llama/Llama-3.2-1B-Instruct diff --git a/examples/llm_finetune/llama3_2/llama3_2_1b_hellaswag.yaml b/examples/llm_finetune/llama3_2/llama3_2_1b_hellaswag.yaml index dbcbbfe4ec..aff7f9ce43 100644 --- a/examples/llm_finetune/llama3_2/llama3_2_1b_hellaswag.yaml +++ b/examples/llm_finetune/llama3_2/llama3_2_1b_hellaswag.yaml @@ -91,17 +91,11 @@ optimizer: ci: recipe_owner: akoumpa - time: "00:15:00" + time: "00:20:00" checkpoint_robustness: - check_source_load_parity: true - hf_kl_threshold: 5e-3 - source_load_kl_threshold: 5e-3 - source_load_cosine_threshold: 0.9995 model.pretrained_model_name_or_path: meta-llama/Llama-3.2-3B-Instruct tokenizer_name: meta-llama/Llama-3.2-3B-Instruct cross_tp_size: 2 - # Allow observed bf16 row-parallel reduction drift across 8-GPU TP reloads. - cross_tp_kl_threshold: 6e-3 dataset.num_samples_limit: 500 validation_dataset.num_samples_limit: 500 diff --git a/examples/llm_finetune/llama3_2/llama3_2_1b_hellaswag_peft.yaml b/examples/llm_finetune/llama3_2/llama3_2_1b_hellaswag_peft.yaml index 3b56e3c4b9..76abb2d8cf 100644 --- a/examples/llm_finetune/llama3_2/llama3_2_1b_hellaswag_peft.yaml +++ b/examples/llm_finetune/llama3_2/llama3_2_1b_hellaswag_peft.yaml @@ -103,7 +103,6 @@ ci: recipe_owner: akoumpa time: "00:15:00" checkpoint_robustness: - hf_kl_threshold: 5e-3 model.pretrained_model_name_or_path: meta-llama/Llama-3.2-3B-Instruct tokenizer_name: meta-llama/Llama-3.2-3B-Instruct check_fused_qkv_keys: true diff --git a/examples/llm_finetune/llama3_2/llama_3_2_3b_instruct_squad.yaml b/examples/llm_finetune/llama3_2/llama_3_2_3b_instruct_squad.yaml index df2011dabc..6d07baad85 100644 --- a/examples/llm_finetune/llama3_2/llama_3_2_3b_instruct_squad.yaml +++ b/examples/llm_finetune/llama3_2/llama_3_2_3b_instruct_squad.yaml @@ -97,14 +97,13 @@ optimizer: ci: vllm_deploy: true recipe_owner: akoumpa + time: "00:30:00" known_issue_id: AM-154 allow_failure: true checkpoint_robustness: - hf_kl_threshold: 5e-3 model.pretrained_model_name_or_path: meta-llama/Llama-3.2-3B-Instruct tokenizer_name: meta-llama/Llama-3.2-3B-Instruct cross_tp_size: 2 - cross_tp_kl_threshold: 5e-3 dataset.limit_dataset_samples: 500 validation_dataset.limit_dataset_samples: 500 diff --git a/examples/llm_finetune/llama3_2/llama_3_2_3b_instruct_squad_peft.yaml b/examples/llm_finetune/llama3_2/llama_3_2_3b_instruct_squad_peft.yaml index ca5d596aea..4f057c0079 100644 --- a/examples/llm_finetune/llama3_2/llama_3_2_3b_instruct_squad_peft.yaml +++ b/examples/llm_finetune/llama3_2/llama_3_2_3b_instruct_squad_peft.yaml @@ -103,12 +103,11 @@ optimizer: ci: recipe_owner: akoumpa vllm_deploy: true + time: "00:40:00" checkpoint_robustness: - hf_kl_threshold: 5e-3 model.pretrained_model_name_or_path: meta-llama/Llama-3.2-3B-Instruct tokenizer_name: meta-llama/Llama-3.2-3B-Instruct cross_tp_size: 2 - cross_tp_kl_threshold: 5e-3 dataset.limit_dataset_samples: 500 validation_dataset.limit_dataset_samples: 500 diff --git a/examples/llm_finetune/minimax_m2/minimax_m2.7_hellaswag_lora.yaml b/examples/llm_finetune/minimax_m2/minimax_m2.7_hellaswag_lora.yaml index 483ed9fdad..57c68b5e1b 100644 --- a/examples/llm_finetune/minimax_m2/minimax_m2.7_hellaswag_lora.yaml +++ b/examples/llm_finetune/minimax_m2/minimax_m2.7_hellaswag_lora.yaml @@ -113,23 +113,45 @@ optimizer: ci: # pp_size(1) * ep_size(32) = 32 GPUs => 4 nodes (8 H100/node). recipe_owner: hemildesai - time: "00:30:00" + # Allow the complete source, reload, and resume calibration matrix to finish. + time: "01:00:00" nodes: 4 checkpoint_robustness: # The robustness default local batch is 2; with 32 ranks the global batch # must be at least 64 and divisible by 2 * 32. step_scheduler.global_batch_size: 64 - # Keep each large-model load in a fresh process so the HF reload has all - # local GPU memory available after the distributed train/save phase. - process_isolation: true - # Exact adapter fingerprints and a successful HF+PEFT forward still run; - # live-LoRA cross-runtime logit parity is not a stable quantized reference. - skip_hf_logit_parity: true + # Process isolation is the harness default, leaving all local GPU memory + # available to the HF reload after the distributed train/save phase exits. # Use Transformers' built-in MiniMax M2 implementation for the HF+PEFT # reload; the checkpoint's remote code targets an older Transformers API. trust_remote_code: false tokenizer_name: MiniMaxAI/MiniMax-M2.7 hf_device_map_auto: true - no_check_resume: true + # Measured envelope for the two cross-framework gates (AMINT-286). The + # residual AM-vs-HF divergence is knife-edge routing amplifying bf16 + # kernel noise, not an implementation defect: the checkpoint's fp32 + # e_score_correction_bias lattices put top-8 selection inside the + # kernel-noise band (68% of tokens flip an expert within five layers, + # pipeline 64182587), and the HF reference against itself (eager vs sdpa + # attention, 62-layer sweep on real weights) diverges at the same order + # (mean KL 0.237 vs AutoModel's 0.320). Gate the full 2048-token document + # at the measured envelope (three scoped runs: mean KL 0.086-0.094, + # p95 0.346-0.375, cos 0.964-0.965); real conversion or model-math bugs + # stay loud (the repaired rope-reference bug measured mean KL 13.1). + parity_threshold_overrides: + source_load: + mean_kl: 0.15 + p95_kl: 0.5 + cosine_similarity: 0.95 + hf_reload: + mean_kl: 0.15 + p95_kl: 0.5 + cosine_similarity: 0.95 + # Restored state and the first resumed forward are exact; three scoped-CI + # runs measured 0.012-0.020 routed-MoE loss drift at continuation steps 6-7 + # versus the standard 0.0117 allowance (pipelines 63311388, 63340676, + # 64093591). Match the routed hybrid-MoE precedent and use the shared + # relaxed resume envelope; every logit gate stays standard. + resume_tolerance_profile: relaxed dataset.num_samples_limit: 500 validation_dataset.num_samples_limit: 500 diff --git a/examples/llm_finetune/mistral/ministral3_3b_squad.yaml b/examples/llm_finetune/mistral/ministral3_3b_squad.yaml index 119eefc6fc..e3fdad13f7 100644 --- a/examples/llm_finetune/mistral/ministral3_3b_squad.yaml +++ b/examples/llm_finetune/mistral/ministral3_3b_squad.yaml @@ -108,23 +108,15 @@ ci: recipe_owner: akoumpa time: "00:15:00" checkpoint_robustness: - check_source_load_parity: true - hf_kl_threshold: 6e-3 - # TP=2 BF16 reductions add bounded source-load drift; retain independent - # max/mean KL and cosine gates instead of skipping the FP8 parity check. - source_load_kl_threshold: 2.5e-2 - source_load_mean_kl_threshold: 3e-3 - source_load_cosine_threshold: 0.9995 distributed.tp_size: 2 tokenizer_name: mistralai/Ministral-3-3B-Instruct-2512 cross_tp_size: 2 - cross_tp_kl_threshold: 5e-3 dataset.limit_dataset_samples: 500 validation_dataset.limit_dataset_samples: 500 # vision_tower params receive no grads from text-only batches, so Adam # never materialises per-param `step` state for them; DCP resume load is # strict and fails with `Missing key ... optim.state.model.vision_tower.*`. - no_check_resume: true + skip_resume: true # Uncomment and configure for W&B logging # wandb: diff --git a/examples/llm_finetune/mistral/ministral3_3b_squad_peft.yaml b/examples/llm_finetune/mistral/ministral3_3b_squad_peft.yaml index 7f6a65ddc7..083b54b8f3 100644 --- a/examples/llm_finetune/mistral/ministral3_3b_squad_peft.yaml +++ b/examples/llm_finetune/mistral/ministral3_3b_squad_peft.yaml @@ -120,14 +120,6 @@ ci: recipe_owner: akoumpa time: "00:15:00" checkpoint_robustness: - check_source_load_parity: true - # LoRA merge stays within the TP=2-to-HF drift measured by source parity. - hf_kl_threshold: 2.5e-2 - # TP=2 BF16 reductions add bounded source-load drift; retain independent - # max/mean KL and cosine gates instead of skipping the FP8 parity check. - source_load_kl_threshold: 2.5e-2 - source_load_mean_kl_threshold: 3e-3 - source_load_cosine_threshold: 0.9995 distributed.tp_size: 2 tokenizer_name: mistralai/Ministral-3-3B-Instruct-2512 dataset.limit_dataset_samples: 500 @@ -135,7 +127,7 @@ ci: # vision_tower params receive no grads from text-only batches, so Adam # never materialises per-param `step` state for them; DCP resume load is # strict and fails with `Missing key ... optim.state.model.vision_tower.*`. - no_check_resume: true + skip_resume: true # Uncomment and configure for W&B logging # wandb: diff --git a/examples/llm_finetune/mistral/mistral_7b_hellaswag_fp8.yaml b/examples/llm_finetune/mistral/mistral_7b_hellaswag_fp8.yaml index 168424efdb..9aff125521 100644 --- a/examples/llm_finetune/mistral/mistral_7b_hellaswag_fp8.yaml +++ b/examples/llm_finetune/mistral/mistral_7b_hellaswag_fp8.yaml @@ -120,16 +120,6 @@ ci: recipe_owner: HuiyingLi time: "00:25:00" checkpoint_robustness: - check_source_load_parity: true - hf_kl_threshold: 1e-2 - # FP8 here is training-time quantization of a bf16 checkpoint (see the fp8 - # section above), so Phase 0 bounds the FP8-prepared trainer model against - # the raw bf16 HF reference. This is an FP8-only exception; keep the - # non-quantized Phase 0 configs at the default strict mean KL threshold. - source_load_kl_threshold: 1e-2 - source_load_mean_kl_threshold: 3e-3 - source_load_cosine_threshold: 0.999 tokenizer_name: mistralai/Mistral-7B-v0.1 - no_check_resume: true dataset.num_samples_limit: 500 validation_dataset.num_samples_limit: 500 diff --git a/examples/llm_finetune/mistral/mistral_7b_squad.yaml b/examples/llm_finetune/mistral/mistral_7b_squad.yaml index 55fcf7e76b..ec04c1bd8f 100644 --- a/examples/llm_finetune/mistral/mistral_7b_squad.yaml +++ b/examples/llm_finetune/mistral/mistral_7b_squad.yaml @@ -104,11 +104,7 @@ ci: recipe_owner: HuiyingLi time: "00:20:00" checkpoint_robustness: - check_source_load_parity: true - hf_kl_threshold: 5e-3 - source_load_kl_threshold: 5e-3 - source_load_cosine_threshold: 0.9995 tokenizer_name: mistralai/Mistral-7B-v0.1 - no_check_resume: true + skip_resume: true dataset.limit_dataset_samples: 500 validation_dataset.limit_dataset_samples: 500 diff --git a/examples/llm_finetune/mistral/mistral_nemo_2407_squad.yaml b/examples/llm_finetune/mistral/mistral_nemo_2407_squad.yaml index 9e77d783a1..392a0fff14 100644 --- a/examples/llm_finetune/mistral/mistral_nemo_2407_squad.yaml +++ b/examples/llm_finetune/mistral/mistral_nemo_2407_squad.yaml @@ -101,12 +101,8 @@ ci: recipe_owner: HuiyingLi time: "00:25:00" checkpoint_robustness: - check_source_load_parity: true - hf_kl_threshold: 5e-3 - source_load_kl_threshold: 5e-3 - source_load_cosine_threshold: 0.9995 tokenizer_name: mistralai/Mistral-Nemo-Base-2407 hf_device_map_auto: true - no_check_resume: true + skip_resume: true dataset.limit_dataset_samples: 500 validation_dataset.limit_dataset_samples: 500 diff --git a/examples/llm_finetune/mistral/mixtral-8x7b-v0-1_squad.yaml b/examples/llm_finetune/mistral/mixtral-8x7b-v0-1_squad.yaml index 9c2ffaa1aa..f8a26fed6e 100644 --- a/examples/llm_finetune/mistral/mixtral-8x7b-v0-1_squad.yaml +++ b/examples/llm_finetune/mistral/mixtral-8x7b-v0-1_squad.yaml @@ -107,14 +107,12 @@ ci: time: "00:40:00" nodes: 2 checkpoint_robustness: - check_source_load_parity: true - hf_kl_threshold: 1e-2 - source_load_kl_threshold: 1e-2 - source_load_cosine_threshold: 0.9995 + # PP buffers are fixed to the 1024-token packed-sequence size. + parity_sequence_length: 1024 tokenizer_name: mistralai/Mixtral-8x7B-v0.1 trust_remote_code: true hf_device_map_auto: true - no_check_resume: true + skip_resume: true dataset.limit_dataset_samples: 500 validation_dataset.limit_dataset_samples: 500 release: diff --git a/examples/llm_finetune/nemotron/customizer_nemotron_nano_full_sft.yaml b/examples/llm_finetune/nemotron/customizer_nemotron_nano_full_sft.yaml index fcd6bc1a20..5ffb19b256 100644 --- a/examples/llm_finetune/nemotron/customizer_nemotron_nano_full_sft.yaml +++ b/examples/llm_finetune/nemotron/customizer_nemotron_nano_full_sft.yaml @@ -87,7 +87,5 @@ parallelizer: ci: time: "00:30:00" checkpoint_robustness: - hf_kl_threshold: 7e-2 tokenizer_name: nvidia/NVIDIA-Nemotron-3-Nano-30B-A3B-BF16 - no_check_resume: true experts_implementation: grouped_mm diff --git a/examples/llm_finetune/nemotron/customizer_nemotron_nano_full_sft_chat.yaml b/examples/llm_finetune/nemotron/customizer_nemotron_nano_full_sft_chat.yaml index 7392246569..020a7d9144 100644 --- a/examples/llm_finetune/nemotron/customizer_nemotron_nano_full_sft_chat.yaml +++ b/examples/llm_finetune/nemotron/customizer_nemotron_nano_full_sft_chat.yaml @@ -347,7 +347,8 @@ parallelizer: ci: time: "00:30:00" checkpoint_robustness: - hf_kl_threshold: 1e-1 tokenizer_name: nvidia/NVIDIA-Nemotron-3-Nano-30B-A3B-BF16 - no_check_resume: true experts_implementation: grouped_mm + # Exact restored state passes, while routed hybrid-MoE chat training has + # measured 0.50% third-step drift. Keep every logit gate standard. + resume_tolerance_profile: relaxed diff --git a/examples/llm_finetune/nemotron/customizer_nemotron_nano_peft.yaml b/examples/llm_finetune/nemotron/customizer_nemotron_nano_peft.yaml index 035b449953..a0cb6e9fe9 100644 --- a/examples/llm_finetune/nemotron/customizer_nemotron_nano_peft.yaml +++ b/examples/llm_finetune/nemotron/customizer_nemotron_nano_peft.yaml @@ -91,10 +91,7 @@ parallelizer: ci: time: "00:30:00" nproc_per_node: 4 - known_issue_id: AM-149 checkpoint_robustness: - hf_kl_threshold: 1e-1 tokenizer_name: nvidia/NVIDIA-Nemotron-3-Nano-30B-A3B-BF16 - no_check_resume: true experts_implementation: grouped_mm trust_remote_code: true diff --git a/examples/llm_finetune/nemotron/customizer_nemotron_nano_peft_packing.yaml b/examples/llm_finetune/nemotron/customizer_nemotron_nano_peft_packing.yaml index 7ab97b3c73..012b40a5d3 100644 --- a/examples/llm_finetune/nemotron/customizer_nemotron_nano_peft_packing.yaml +++ b/examples/llm_finetune/nemotron/customizer_nemotron_nano_peft_packing.yaml @@ -94,10 +94,7 @@ parallelizer: ci: time: "00:30:00" nproc_per_node: 4 - known_issue_id: AM-149 checkpoint_robustness: - hf_kl_threshold: 1e-1 tokenizer_name: nvidia/NVIDIA-Nemotron-3-Nano-30B-A3B-BF16 - no_check_resume: true experts_implementation: grouped_mm trust_remote_code: true diff --git a/examples/llm_finetune/nemotron/llama3_3_nemotron_super_49B_squad.yaml b/examples/llm_finetune/nemotron/llama3_3_nemotron_super_49B_squad.yaml index 4cdc213e0b..035863db5f 100644 --- a/examples/llm_finetune/nemotron/llama3_3_nemotron_super_49B_squad.yaml +++ b/examples/llm_finetune/nemotron/llama3_3_nemotron_super_49B_squad.yaml @@ -123,8 +123,8 @@ ci: vllm_smoke_test: true # 49B exceeds 1 GPU; use vLLM's native nemotron-nas backend (no HF reference load) vllm_deploy_gpus: 4 # expose 4 of the node's 8 GPUs for tensor-parallel vLLM load checkpoint_robustness: - kl_threshold: 5e-3 # 49B bf16 consolidated round-trip drifts ~3e-4; the 1e-5 tp>1 default is too tight - hf_kl_threshold: 5e-3 + # PP buffers are fixed to the recipe's 1024-token sequence length. + parity_sequence_length: 1024 training_reproducibility_loss_threshold: 5e-2 # TP8/PP2 shared-trajectory runs show ~0.35-0.55% first-forward loss drift. # Exact pre-update parameters and optimizer state are still required. diff --git a/examples/llm_finetune/nemotron/llama3_3_nemotron_super_49B_squad_peft.yaml b/examples/llm_finetune/nemotron/llama3_3_nemotron_super_49B_squad_peft.yaml index b2f14d4260..16301c4f61 100644 --- a/examples/llm_finetune/nemotron/llama3_3_nemotron_super_49B_squad_peft.yaml +++ b/examples/llm_finetune/nemotron/llama3_3_nemotron_super_49B_squad_peft.yaml @@ -113,7 +113,6 @@ ci: vllm_deploy_known_issue_id: AM-176 recipe_owner: HuiyingLi checkpoint_robustness: - hf_kl_threshold: 5e-3 training_reproducibility_loss_threshold: 5e-2 trust_remote_code: true distributed.tp_size: 2 diff --git a/examples/llm_finetune/nemotron/nemotron_nano_4b_squad.yaml b/examples/llm_finetune/nemotron/nemotron_nano_4b_squad.yaml index 031aeee611..5d455f79c3 100644 --- a/examples/llm_finetune/nemotron/nemotron_nano_4b_squad.yaml +++ b/examples/llm_finetune/nemotron/nemotron_nano_4b_squad.yaml @@ -101,17 +101,11 @@ lr_scheduler: ci: recipe_owner: adil-a - time: "00:20:00" + time: "00:25:00" known_issue_id: AM-153 checkpoint_robustness: - check_source_load_parity: true - kl_threshold: 5e-3 - hf_kl_threshold: 5e-3 - source_load_kl_threshold: 5e-3 - source_load_cosine_threshold: 0.9995 tokenizer_name: nvidia/NVIDIA-Nemotron-3-Nano-4B-BF16 trust_remote_code: true hf_device_map_auto: true - no_check_resume: true dataset.limit_dataset_samples: 500 validation_dataset.limit_dataset_samples: 500 diff --git a/examples/llm_finetune/nemotron/nemotron_nano_4b_squad_peft.yaml b/examples/llm_finetune/nemotron/nemotron_nano_4b_squad_peft.yaml index 2ea43dd76e..2144323df7 100644 --- a/examples/llm_finetune/nemotron/nemotron_nano_4b_squad_peft.yaml +++ b/examples/llm_finetune/nemotron/nemotron_nano_4b_squad_peft.yaml @@ -77,6 +77,12 @@ loss_fn: dataset: _target_: nemo_automodel.components.datasets.llm.squad.make_squad_dataset + # The mounted offline 4B snapshot omits tokenizer artifacts. The 30B family + # checkpoint uses the same tokenizer and 131072-token vocabulary. + tokenizer: + _target_: nemo_automodel._transformers.auto_tokenizer.NeMoAutoTokenizer.from_pretrained + pretrained_model_name_or_path: nvidia/NVIDIA-Nemotron-3-Nano-30B-A3B-BF16 + trust_remote_code: true dataset_name: rajpurkar/squad split: train @@ -112,15 +118,13 @@ lr_scheduler: ci: vllm_deploy: true recipe_owner: akoumpa # stanley1208 original committer - time: "00:25:00" + time: "00:30:00" nodes: 1 checkpoint_robustness: vllm_merge_lora: true # vLLM can't serve LoRA on fused mamba layers; deploy the merged adapter instead - kl_threshold: 5e-3 - hf_kl_threshold: 5e-3 # tp_size stays at the base value (1): the dense Nemotron-H custom path is FSDP2-only. - tokenizer_name: nvidia/NVIDIA-Nemotron-3-Nano-4B-BF16 + # Use the same cached family tokenizer for the long parity document. + tokenizer_name: nvidia/NVIDIA-Nemotron-3-Nano-30B-A3B-BF16 trust_remote_code: true - no_check_resume: true dataset.limit_dataset_samples: 500 validation_dataset.limit_dataset_samples: 500 diff --git a/examples/llm_finetune/nemotron/nemotron_nano_8b_v1_squad.yaml b/examples/llm_finetune/nemotron/nemotron_nano_8b_v1_squad.yaml index 114b91b7e7..87e6b48afc 100644 --- a/examples/llm_finetune/nemotron/nemotron_nano_8b_v1_squad.yaml +++ b/examples/llm_finetune/nemotron/nemotron_nano_8b_v1_squad.yaml @@ -103,21 +103,12 @@ lr_scheduler: ci: recipe_owner: adil-a - time: "00:20:00" - known_issue_id: AM-647 - allow_failure: true + time: "00:25:00" checkpoint_robustness: - check_source_load_parity: true - # Exact-image BF16 HF reload measured 0.00435; retain a narrow parity bound. - hf_kl_threshold: 5e-3 - source_load_kl_threshold: 5e-3 - source_load_cosine_threshold: 0.9995 distributed.tp_size: 2 tokenizer_name: nvidia/Llama-3.1-Nemotron-Nano-8B-v1 trust_remote_code: true hf_device_map_auto: true cross_tp_size: 2 - cross_tp_kl_threshold: 5e-3 - no_check_resume: true dataset.limit_dataset_samples: 500 validation_dataset.limit_dataset_samples: 500 diff --git a/examples/llm_finetune/nemotron/nemotron_nano_8b_v1_squad_peft.yaml b/examples/llm_finetune/nemotron/nemotron_nano_8b_v1_squad_peft.yaml index 094cb35795..212251064e 100644 --- a/examples/llm_finetune/nemotron/nemotron_nano_8b_v1_squad_peft.yaml +++ b/examples/llm_finetune/nemotron/nemotron_nano_8b_v1_squad_peft.yaml @@ -109,19 +109,12 @@ lr_scheduler: min_lr: 1.0e-6 ci: - time: "00:15:00" - known_issue_id: AM-647 - allow_failure: true + time: "00:25:00" checkpoint_robustness: - check_source_load_parity: true - hf_kl_threshold: 5e-3 - source_load_kl_threshold: 5e-3 - source_load_cosine_threshold: 0.9995 distributed.tp_size: 2 tokenizer_name: nvidia/Llama-3.1-Nemotron-Nano-8B-v1 trust_remote_code: true hf_device_map_auto: true check_fused_qkv_keys: true - no_check_resume: true dataset.limit_dataset_samples: 500 validation_dataset.limit_dataset_samples: 500 diff --git a/examples/llm_finetune/nemotron/nemotron_nano_9b_squad.yaml b/examples/llm_finetune/nemotron/nemotron_nano_9b_squad.yaml index e40443437e..b0de206793 100644 --- a/examples/llm_finetune/nemotron/nemotron_nano_9b_squad.yaml +++ b/examples/llm_finetune/nemotron/nemotron_nano_9b_squad.yaml @@ -104,18 +104,11 @@ lr_scheduler: ci: vllm_deploy: true recipe_owner: HuiyingLi - time: "00:25:00" - known_issue_id: AM-156 + time: "00:30:00" checkpoint_robustness: - check_source_load_parity: true - kl_threshold: 5e-3 - hf_kl_threshold: 5e-3 - source_load_kl_threshold: 5e-3 - source_load_cosine_threshold: 0.9995 distributed.tp_size: 2 tokenizer_name: nvidia/NVIDIA-Nemotron-Nano-9B-v2 trust_remote_code: true hf_device_map_auto: true - no_check_resume: true dataset.limit_dataset_samples: 500 validation_dataset.limit_dataset_samples: 500 diff --git a/examples/llm_finetune/nemotron/nemotron_nano_9b_squad_peft.yaml b/examples/llm_finetune/nemotron/nemotron_nano_9b_squad_peft.yaml index aa7c34fbe3..eed490e7f2 100644 --- a/examples/llm_finetune/nemotron/nemotron_nano_9b_squad_peft.yaml +++ b/examples/llm_finetune/nemotron/nemotron_nano_9b_squad_peft.yaml @@ -111,18 +111,12 @@ lr_scheduler: ci: vllm_deploy: true recipe_owner: HuiyingLi - time: "00:25:00" + time: "00:30:00" checkpoint_robustness: - check_source_load_parity: true vllm_merge_lora: true # vLLM can't serve LoRA on fused mamba layers; deploy the merged adapter instead - kl_threshold: 5e-3 - hf_kl_threshold: 5e-3 - source_load_kl_threshold: 5e-3 - source_load_cosine_threshold: 0.9995 distributed.tp_size: 2 tokenizer_name: nvidia/NVIDIA-Nemotron-Nano-9B-v2 trust_remote_code: true hf_device_map_auto: true - no_check_resume: true dataset.limit_dataset_samples: 500 validation_dataset.limit_dataset_samples: 500 diff --git a/examples/llm_finetune/nemotron/nemotron_nano_v3_hellaswag.yaml b/examples/llm_finetune/nemotron/nemotron_nano_v3_hellaswag.yaml index 1c40a616ec..868f574d06 100644 --- a/examples/llm_finetune/nemotron/nemotron_nano_v3_hellaswag.yaml +++ b/examples/llm_finetune/nemotron/nemotron_nano_v3_hellaswag.yaml @@ -98,13 +98,11 @@ ci: vllm_deploy: true vllm_smoke_test: true recipe_owner: adil-a - time: "00:15:00" + time: "00:30:00" checkpoint_robustness: - hf_kl_threshold: 7e-2 hf_device_map_auto: true experts_implementation: grouped_mm tokenizer_name: nvidia/NVIDIA-Nemotron-3-Nano-30B-A3B-BF16 - no_check_resume: true dataset.num_samples_limit: 500 validation_dataset.num_samples_limit: 500 diff --git a/examples/llm_finetune/nemotron/nemotron_nano_v3_hellaswag_peft.yaml b/examples/llm_finetune/nemotron/nemotron_nano_v3_hellaswag_peft.yaml index 9650434ed0..907eaf9a14 100644 --- a/examples/llm_finetune/nemotron/nemotron_nano_v3_hellaswag_peft.yaml +++ b/examples/llm_finetune/nemotron/nemotron_nano_v3_hellaswag_peft.yaml @@ -115,16 +115,13 @@ ci: vllm_deploy: true vllm_smoke_test: true recipe_owner: adil-a - time: "00:15:00" - known_issue_id: AM-149 + time: "00:30:00" checkpoint_robustness: - hf_kl_threshold: 1e-1 hf_device_map_auto: true experts_implementation: grouped_mm trust_remote_code: true tokenizer_name: nvidia/NVIDIA-Nemotron-3-Nano-30B-A3B-BF16 distributed.ep_size: 8 - no_check_resume: true dataset.num_samples_limit: 500 validation_dataset.num_samples_limit: 500 diff --git a/examples/llm_finetune/nemotron/nemotron_super_v3_hellaswag.yaml b/examples/llm_finetune/nemotron/nemotron_super_v3_hellaswag.yaml index 937348b26e..7ad7df5f13 100644 --- a/examples/llm_finetune/nemotron/nemotron_super_v3_hellaswag.yaml +++ b/examples/llm_finetune/nemotron/nemotron_super_v3_hellaswag.yaml @@ -109,12 +109,10 @@ ci: time: "00:25:00" known_issue_id: AM-156 checkpoint_robustness: - hf_kl_threshold: 7e-2 experts_implementation: grouped_mm trust_remote_code: true tokenizer_name: nvidia/NVIDIA-Nemotron-3-Super-120B-A12B-BF16 hf_device_map_auto: true - no_check_resume: true step_scheduler.global_batch_size: 64 dataset.num_samples_limit: 500 validation_dataset.num_samples_limit: 500 diff --git a/examples/llm_finetune/nemotron/nemotron_super_v3_hellaswag_peft.yaml b/examples/llm_finetune/nemotron/nemotron_super_v3_hellaswag_peft.yaml index 28564c8b01..c3a45dc4fc 100644 --- a/examples/llm_finetune/nemotron/nemotron_super_v3_hellaswag_peft.yaml +++ b/examples/llm_finetune/nemotron/nemotron_super_v3_hellaswag_peft.yaml @@ -114,17 +114,19 @@ ci: vllm_deploy_time: "00:30:00" vllm_enable_expert_parallel: false # vLLM 0.19 does not support expert parallelism with runtime LoRA recipe_owner: adil-a - time: "00:25:00" + time: "00:30:00" checkpoint_robustness: - hf_kl_threshold: 7e-2 experts_implementation: grouped_mm trust_remote_code: true tokenizer_name: nvidia/NVIDIA-Nemotron-3-Super-120B-A12B-BF16 hf_device_map_auto: true + # EP8 grouped-MoE reloads have exact trainable weights and exact pre-update + # resume state, but their measured logit/loss drift exceeds the standard profile. + parity_tolerance_profile: relaxed + resume_tolerance_profile: relaxed # Vanilla Transformers intentionally omits the training-only MTP modules. # Exact-match every adapter tensor it can instantiate and allow only MTP keys. hf_adapter_ignored_key_prefix: "base_model.model.mtp." - no_check_resume: true dataset.num_samples_limit: 500 validation_dataset.num_samples_limit: 500 diff --git a/examples/llm_finetune/nemotron_flash/nemotron_flash_1b_squad.yaml b/examples/llm_finetune/nemotron_flash/nemotron_flash_1b_squad.yaml index b940674f27..847f4da5af 100755 --- a/examples/llm_finetune/nemotron_flash/nemotron_flash_1b_squad.yaml +++ b/examples/llm_finetune/nemotron_flash/nemotron_flash_1b_squad.yaml @@ -111,12 +111,10 @@ ci: # list-form-tied-keys remote-code model (the #1817 case), so this run # proves the force-HF load path still ties when the config requests it — # the mirror of the untied Nemotron-Nano-9B-v2 coverage. - check_source_load_parity: true - hf_kl_threshold: 5e-3 - source_load_kl_threshold: 5e-3 - source_load_cosine_threshold: 0.9995 tokenizer_name: nvidia/Nemotron-Flash-1B trust_remote_code: true + # Exact restored state passes; the first shared forward has small hybrid-kernel drift. + resume_tolerance_profile: relaxed # Preserve the independently calibrated hybrid-attention/Mamba/DeltaNet # reproducibility envelope. This metric is observational and does not # relax the exact-state shared-trajectory resume checks. diff --git a/examples/llm_finetune/nemotron_flash/nemotron_flash_1b_squad_peft.yaml b/examples/llm_finetune/nemotron_flash/nemotron_flash_1b_squad_peft.yaml index 9ee657e3c1..48d638d771 100755 --- a/examples/llm_finetune/nemotron_flash/nemotron_flash_1b_squad_peft.yaml +++ b/examples/llm_finetune/nemotron_flash/nemotron_flash_1b_squad_peft.yaml @@ -112,12 +112,14 @@ ci: recipe_owner: akoumpa time: "00:15:00" checkpoint_robustness: - hf_kl_threshold: 5e-3 - kl_threshold: 5e-3 # Nemotron Flash's hybrid Mamba/DeltaNet layers do not have a native TP - # plan. Keep checkpoint reload and strict resume parity on the recipe's + # plan. Keep checkpoint reload and exact resume-state gates on the recipe's # supported TP=1 topology instead of introducing row-parallel BF16 drift. distributed.tp_size: 1 + # Full and PEFT runs restore exact pre-update state but show 0.25-0.47% + # first-forward hybrid-kernel drift. Keep logit parity standard and widen + # only the shared-trajectory loss envelope. + resume_tolerance_profile: relaxed tokenizer_name: nvidia/Nemotron-Flash-1B trust_remote_code: true check_fused_qkv_keys: true diff --git a/examples/llm_finetune/phi/phi_4_squad.yaml b/examples/llm_finetune/phi/phi_4_squad.yaml index c01748d161..e5ea8200b3 100644 --- a/examples/llm_finetune/phi/phi_4_squad.yaml +++ b/examples/llm_finetune/phi/phi_4_squad.yaml @@ -108,7 +108,6 @@ ci: node_multiplier: true vllm_deploy: true checkpoint_robustness: - hf_kl_threshold: 2e-3 tokenizer_name: microsoft/phi-4 dataset.limit_dataset_samples: 500 validation_dataset.limit_dataset_samples: 500 diff --git a/examples/llm_finetune/phi/phi_4_squad_peft.yaml b/examples/llm_finetune/phi/phi_4_squad_peft.yaml index 0583f78174..e442921f81 100644 --- a/examples/llm_finetune/phi/phi_4_squad_peft.yaml +++ b/examples/llm_finetune/phi/phi_4_squad_peft.yaml @@ -106,7 +106,6 @@ ci: known_issue_id: AM-154 allow_failure: true checkpoint_robustness: - hf_kl_threshold: 1e-3 tokenizer_name: microsoft/phi-4 dataset.limit_dataset_samples: 500 validation_dataset.limit_dataset_samples: 500 diff --git a/examples/llm_finetune/qwen/qwen2_5_7b_squad.yaml b/examples/llm_finetune/qwen/qwen2_5_7b_squad.yaml index 100f768025..52d739790f 100644 --- a/examples/llm_finetune/qwen/qwen2_5_7b_squad.yaml +++ b/examples/llm_finetune/qwen/qwen2_5_7b_squad.yaml @@ -107,11 +107,9 @@ ci: known_issue_id: AM-154 allow_failure: true checkpoint_robustness: - hf_kl_threshold: 1e-1 distributed.tp_size: 2 tokenizer_name: Qwen/Qwen2.5-7B cross_tp_size: 2 - cross_tp_kl_threshold: 9e-3 training_reproducibility_loss_threshold: 5e-2 dataset.limit_dataset_samples: 500 validation_dataset.limit_dataset_samples: 500 diff --git a/examples/llm_finetune/qwen/qwen2_5_7b_squad_peft.yaml b/examples/llm_finetune/qwen/qwen2_5_7b_squad_peft.yaml index 07b8f12e30..0363afc283 100644 --- a/examples/llm_finetune/qwen/qwen2_5_7b_squad_peft.yaml +++ b/examples/llm_finetune/qwen/qwen2_5_7b_squad_peft.yaml @@ -107,7 +107,6 @@ ci: known_issue_id: AM-154 allow_failure: true checkpoint_robustness: - hf_kl_threshold: 8e-2 distributed.tp_size: 2 tokenizer_name: Qwen/Qwen2.5-7B check_fused_qkv_keys: true diff --git a/examples/llm_finetune/qwen/qwen3_moe_30b_hellaswag.yaml b/examples/llm_finetune/qwen/qwen3_moe_30b_hellaswag.yaml index ab16d31580..3216f22117 100644 --- a/examples/llm_finetune/qwen/qwen3_moe_30b_hellaswag.yaml +++ b/examples/llm_finetune/qwen/qwen3_moe_30b_hellaswag.yaml @@ -92,23 +92,12 @@ optimizer: ci: recipe_owner: hemildesai - time: "00:20:00" + time: "00:25:00" known_issue_id: AM-648 allow_failure: true checkpoint_robustness: - check_source_load_parity: true - # Keep source-reference, training, and reload forwards independent. Same-process - # release CI retained accelerator state and produced a spurious source KL of 17.52 (AMINT-238). - process_isolation: true - hf_kl_threshold: 1e-1 - # Optimized BF16 MoE execution differs broadly from vanilla HF at source load; - # current-stack CI observed max/mean KL 0.0790/0.0113 and cosine 0.9957. - source_load_kl_threshold: 1e-1 - source_load_mean_kl_threshold: 1.5e-2 - source_load_cosine_threshold: 0.995 tokenizer_name: Qwen/Qwen3-30B-A3B trust_remote_code: true hf_device_map_auto: true - no_check_resume: true dataset.num_samples_limit: 500 validation_dataset.num_samples_limit: 500 diff --git a/examples/llm_finetune/qwen/qwen3_moe_30b_lora.yaml b/examples/llm_finetune/qwen/qwen3_moe_30b_lora.yaml index 60c1094f35..d54097d870 100644 --- a/examples/llm_finetune/qwen/qwen3_moe_30b_lora.yaml +++ b/examples/llm_finetune/qwen/qwen3_moe_30b_lora.yaml @@ -102,25 +102,14 @@ optimizer: ci: recipe_owner: adil-a + # Allow the complete source, reload, and resume matrix to finish. time: "00:30:00" nodes: 1 env_vars: PYTORCH_CUDA_ALLOC_CONF: "expandable_segments:True" checkpoint_robustness: - check_source_load_parity: true - # HF+PEFT reload remains gated by exact adapter fingerprints and a successful - # forward; its live-LoRA logits are not a stable cross-runtime reference (AMINT-216). - skip_hf_logit_parity: true - # Derive isolated source, train/save, and reload phases from the standard - # check/skip options below; no resume phases are selected. - process_isolation: true - # Exact-image BF16 source parity measured max/mean KL 0.0161/0.00407 and cosine 0.99764. - source_load_kl_threshold: 3e-2 - source_load_mean_kl_threshold: 6e-3 - source_load_cosine_threshold: 0.997 tokenizer_name: Qwen/Qwen3-30B-A3B trust_remote_code: true hf_device_map_auto: true - no_check_resume: true dataset.num_samples_limit: 500 validation_dataset.num_samples_limit: 500 diff --git a/examples/llm_finetune/qwen/qwen3_moe_30b_te_deepep.yaml b/examples/llm_finetune/qwen/qwen3_moe_30b_te_deepep.yaml index e492d987d4..9a116b67db 100644 --- a/examples/llm_finetune/qwen/qwen3_moe_30b_te_deepep.yaml +++ b/examples/llm_finetune/qwen/qwen3_moe_30b_te_deepep.yaml @@ -122,24 +122,17 @@ optimizer: ci: recipe_owner: hemildesai node_multiplier: true - time: "00:20:00" + # The full 2K source/reload matrix completes inside 20 minutes, but the + # default native-resume phase needs additional setup and continuation time. + time: "00:30:00" vllm_deploy: true vllm_smoke_test: true checkpoint_robustness: - check_source_load_parity: true - # Exact-stack runs measured max/mean KL up to 0.0167/0.00437 and cosine down to 0.99750. - source_load_kl_threshold: 3e-2 - source_load_mean_kl_threshold: 6e-3 - source_load_cosine_threshold: 0.997 - # Same-process CI reached KL 0.2667, while controlled fresh-process reloads were stable at 0.0199. - hf_kl_threshold: 5e-2 - # Derive isolated source, train/save, and reload phases from the standard - # check/skip options below; no resume phases are selected. - process_isolation: true + # Long-context mean, p95, and cosine metrics fit the shared standard profile. + # Max KL remains diagnostic rather than using the former per-model max-only overrides. tokenizer_name: Qwen/Qwen3-30B-A3B # Match the other Qwen3-MoE references and avoid the real-device fallback # selected by trust_remote_code when no HF device map is provided. trust_remote_code: true hf_device_map_auto: true - no_check_resume: true dataset.num_samples_limit: 500 diff --git a/examples/retrieval/bi_encoder/nemotron_vl_1b/nemotron_vl_1b_example.yaml b/examples/retrieval/bi_encoder/nemotron_vl_1b/nemotron_vl_1b_example.yaml index db87aabe7b..a99d8adfb6 100644 --- a/examples/retrieval/bi_encoder/nemotron_vl_1b/nemotron_vl_1b_example.yaml +++ b/examples/retrieval/bi_encoder/nemotron_vl_1b/nemotron_vl_1b_example.yaml @@ -129,10 +129,7 @@ ci: lr_scheduler.lr_warmup_steps: 1 step_scheduler.global_batch_size: 8 step_scheduler.local_batch_size: 1 - cosine_threshold: 0.999 - hf_cosine_threshold: 0.999 - check_hf_reload: true - check_resume: true + parity_tolerance_profile: standard # Shared-trajectory CI observed later-step loss deltas of 0.0137 and 0.0215. # Keep the exact first resumed step strict, but allow BF16 distributed # contrastive-loss drift to accumulate across the following optimizer steps. diff --git a/examples/vlm_finetune/gemma4/gemma4_26b_a4b_moe.yaml b/examples/vlm_finetune/gemma4/gemma4_26b_a4b_moe.yaml index 858b682d85..bab78d17a4 100644 --- a/examples/vlm_finetune/gemma4/gemma4_26b_a4b_moe.yaml +++ b/examples/vlm_finetune/gemma4/gemma4_26b_a4b_moe.yaml @@ -126,24 +126,10 @@ ci: # The shared robustness batch of 2 exceeds this 26B MoE recipe's memory # envelope; 1 matches the recipe's validated per-rank image batch. step_scheduler.local_batch_size: 1 - # Source-load logit parity is not a checkpoint-integrity signal for this - # recipe. The same commit gives max KL 0.28 locally and 4.xx on nemo-ci with - # the same container, byte-identical weights, config, tokenizer and token - # IDs, so the value tracks the hardware's reduction order rather than the - # checkpoint: top-k MoE routing amplifies normal BF16 implementation-order - # differences through 30 layers. Everything else stays gated -- checkpoint - # load, adapter key coverage, the AutoModel reload, the HF reload forward - # and resume. - check_source_load_parity: false - # Keep training and reload forwards independent; same-process CI retains - # accelerator state and inflates parity numbers. - process_isolation: true - check_resume: true hf_device_map_auto: true tokenizer_name: google/gemma-4-26B-A4B-it - # Attention, dense MLP and every norm are bit-exact against HF; the residual - # is the routed-expert path, where AutoModel dispatches through DeepEP and - # HF uses grouped_mm with an fp32 top-k reduction. That drift is inherent to - # the two dispatch routes, so the dense default (2e-2) does not apply. - # Measured on nemo-ci: HF-reload max KL 0.409 (main: 0.686). - hf_kl_threshold: 6e-1 + # At 2K tokens the HF model is not self-repeatable, while AutoModel is bitwise + # repeatable. Keep both HF forwards and cross-framework metrics as diagnostics, + # but do not use the unstable reference as a blocking numerical oracle. + skip_source_load_logit_parity: true + skip_hf_reload_logit_parity: true diff --git a/examples/vlm_finetune/gemma4/gemma4_2b.yaml b/examples/vlm_finetune/gemma4/gemma4_2b.yaml index 7ce7547657..35bacea221 100644 --- a/examples/vlm_finetune/gemma4/gemma4_2b.yaml +++ b/examples/vlm_finetune/gemma4/gemma4_2b.yaml @@ -107,7 +107,4 @@ ci: nodes: 1 time: "00:20:00" checkpoint_robustness: - check_source_load_parity: true - hf_kl_threshold: 2e-2 - source_load_cosine_threshold: 0.9995 tokenizer_name: google/gemma-4-E2B-it diff --git a/examples/vlm_finetune/mistral/ministral3_3b_medpix.yaml b/examples/vlm_finetune/mistral/ministral3_3b_medpix.yaml index ae92352081..b3afa0480c 100644 --- a/examples/vlm_finetune/mistral/ministral3_3b_medpix.yaml +++ b/examples/vlm_finetune/mistral/ministral3_3b_medpix.yaml @@ -106,8 +106,4 @@ ci: nodes: 1 time: "00:15:00" checkpoint_robustness: - check_source_load_parity: true - # Source max/mean KL are exactly zero in scoped CI; BF16 logits retain a - # bounded cosine of 0.9998064 after the FP8 checkpoint is dequantized. - source_load_cosine_threshold: 0.9995 tokenizer_name: mistralai/Ministral-3-3B-Reasoning-2512 diff --git a/examples/vlm_finetune/mistral4/mistral4_medpix.yaml b/examples/vlm_finetune/mistral4/mistral4_medpix.yaml index f8033c25e4..e1d3f574f9 100644 --- a/examples/vlm_finetune/mistral4/mistral4_medpix.yaml +++ b/examples/vlm_finetune/mistral4/mistral4_medpix.yaml @@ -124,15 +124,10 @@ ci: checkpoint_robustness: # PP=4 with pp_microbatch_size=1 requires at least four pipeline microbatches. step_scheduler.local_batch_size: 4 - check_source_load_parity: true - # Fresh PP4/EP8 DCP and consolidated reloads are bit-identical to each other, - # but both differ from the still-warm in-memory runtime. Bound that distributed - # BF16 execution delta separately from the native-vs-HF implementation delta. - kl_threshold: 5e-2 - source_load_kl_threshold: 1e-2 - source_load_mean_kl_threshold: 2e-3 - source_load_cosine_threshold: 0.999 - hf_kl_threshold: 5e-2 + # The original source and AutoModel reload pass standard. Only the PP4/EP8 + # post-training HF reload has measured long-context BF16 drift. + parity_tolerance_profile_overrides: + hf_reload: relaxed hf_device_map_auto: true # Transformers' load-time dequantization leaves Mistral4 expert weights in # FP8, while native forward requires an optional kernel absent from CI. diff --git a/examples/vlm_finetune/qwen3/qwen3_vl_moe_30b_te_deepep.yaml b/examples/vlm_finetune/qwen3/qwen3_vl_moe_30b_te_deepep.yaml index c9ad7bdca6..ca74a917eb 100644 --- a/examples/vlm_finetune/qwen3/qwen3_vl_moe_30b_te_deepep.yaml +++ b/examples/vlm_finetune/qwen3/qwen3_vl_moe_30b_te_deepep.yaml @@ -119,13 +119,7 @@ ci: nodes: 1 time: "00:30:00" checkpoint_robustness: - check_source_load_parity: true - # Current-stack BF16 CI observed source max/mean KL up to 0.0293/0.0052 and HF reload KL 0.0201. - hf_kl_threshold: 2.5e-2 hf_device_map_auto: true # Independently calibrated normal-run envelope; it does not relax resume correctness. training_reproducibility_loss_threshold: 2e-2 - source_load_kl_threshold: 4e-2 - source_load_mean_kl_threshold: 7e-3 - source_load_cosine_threshold: 0.9985 tokenizer_name: Qwen/Qwen3-VL-30B-A3B-Instruct diff --git a/examples/vlm_finetune/qwen3_5_moe/qwen3_5_35b.yaml b/examples/vlm_finetune/qwen3_5_moe/qwen3_5_35b.yaml index 1e6115c50a..feea3dbff8 100644 --- a/examples/vlm_finetune/qwen3_5_moe/qwen3_5_35b.yaml +++ b/examples/vlm_finetune/qwen3_5_moe/qwen3_5_35b.yaml @@ -119,19 +119,10 @@ ci: time: "00:30:00" # Keep failures blocking; source/reload defects are tracked by AM-711/AM-717. checkpoint_robustness: - check_source_load_parity: true experts_implementation: grouped_mm hf_device_map_auto: true - # Qwen3.5 BF16 parity has known model/backend drift: AutoModel applies routing - # weights before down_proj while HF applies them after, and DeepEP combine uses - # different reduction numerics from HF grouped_mm's FP32 reshape-and-sum. - # Focused parity runs measured max/mean source KL of ~0.068/~0.0086 (AM-711). - hf_kl_threshold: 1e-1 loss_fn._target_: nemo_automodel.components.loss.chunked_ce.ChunkedCrossEntropy model.backend.experts: torch_mm - source_load_cosine_threshold: 0.9985 - source_load_kl_threshold: 1e-1 - source_load_mean_kl_threshold: 1e-2 step_scheduler.global_batch_size: 16 step_scheduler.local_batch_size: 1 tokenizer_name: Qwen/Qwen3.5-35B-A3B diff --git a/examples/vlm_finetune/qwen3_8/qwen3_8_27b.yaml b/examples/vlm_finetune/qwen3_8/qwen3_8_27b.yaml index a656bdc45d..18a6b95fe1 100644 --- a/examples/vlm_finetune/qwen3_8/qwen3_8_27b.yaml +++ b/examples/vlm_finetune/qwen3_8/qwen3_8_27b.yaml @@ -125,6 +125,13 @@ wandb: ci: recipe_owner: akoumpa nodes: 1 - time: "00:15:00" + time: "00:40:00" env_vars: MAX_STEPS: 10 + checkpoint_robustness: + # The shared robustness batch of two OOMs this 27B full-SFT recipe; retain + # the recipe's validated per-rank batch of one. + step_scheduler.global_batch_size: 8 + step_scheduler.local_batch_size: 1 + tokenizer_name: Qwen/Qwen3.8-27B + hf_device_map_auto: true diff --git a/examples/vlm_finetune/stepfun/step3p7_medpix_200b_lora_pp8ep8_8node.yaml b/examples/vlm_finetune/stepfun/step3p7_medpix_200b_lora_pp8ep8_8node.yaml index effa21485e..af8a2adaaa 100644 --- a/examples/vlm_finetune/stepfun/step3p7_medpix_200b_lora_pp8ep8_8node.yaml +++ b/examples/vlm_finetune/stepfun/step3p7_medpix_200b_lora_pp8ep8_8node.yaml @@ -148,7 +148,8 @@ wandb: ci: # pp_size(8) * ep_size(8) = 64 GPUs => 8 nodes (8 H100/node). nodes: 8 - time: "01:00:00" + # The complete source, reload, and resume matrix repeatedly constructs this 200B model. + time: "02:00:00" checkpoint_robustness: # PP=8 with pp_microbatch_size=1 requires at least eight pipeline microbatches. step_scheduler.global_batch_size: 64 @@ -156,24 +157,31 @@ ci: # One optimizer step materializes Adam state; AutoModel setup then restores it from DCP. step_scheduler.max_steps: 1 step_scheduler.ckpt_every_steps: 1 - # Match a real restart and let each large trainer exit before the next one starts. - process_isolation: true - # AutoModel reload restores model and optimizer state; skip the separate resume/loss-continuity check. - no_check_resume: true - # Cross-process PP+EP BF16 logits are not a stable checkpoint-integrity signal (AMINT-238). - # Keep the exact per-rank LoRA fingerprints as the blocking AutoModel reload gate. - skip_automodel_logit_parity: true - # HF reload still exact-matches every supported adapter tensor and runs a forward smoke. - # Cross-implementation KL is not a checkpoint-integrity gate: Step-3.7's top-k MoE routing - # amplifies normal BF16 implementation-order differences into different expert choices. - skip_hf_logit_parity: true + # Process isolation is the harness default, so each large trainer exits before the next one starts. # Vanilla HF needs remote model code and shards the 403 GB base over rank 0's local GPU node. trust_remote_code: true hf_device_map_auto: true # Keep overflow in RAM and leave GPU headroom for PEFT's adapter re-dispatch. hf_device_map_max_memory_gib: 55 hf_device_map_cpu_max_memory_gib: 128 + # CPU-offloaded rank-0 HF reload can exceed the shared 30-minute wait on a busy cluster. + hf_reload_timeout_seconds: 3600 # The published HF model intentionally omits training-only MTP layers 45-47. Exact-match every # adapter tensor HF can instantiate, while allowing only the saved MTP adapter namespace. hf_adapter_ignored_key_prefix: "base_model.model.mtp." + # HF is not self-repeatable, so only the two HF-based numerical gates are + # informational. AutoModel reload remains blocking. Across three exact-state + # scoped runs its mean KL varied from 0.00447 to 0.03215 while within-process + # self-repeat stayed bitwise zero. Keep the shared relaxed profile and override + # only the two Phase 2 gates exceeded by this model; relaxed p95 remains active. + parity_tolerance_profile: relaxed + parity_threshold_overrides: + automodel_reload: + mean_kl: 0.04 + cosine_similarity: 0.99 + skip_source_load_logit_parity: true + skip_hf_reload_logit_parity: true + # Restored model/optimizer state and the first forward are exact. The + # process-sensitive MoE path diverges only after computing the next update. + resume_tolerance_profile: relaxed tokenizer_name: stepfun-ai/Step-3.7-Flash diff --git a/nemo_automodel/components/models/minimax_m2/model.py b/nemo_automodel/components/models/minimax_m2/model.py index f80a18ea60..15123f97ec 100644 --- a/nemo_automodel/components/models/minimax_m2/model.py +++ b/nemo_automodel/components/models/minimax_m2/model.py @@ -12,7 +12,7 @@ # See the License for the specific language governing permissions and # limitations under the License. -from dataclasses import dataclass +from dataclasses import dataclass, replace from typing import Any, Union import torch @@ -98,6 +98,14 @@ def __init__( moe_overrides: dict | None = None, ): super().__init__() + # Released MiniMax-M2 checkpoints store the router gate weight in fp32, + # and the HF reference projects with hidden_states.to(weight.dtype), so + # the checkpoint-faithful router runs an fp32 projection, fp32 scoring, + # and fp32 selected weights. Keep that default while preserving an + # explicit backend override (see AMINT-286; ERNIE follows the same + # pattern for its fp32 router). + if backend.gate_precision is None: + backend = replace(backend, gate_precision=torch.float32) self.backend = backend self.config = config if moe_config is not None and moe_overrides is not None: @@ -133,6 +141,14 @@ def __init__( expert_activation="swiglu", softmax_before_topk=(score_func == "softmax"), force_e_score_correction_bias=True, + # The HF reference returns selected weights in the fp32 router + # logits dtype; keep them fp32 through the expert combine. + router_weights_fp32=True, + # The checkpoint stores the gate weight in fp32; allocate it fp32 so + # every construction path (including meta-device init before FSDP + # sharding) keeps the gate's dtype group uniform with its fp32 + # correction-bias buffer. + gate_dtype=torch.float32, dtype=model_dtype, ) if moe_overrides: @@ -229,7 +245,7 @@ def init_weights(self, buffer_device: torch.device | None = None) -> None: class MiniMaxM2ForCausalLM(HFCheckpointingMixin, nn.Module, MoEFSDPSyncMixin): tie_word_embeddings_support: TieSupport = TieSupport.UNTIED_ONLY - _keep_in_fp32_modules_strict = ["mlp.gate.e_score_correction_bias"] + _keep_in_fp32_modules_strict = ["mlp.gate.weight", "mlp.gate.e_score_correction_bias"] @dataclass(frozen=True) class ModelCapabilities: diff --git a/nemo_automodel/components/models/nemotron_v3/state_dict_adapter.py b/nemo_automodel/components/models/nemotron_v3/state_dict_adapter.py index 68f6dd9767..4a71be7ef9 100644 --- a/nemo_automodel/components/models/nemotron_v3/state_dict_adapter.py +++ b/nemo_automodel/components/models/nemotron_v3/state_dict_adapter.py @@ -105,7 +105,7 @@ def __init__( self.moe_config = moe_config self.backend = backend self.dtype = dtype - self._uses_model_prefix = True + self._uses_model_prefix = False # Mapping for expert weights (HF split → internal merged) self.from_hf_map = { @@ -115,8 +115,8 @@ def __init__( @property def _hf_prefix(self) -> str: - """NemotronV3 HF format uses 'backbone.' prefix.""" - return "backbone." + """Return the source checkpoint's public Nemotron-H model prefix.""" + return "model." if self._uses_model_prefix else "backbone." @property def _expert_path_segment(self) -> str: @@ -128,22 +128,38 @@ def _v5_peft_target_parameters(self) -> tuple[str, ...]: """Nemotron V3 exposes fused non-gated expert parameters in Transformers v5.""" return ("mixer.experts.up_proj", "mixer.experts.down_proj") - @staticmethod - def _native_key_to_hf(key: str) -> str: + def _native_key_to_hf(self, key: str) -> str: """Normalize a native Nemotron V3 key to its public HF namespace.""" key = _strip_mamba_fp32_holder_key(key) - key = re.sub(r"^model\.", "backbone.", key) - key = re.sub(r"^backbone\.norm\.weight$", "backbone.norm_f.weight", key) - key = re.sub(r"^backbone\.embed_tokens\.weight$", "backbone.embeddings.weight", key) + key = re.sub( + r"^(?Pbase_model\.model\.)?model\.", + lambda match: f"{match.group('outer') or ''}{self._hf_prefix}", + key, + ) + hf_root = re.escape(self._hf_prefix.rstrip(".")) + key = re.sub(rf"^{hf_root}\.norm\.weight$", f"{self._hf_prefix}norm_f.weight", key) + key = re.sub(rf"^{hf_root}\.embed_tokens\.weight$", f"{self._hf_prefix}embeddings.weight", key) return key - @staticmethod - def _hf_key_to_native(key: str) -> str: + def map_peft_target_module_to_hf(self, module_name: str) -> str: + """Map native PEFT target modules to the public Nemotron-H namespace.""" + return self._native_key_to_hf(module_name) + + def _hf_key_to_native(self, key: str) -> str: """Normalize a public HF Nemotron V3 key to its native namespace.""" - key = re.sub(r"^((?:base_model\.model\.)?backbone)\.norm_f\.weight$", r"\1.norm.weight", key) - key = re.sub(r"^((?:base_model\.model\.)?backbone)\.embeddings\.weight$", r"\1.embed_tokens.weight", key) + hf_root = re.escape(self._hf_prefix.rstrip(".")) + key = re.sub( + rf"^(?Pbase_model\.model\.)?{hf_root}\.norm_f\.weight$", + lambda match: f"{match.group('outer') or ''}model.norm.weight", + key, + ) + key = re.sub( + rf"^(?Pbase_model\.model\.)?{hf_root}\.embeddings\.weight$", + lambda match: f"{match.group('outer') or ''}model.embed_tokens.weight", + key, + ) return re.sub( - r"^(?Pbase_model\.model\.)?backbone\.", + rf"^(?Pbase_model\.model\.)?{hf_root}\.", lambda match: f"{match.group('outer') or ''}model.", key, ) @@ -216,11 +232,16 @@ def from_hf( else: backbone_state_dict[key] = value - # Detect if HF checkpoint uses 'backbone' or 'model' prefix. Only - # look at backbone keys; MTP keys never carry a backbone/model prefix. + # Detect whether the source checkpoint uses the remote-code ``backbone`` + # namespace or Transformers v5's native ``model`` namespace. MTP keys + # never carry either prefix. for key in backbone_state_dict.keys(): - if ".mixer.experts." in key: - self._uses_model_prefix = not key.startswith("backbone.") + bare_key = key.removeprefix("base_model.model.") + if bare_key.startswith("backbone."): + self._uses_model_prefix = False + break + if bare_key.startswith("model."): + self._uses_model_prefix = True break # First, rename backbone → model and norm_f → norm diff --git a/nemo_automodel/components/moe/config.py b/nemo_automodel/components/moe/config.py index 6255cb8150..e1bf85b235 100644 --- a/nemo_automodel/components/moe/config.py +++ b/nemo_automodel/components/moe/config.py @@ -56,6 +56,11 @@ class MoEConfig: router_weights_fp32: bool = False router_weight_uses_score_correction_bias: bool = False dtype: str | torch.dtype = torch.bfloat16 + # Storage dtype for the router gate parameters. None inherits ``dtype``. + # Models whose checkpoints store the gate in fp32 (e.g. MiniMax-M2) set + # this so the gate is fp32 from allocation on every construction path, + # keeping FSDP dtype groups uniform with the fp32 correction-bias buffer. + gate_dtype: str | torch.dtype | None = None shared_expert_gate: bool = False shared_expert_inter_dim: int | None = None shared_expert_activation: str = "swiglu" # Activation for shared experts ("swiglu" or "relu2") @@ -74,6 +79,8 @@ def expert_dim(self) -> int: def __post_init__(self): if isinstance(self.dtype, str): self.dtype = dtype_from_str(self.dtype, default=torch.bfloat16) + if isinstance(self.gate_dtype, str): + self.gate_dtype = dtype_from_str(self.gate_dtype, default=torch.bfloat16) @dataclass diff --git a/nemo_automodel/components/moe/layers.py b/nemo_automodel/components/moe/layers.py index 30e3a39773..3020be135e 100644 --- a/nemo_automodel/components/moe/layers.py +++ b/nemo_automodel/components/moe/layers.py @@ -270,12 +270,13 @@ def __init__( if self.bias_update_factor > 0: assert self.train_gate, "Require train_gate to be set to True to apply the bias update" + gate_dtype = config.gate_dtype or config.dtype self.weight = nn.Parameter( - torch.empty(config.n_routed_experts, config.dim, dtype=config.dtype), requires_grad=self.train_gate + torch.empty(config.n_routed_experts, config.dim, dtype=gate_dtype), requires_grad=self.train_gate ) if config.router_bias: self.bias = nn.Parameter( - torch.empty(config.n_routed_experts, dtype=config.dtype), requires_grad=self.train_gate + torch.empty(config.n_routed_experts, dtype=gate_dtype), requires_grad=self.train_gate ) else: self.bias = None @@ -458,7 +459,11 @@ def _route_scores(self, scores: torch.Tensor) -> tuple[torch.Tensor, torch.Tenso indices = replay_selection(self.router_replay, indices) weights = original_scores.gather(1, indices) elif self.score_func == "sigmoid_with_bias": - scores = scores.sigmoid() + # Score in fp32 like the softmax path: HF sigmoid-router references + # compute sigmoid(logits.float()), and bf16 sigmoid quantizes scores + # at ~2e-3 — enough to flip knife-edge e_score_correction_bias + # selections (AMINT-286). + scores = torch.sigmoid(scores.to(dtype=self.gate_precision or torch.float32)) original_scores = scores scores_for_choice = scores @@ -479,7 +484,8 @@ def _route_scores(self, scores: torch.Tensor) -> tuple[torch.Tensor, torch.Tenso indices = replay_selection(self.router_replay, indices) weights = original_scores.gather(1, indices) else: - scores = scores.sigmoid() + # Score in fp32 like the softmax path (see sigmoid_with_bias above). + scores = torch.sigmoid(scores.to(dtype=self.gate_precision or torch.float32)) original_scores = scores # Add correction bias to balance tokens across gates. diff --git a/tests/ci_tests/README.md b/tests/ci_tests/README.md index 186719c56e..4948c5adcf 100644 --- a/tests/ci_tests/README.md +++ b/tests/ci_tests/README.md @@ -61,25 +61,35 @@ ci: vllm_deploy: true # Optional. Enable vLLM deployment test vllm_deploy_time: "00:30:00" # Optional. Override the vLLM deploy SLURM wall time (defaults to 00:10:00) checkpoint_robustness: # Optional. Enable robustness testing - hf_kl_threshold: 1e-3 tokenizer_name: org/model - check_source_load_parity: true # Optional. Compare raw HF source load vs constructed trainer before training + parity_sequence_length: 2048 # Optional. Full-logit parity prompt length (default: 2048; 1K-4K recommended) + parity_tolerance_profile: standard # Optional: strict, standard (default), or relaxed hf_device_map_auto: true # Optional. Use for large HF reference loads that do not fit on one GPU - no_check_resume: true # Skip phase 6 (training resumption) + # skip_resume: true # Exceptional: skip native-checkpoint resume (Phase 4) # See checkpoint robustness section for all options ``` ## Checkpoint Robustness -When `checkpoint_robustness` is present, the robustness test runs after the finetune under the same SLURM allocation. It trains for 5 steps, saves a checkpoint, then validates through: +When `checkpoint_robustness` is present, the robustness test runs after the finetune under the same SLURM allocation. +LLM and VLM tests run each lifecycle phase in a fresh process by default so the test models a real restart and does not +depend on Python object teardown. `process_isolation: false` retains the old single-process path as a compatibility +fallback. -0. **Source-load parity** (optional) -- With `check_source_load_parity: true`, capture logits from the raw HF source load, release the HF model, construct a parity-only trainer model, compare the constructed pre-training model against those HF logits, then release it so training starts from a fresh trainer -1. **Reference logits** -- Capture logits before teardown -2. **AutoModel reload** -- Reload from consolidated checkpoint, verify KL = 0 -3. **HF reload** -- Load into vanilla `transformers`/`peft`, verify KL below `hf_kl_threshold` -4. **Cross-TP** (optional) -- Reload with different `tp_size` -5. **Training resumption** (on by default) -- Continue the checkpoint-producing trajectory, restore the exact boundary - checkpoint in a fresh trainer, and compare identical post-boundary batches and losses +The public phase model is deliberately numbered 0 through 5. The two isolated jobs that produce and consume the +source reference are implementation details of Phase 0, not separate phases. + +| Phase | Default | Operation | Blocking oracle | +|-------|---------|-----------|-----------------| +| 0. Source parity | Yes | Compare the original vanilla-HF source checkpoint with a freshly constructed AutoModel before training. | Full-logit parity plus tied-input/output-embedding alias checks. | +| 1. Train, save, reference | Yes | Train for the configured short trajectory, save the checkpoint, and capture finite reference logits from the trained model. | Training and checkpoint publication complete; reference logits are finite. When Phase 4 is enabled, boundary state and continuation artifacts are also captured. | +| 2. AutoModel model reload | Yes | Reload the saved model payload through AutoModel. Dense models use the exported HF-format consolidated weights. PEFT models restore the AutoModel checkpoint payload. | Full-logit parity. PEFT additionally requires exact trainable-adapter fingerprints. | +| 3. Vanilla-HF model reload | Yes | Load the same exported dense weights with `transformers`, or the exported adapter with `peft`, and run a forward pass. | Full-logit parity. PEFT additionally requires exact saved-adapter tensor fingerprints. | +| 4. Native training resume | Yes | Restore the native distributed checkpoint at the Phase 1 boundary, including model, optimizer, scheduler, RNG, and data state, then replay identical batches. | Exact restored state and pre-update fingerprints, followed by the configured shared-trajectory loss envelope. This phase does not use KL. | +| 5. Cross-TP reload | No | Reload the exported dense weights with `cross_tp_size`. | Full-logit parity against the Phase 1 reference. | + +Phases 0–4 are the core lifecycle and are enabled by default. Phase 5 is an optional topology-portability test enabled +by setting `cross_tp_size`. A phase should be skipped only for a documented incompatibility; see the skip controls below. LLM recipes use the causal-LM harness, while `examples/vlm_finetune/` recipes use the VLM finetune recipe and `AutoModelForImageTextToText`. VLM parity currently exercises the language path with text-only `input_ids`; real-image @@ -92,13 +102,16 @@ it requires exact rank-local model parameters and optimizer tensors. Persistent plus post-update model/optimizer fingerprints, are recorded diagnostically so a numerical divergence can be localized. Exact loss deltas and diagnostic comparisons are written for successful and failed runs. -Select a shared scale-aware loss envelope with `resume_tolerance_profile`. Each stage allows +After the native checkpoint has been restored, Phase 4 disables further checkpoint writes: the continuation is an +oracle for restored state and training trajectory, and its final checkpoint is not consumed by any later phase. + +Select the Phase 4 shared scale-aware loss envelope with `resume_tolerance_profile`. Each stage allows `atol + rtol * max(abs(uninterrupted_loss), abs(resumed_loss))`: `strict` uses `1e-6 + 0%` for both stages; `standard` (default) uses `1e-5 + 0.2%` for the first step and `5e-3 + 0.2%` later; `relaxed` uses `1e-4 + 0.75%` first and `1e-2 + 0.75%` later. Prefer profiles over model-specific calibration, and use `relaxed` only for demonstrated distributed or low-precision drift after the exact state gates pass. `resume_first_loss_threshold` and `resume_loss_threshold` remain authoritative absolute-only overrides for exceptional cases. Use -`no_check_resume: true` only for an explicitly documented restore blocker. +`skip_resume: true` only for an explicitly documented restore blocker. CI also reuses the normal finetune that already precedes checkpoint robustness as a separate, non-blocking training- reproducibility metric; it does not launch another baseline. Normal finetune and checkpoint Phase 1 record per-rank @@ -110,17 +123,128 @@ emits a prominent `ALERT` and saves a machine-readable `report.json` in the repr an opportunistic diagnostic rather than required coverage: phase-specific overrides may make the two existing runs incomparable, while the shared-trajectory resume check remains the blocking reproducibility oracle. -Use source-load parity for recipes where the initial HF checkpoint load is itself part of the contract, especially -remote-code, force-HF, custom model, or tied/untied `lm_head` paths. The raw HF reference model is loaded only long -enough to capture logits and is released before the trainer model is constructed. +Phase 0 makes the initial HF checkpoint load part of the default contract. The raw HF reference model is loaded only +long enough to capture logits and is released before the trainer model is constructed. This catches remote-code, +force-HF, custom-model, and tied/untied `lm_head` regressions before training can obscure them. + +The AutoModel side always keeps the recipe's configured attention backend. The independent vanilla-HF reference uses +that backend when the pinned Transformers model declares support for it; otherwise it uses `eager` and logs an +attention-compatibility fallback. This preserves a working HF reference instead of making a recipe backend that HF +cannot execute look like a checkpoint failure. For large reference models, set `hf_device_map_auto: true` so HF can use `device_map="auto"` instead of placing the -whole reference load on one rank's GPU. This is intentionally opt-in rather than the default: small models should keep -the simpler single-device HF load for deterministic behavior, while large models (for example 9B+ or configs that -already require multi-GPU HF reloads) should enable it to avoid rank-0 OOM. Tune `source_load_kl_threshold` and -`source_load_mean_kl_threshold` only when backend or dtype differences are expected. The first threshold bounds the -worst token, while the stricter mean threshold prevents broad drift; Phase 0 also reports p95 KL for diagnosis. -`source_load_cosine_threshold` remains an independent full-logit check. +whole reference load on one rank's GPU. This remains opt-in: small models keep the simpler single-device HF load, +while large models (for example 9B+ or configs that already require multi-GPU HF reloads) should enable it to avoid +rank-0 OOM. + +The other ranks wait up to 1,800 seconds for the rank-0-only vanilla-HF reload by default. Set +`hf_reload_timeout_seconds` only when a documented large or CPU-offloaded reference can legitimately take longer; +this changes the synchronization timeout, not any numerical gate. + +### Full-Logit Metrics and Profiles + +Phases 0, 2, 3, and 5 compare every vocabulary logit for every prompt token. A version-controlled snapshot of the +long-form finetuning guide is tokenized with each model's tokenizer, then truncated to `parity_sequence_length` tokens +(default 2048; 1K-4K is the recommended range). The snapshot contains more than 6,000 words and is protected by a +checked SHA-256 digest, so a 4K test uses stable, unique document content rather than a repeated short prompt. The +harness fails with an actionable error instead of repeating content if a requested length exceeds the tokenized +document. Pipeline-parallel runs resize their stage activation buffers to the configured parity length; reduce the +length only when a model has a documented memory limit. + +Every comparison reports mean, p95, and max per-token `KL(reference || candidate)`; whole-tensor cosine similarity; +and mean/max absolute logit difference. The full record is printed as `CHECKPOINT_PARITY_METRICS ` and saved +under `/.checkpoint_robustness/parity_metrics/`. Named profiles gate mean KL, p95 KL, and cosine +similarity. Max KL and absolute logit differences remain diagnostics, allowing a single extreme token to remain +visible without making the default gate as unstable as max KL. + +Each vanilla-HF reference is forwarded twice through the same loaded model. The resulting `hf_source_self_repeat` +or `hf_export_self_repeat` record is informational and distinguishes cross-framework drift from an unstable reference. +Phase 1 likewise always emits an informational `automodel_reference_self_repeat` record. Phase 2 emits +`automodel_reload_self_repeat` for the `relaxed` profile, an informational reload gate, or a reload +comparison that exceeds its active thresholds. This keeps the dense passing path to one additional AutoModel forward +while capturing both sides of the repeatability diagnosis for sensitive or failing configurations. + +| Self-repeat record | What it measures | +|--------------------|------------------| +| `hf_source_self_repeat` | Repeatability of the original loaded HF checkpoint. | +| `automodel_reference_self_repeat` | Repeatability of the trained Phase 1 AutoModel reference. | +| `automodel_reload_self_repeat` | Repeatability of the independently reloaded Phase 2 AutoModel. | +| `hf_export_self_repeat` | Repeatability of the exported checkpoint reloaded in vanilla HF. | + +All self-repeat records have `enforced: false`: they cannot fail the job, select a profile, or change an active +threshold. Use their logged JSON metrics for offline diagnosis and profile calibration. If a primary reload comparison +is large while both relevant self-repeat comparisons are small, investigate checkpoint/load correctness. If a +self-repeat comparison is already large, the model or reference execution is itself numerically variable and the +primary comparison includes that variability. + +| Profile | Same implementation mean / p95 / cosine | Cross-framework mean / p95 / cosine | Cross-topology mean / p95 / cosine | +|---------|-----------------------------------------|---------------------------------------|-------------------------------------| +| `strict` | `1e-7` / `1e-6` / `0.999999` | `1e-4` / `1e-3` / `0.9999` | `1e-6` / `1e-5` / `0.99999` | +| `standard` (default) | `3e-3` / `1.2e-2` / `0.999` | `6e-3` / `3e-2` / `0.998` | `6e-3` / `3e-2` / `0.998` | +| `relaxed` | `2e-2` / `5e-2` / `0.995` | `2.5e-2` / `1e-1` / `0.99` | `2e-2` / `5e-2` / `0.995` | +Use `strict` for deterministic same-kernel paths and `standard` for dense models and numerically stable MoE paths. +Reserve `relaxed` for demonstrated discontinuous distributed behavior, normally expert-parallel MoE routing, and +document the evidence in the recipe. Model size, TP/PP, or MoE status alone does not justify it. A dense model that +exceeds `standard` should be investigated. `parity_tolerance_profile` is the fallback for every comparison. If only +one verified comparison needs another shared profile after exact checkpoint-state and within-process repeatability +checks pass, select it with `parity_tolerance_profile_overrides`; every unspecified comparison retains the fallback. +Use `parity_threshold_overrides` only when that comparison also exceeds the closest shared profile, and override only +the necessary gate. Every unspecified metric remains inherited from the active comparison profile. +The comparison class is selected by the harness. For every profile, a cross-topology comparison is never stricter +than the same-implementation comparison because changing topology adds a numerical variation source. + +```yaml +parity_tolerance_profile: standard +parity_tolerance_profile_overrides: + hf_reload: relaxed +``` + +The global `standard` line above is optional because it is the default. Supported comparison names are `source_load`, +`automodel_reload`, `hf_reload`, and `cross_tp`. Phase 4 uses the separate `resume_tolerance_profile` because it gates +the restored training loss trajectory rather than full-logit metrics. + +For a measured exception that exceeds even its selected comparison profile: + +```yaml +parity_tolerance_profile: relaxed +parity_threshold_overrides: + automodel_reload: + mean_kl: 0.04 + cosine_similarity: 0.99 +``` + +Supported metric names are `mean_kl`, `p95_kl`, and `cosine_similarity`. Numeric overrides are exceptional calibration +escape hatches, not additional profiles. Max KL remains diagnostic and cannot be overridden. + +Legacy positive `check_*` controls, generic numeric cosine fields, and max-KL threshold fields are no longer accepted. +All live recipes use default-on phases, semantic `skip_*` controls, and named profiles. The optional structured +profile and numeric override mappings remain available for measured one-model exceptions. + +Retrieval checkpoint robustness uses the same phase contract for Phases 1–4. Because a biencoder produces embeddings +rather than language-model logits, its Phase 2 AutoModel reload gates the selected profile's same-implementation +cosine threshold, and its Phase 3 vanilla-HF reload gates the cross-framework cosine threshold. KL gates do not apply +to embedding outputs. Retrieval profile and numeric override mappings therefore support only `automodel_reload` and +`hf_reload`; retrieval currently has no Phase 0 source-load or Phase 5 cross-TP comparison. + +### Phase Controls + +| Field | Effect | +|-------|--------| +| `skip_source_load_parity: true` | Skip all of Phase 0. | +| `skip_source_load_logit_parity: true` | Keep the Phase 0 HF load/forward smoke and report full metrics, but make source-vs-AutoModel logit parity informational. | +| `skip_automodel_reload_logit_parity: true` | Keep the Phase 2 reload, forward smoke, and PEFT fingerprints, but make its logit metrics informational. | +| `skip_hf_reload: true` | Skip all of Phase 3, including its load and forward smoke. | +| `skip_hf_reload_logit_parity: true` | Keep the Phase 3 load, forward smoke, and PEFT fingerprints, but make its logit metrics informational. | +| `skip_resume: true` | Skip Phase 4. | +| `cross_tp_size: N` | Enable Phase 5 with tensor-parallel size `N` for dense models. | +| `process_isolation: false` | Use the legacy single-process lifecycle as a compatibility fallback. | + +Removed fields map to the current contract as follows: omit `check_source_load_parity`, `check_hf_reload`, and +`check_resume` to keep their phases enabled; use `skip_source_load_parity`, `skip_hf_reload`, or `skip_resume` to +disable one. Use `skip_automodel_reload_logit_parity` and `skip_hf_reload_logit_parity` instead of their shorter legacy +aliases. Replace generic or phase-specific legacy KL/cosine thresholds with `parity_tolerance_profile` and, when only +one comparison needs a different shared profile, `parity_tolerance_profile_overrides`. Use +`parity_threshold_overrides` only for a measured exception that does not fit a shared profile. `ci.time` must cover both finetune and robustness. Resume adds one short restored continuation; it no longer launches a separate fresh baseline. @@ -135,9 +259,9 @@ separate fresh baseline. ### Enable Checkpoint Robustness -1. Add `checkpoint_robustness:` under `ci:` with at least `hf_kl_threshold` and `tokenizer_name` +1. Add `checkpoint_robustness:` under `ci:` and set `tokenizer_name` when the model does not use the default Llama tokenizer IDs 2. Increase `ci.time` per the guidelines below -3. For large models, consider `no_check_resume: true` +3. For large vanilla-HF loads, consider `hf_device_map_auto: true`; add a `skip_*` field only with a documented blocker ### Enable vLLM Deploy @@ -173,7 +297,7 @@ known_issue: `ci.time` covers the entire SLURM job: finetune, robustness (if enabled), model downloads, setup, and teardown. -| Model Size | Finetune Only | Robustness (`no_check_resume`) | Robustness (full) | +| Model Size | Finetune Only | Robustness (`skip_resume`) | Robustness (full) | |------------|---------------|--------------------------------|-------------------| | < 2B | 10 min | 15 min | 15 min | | 2-5B | 12 min | 15 min | 20 min | diff --git a/tests/ci_tests/scripts/config_templates/ci_config.yaml b/tests/ci_tests/scripts/config_templates/ci_config.yaml index 0793f370ed..a8c21af657 100644 --- a/tests/ci_tests/scripts/config_templates/ci_config.yaml +++ b/tests/ci_tests/scripts/config_templates/ci_config.yaml @@ -113,24 +113,29 @@ computed: # resolved YAML's ci. block but does not apply them to the top-level. fixture_keys: checkpoint_robustness: - - check_hf_reload - - check_source_load_parity + # Specialized structural and loading controls. - check_fused_qkv_keys - check_phantom_keys - - check_resume - - cross_tp_kl_threshold - cross_tp_size - - cosine_threshold - experts_implementation - hf_adapter_ignored_key_prefix - hf_device_map_auto - hf_device_map_cpu_max_memory_gib - hf_device_map_max_memory_gib - - hf_cosine_threshold - - hf_kl_threshold - - kl_threshold + - hf_reload_timeout_seconds - hf_source_post_load_dequantize - - no_check_resume + # Full-logit parity defaults to 2048 tokens from the repository's fixed + # long-form document. Named profiles gate mean and p95 KL; all metrics are + # emitted for calibration. + - parity_sequence_length + - parity_tolerance_profile + # Optional source_load, automodel_reload, hf_reload, or cross_tp profile + # selections. Unspecified comparisons inherit parity_tolerance_profile. + - parity_tolerance_profile_overrides + # Prefer a profile. This optional mapping overrides only specified gates + # for a measured one-model exception in source_load, automodel_reload, + # hf_reload, or cross_tp; every other gate remains profile-owned. + - parity_threshold_overrides - process_isolation # Independent normal-finetune versus checkpoint-Phase-1 metric. This is # reported separately and never gates checkpoint resume correctness. @@ -142,11 +147,16 @@ fixture_keys: - resume_tolerance_profile - resume_first_loss_threshold - resume_loss_threshold + # Core phases are enabled by default. Use only these semantic skip controls + # for a documented model or backend blocker. + - skip_resume + - skip_source_load_parity + - skip_source_load_logit_parity - skip_hf_reload - - skip_automodel_logit_parity - - skip_hf_logit_parity - - source_load_cosine_threshold - - source_load_kl_threshold - - source_load_mean_kl_threshold + - skip_automodel_reload_logit_parity + - skip_hf_reload_logit_parity - tokenizer_name - trust_remote_code + # Removed positive check aliases and legacy max-KL fields are intentionally + # unsupported. Use the default-on phases, semantic skip fields, and named + # tolerance profiles above. diff --git a/tests/ci_tests/scripts/finetune_launcher.sh b/tests/ci_tests/scripts/finetune_launcher.sh index 599c30c7be..3dc3c6466f 100644 --- a/tests/ci_tests/scripts/finetune_launcher.sh +++ b/tests/ci_tests/scripts/finetune_launcher.sh @@ -147,7 +147,7 @@ if [[ "$HAS_ROBUSTNESS" == "true" ]]; then # recipes also retain enough CUDA, PP, scheduler, and dataloader ownership # that rebuilding several trainers in one interpreter is not reliable. read -r -a ROBUSTNESS_PHASES <<< \ - "${CHECKPOINT_ROBUSTNESS_PHASES:-train_and_save automodel_reload resume}" + "${CHECKPOINT_ROBUSTNESS_PHASES:-source_load_reference source_load_parity train_and_save automodel_reload hf_reload resume}" # Preserve the old harness's deferred-comparison behavior: record a failed # parity phase, continue independent phases, and return the first failure # only after every reachable phase has reported a result. diff --git a/tests/ci_tests/utils/generate_ci_tests.py b/tests/ci_tests/utils/generate_ci_tests.py index 35ef7698e9..210ac3905a 100644 --- a/tests/ci_tests/utils/generate_ci_tests.py +++ b/tests/ci_tests/utils/generate_ci_tests.py @@ -178,7 +178,13 @@ def _build_job( return job -def _enrich_base_job(job: Dict[str, Any], ci_config: Dict[str, Any], scope: str) -> None: +def _enrich_base_job( + job: Dict[str, Any], + ci_config: Dict[str, Any], + scope: str, + test_folder: str, + config: Path, +) -> None: """Add base-only extras: resource overrides, env_vars, HAS_ROBUSTNESS, convergence time.""" for ci_key, ci_var in CI_KEY_TO_VAR.items(): if ci_key not in ci_config: @@ -196,19 +202,26 @@ def _enrich_base_job(job: Dict[str, Any], ci_config: Dict[str, Any], scope: str) for key, value in ci_config.get("env_vars", {}).items(): job["variables"][key] = str(value) + has_robustness = "checkpoint_robustness" in ci_config robustness_config = ci_config.get("checkpoint_robustness") or {} - job["variables"]["HAS_ROBUSTNESS"] = str(bool(robustness_config)).lower() - if robustness_config.get("process_isolation"): + job["variables"]["HAS_ROBUSTNESS"] = str(has_robustness).lower() + supports_process_isolation = test_folder in {"llm_finetune", "vlm_finetune"} + process_isolation = supports_process_isolation and robustness_config.get("process_isolation", True) + if has_robustness and process_isolation: job["variables"]["CHECKPOINT_ROBUSTNESS_PROCESS_ISOLATION"] = "true" if "CHECKPOINT_ROBUSTNESS_PHASES" not in job["variables"]: robustness_phases = [] - if robustness_config.get("check_source_load_parity"): + source_load_parity_enabled = not robustness_config.get("skip_source_load_parity", False) + if source_load_parity_enabled: robustness_phases.extend(("source_load_reference", "source_load_parity")) robustness_phases.extend(("train_and_save", "automodel_reload")) if not robustness_config.get("skip_hf_reload"): robustness_phases.append("hf_reload") - if not robustness_config.get("no_check_resume"): + if not robustness_config.get("skip_resume"): robustness_phases.append("resume") + is_peft = "peft" in config.stem or "lora" in config.stem + if int(robustness_config.get("cross_tp_size") or 0) > 0 and not is_peft: + robustness_phases.append("cross_tp_reload") job["variables"]["CHECKPOINT_ROBUSTNESS_PHASES"] = " ".join(robustness_phases) # Convergence tests run for 2 epochs; double the slurm time allocation. @@ -248,7 +261,7 @@ def generate_job( if known_issue_id and not recipe_allow_failure: return [] - has_robustness = bool(ci_config.get("checkpoint_robustness")) + has_robustness = "checkpoint_robustness" in ci_config base_allow_failure = recipe_allow_failure or config.stem in (config_override.get("known_issue") or []) base_job = _build_job( @@ -259,7 +272,7 @@ def generate_job( allow_failure=base_allow_failure, known_issue_id=known_issue_id, ) - _enrich_base_job(base_job, ci_config, scope) + _enrich_base_job(base_job, ci_config, scope, test_folder, config) variants: list[tuple[str, Dict[str, Any]]] = [("", base_job)] # vLLM deploy variant. `ci.vllm_deploy_known_issue_id` suppresses just this diff --git a/tests/functional_tests/checkpoint_robustness/STATUS.md b/tests/functional_tests/checkpoint_robustness/STATUS.md index be3762726e..0a82afcbef 100644 --- a/tests/functional_tests/checkpoint_robustness/STATUS.md +++ b/tests/functional_tests/checkpoint_robustness/STATUS.md @@ -6,6 +6,15 @@ Historical model matrix last measured: 2026-04-02 UTC > **Note:** vLLM deployment tests moved to separate PR. +> **Archived results:** The model matrix, issue list, and follow-up list below are an April 2026 snapshot retained only +> for interpreting archived logs. They are not current release status. See the +> [checkpoint robustness documentation](../../ci_tests/README.md#checkpoint-robustness) for the current phase and +> configuration contract; current model findings live in recipe comments and tracked issues. + +> **Historical numbering:** The April matrix predates the current six-phase contract. Its “Phases 1–3” cover +> train/reference/AutoModel reload, and its “Phase 4” means the vanilla-HF reload that is now Phase 3. Historical +> labels below are preserved so they continue to match the archived logs. + ## Resume oracle policy Enabled LLM, VLM, and retrieval resume coverage now compares a restored trainer @@ -34,13 +43,15 @@ independent-run report is opportunistic rather than guaranteed coverage when phase-specific overrides make the existing runs incomparable; the blocking reproducibility oracle is the shared-trajectory resume comparison. -Recipes with `no_check_resume: true` remain explicitly exempt rather than being +Recipes with `skip_resume: true` remain explicitly exempt rather than being treated as passing. Those exemptions cover model/topology-specific restore blockers already documented in the recipe or this file (for example DeepEP/MoE state, hybrid Mamba state, or strict optimizer-state loading for unused/frozen parameters); they require focused fixes before resume coverage is enabled. -## Passing Models (8/15) +## Archived April 2026 model matrix + +### Passing models (8/15) | # | Model | SFT | PEFT | TP | Cross-TP | HF KL (SFT) | HF KL (PEFT) | VRAM SFT | VRAM PEFT | Resume | Special Flags | |---|-------|-----|------|----|----------|-------------|--------------|----------|-----------|--------|---------------| @@ -53,7 +64,7 @@ parameters); they require focused fixes before resume coverage is enabled. | 7 | Nemotron-Nano-8B-v1 | PASS | PASS | 2 | TP=2 (KL=0) | 4.2e-4 (t=7e-4) | 2.1e-3 (t=5e-3) | 7.73 GB | 4.42 GB | Disabled (Mamba) | check_fused_qkv_keys ✓, cross-TP ✓ | | 8 | Qwen3-MoE 30B | PASS | **FAIL** | 1 | — | 6.4e-5 (t=1e-4) | — | 28.18 GB | 11.81 GB | — | EP=8. SFT KL extremely low. **PEFT Phase 3 KL=0.84 — broken PEFT checkpoint reload, real bug** | -## Failing Models (5/15) +### Failing models (5/15) | # | Model | TP Tried | Error | Root Cause | Phases Passed | |---|-------|----------|-------|------------|---------------| @@ -63,7 +74,7 @@ parameters); they require focused fixes before resume coverage is enabled. | 12 | Mistral3 3B | TP=2, TP=1 | `fully_shard doesn't support scalar parameters (weight_scale_inv)` | FP8 quantized model has scalar scale params incompatible with FSDP2. Same error at both TP sizes. | Crashes during setup | | 8* | Qwen3-MoE PEFT | TP=1 EP=8 | Phase 3 KL=0.84 (should be 0) | **Real bug**: PEFT checkpoint reload is broken for Qwen3-MoE. SFT works fine. | Phase 1-2 PASS | -## Multi-Node Results (tested 2026-04-02) +### Multi-node results (tested 2026-04-02) | # | Model | Mode | Nodes | Config | Phases 1-3 | Phase 4 (HF) | Resume | Notes | |---|-------|------|-------|--------|-----------|--------------|--------|-------| @@ -73,36 +84,37 @@ parameters); they require focused fixes before resume coverage is enabled. | 15 | Super-120B | SFT | 4 (EP=32) | 32 GPUs | PASS | PASS (device_map=auto) | Disabled (MoE) | All phases pass, 9:27 | | 15 | Super-120B | PEFT | 2 (EP=16) | 16 GPUs | PASS | FAIL (KL=8.5e-2, t=7e-2) | Disabled (MoE) | Combined QKV in PEFT adapter | -## Not Yet Run +### Not yet run (None — all models tested) -## Known Issues +### Known issues recorded in April 2026 -- **MoE resume non-determinism**: DeepEP expert routing causes 3e-2 to 1e-1 loss diff. `--check_resume` disabled for MoE models. +- **MoE resume non-determinism**: DeepEP expert routing caused 3e-2 to 1e-1 loss diff; resume was disabled for those + historical runs. - **Mamba hybrid resume non-determinism**: Nano-8B-v1 has 0.62 loss diff on resume. Mamba layers have non-deterministic state. - **transformers 5.3 compatibility**: Flash 1B (triton_attention.py), Nano V2 (FSDP model attr), Baichuan (meta tensor). - **TP=2 failures**: Gemma 3 (1 KV head), Baichuan (custom layers), Mistral3 (FP8 scalars). Phi-4 TP=2 fixed on main. - **Combined QKV Phase 4 failures**: Super-49B and Super-120B PEFT produce combined projection keys (qkv_proj, gate_up_proj) in consolidated/adapter checkpoints. Vanilla HF models expect separate projections. StateDictAdapter conversion needed for Phase 4 to work. - **Qwen3-MoE PEFT bug**: Phase 3 KL=0.84 indicates broken PEFT checkpoint save/reload. Needs investigation in Qwen3MoeStateDictAdapter. -## TODO +### April 2026 follow-up list -### Investigate failures: +#### Investigate failures 1. **Super-49B Phase 4** — consolidated checkpoint has combined QKV keys. Need StateDictAdapter in Phase 4, or fix save_consolidated to split projections. 2. **Super-120B PEFT Phase 4** — same combined QKV issue for PEFT adapter weights. 3. **Qwen3-MoE PEFT bug** — investigate why Phase 3 KL=0.84 (real checkpoint bug) -### Investigate other failures (may need code fixes): +#### Investigate other failures (may need code fixes) 5. **Nemotron Flash 1B** — consolidated checkpoint missing triton_attention.py 6. **Nemotron Nano V2 9B** — FSDP wrapping issue 7. **Baichuan 2 7B** — meta tensor in Phase 4 HF loading 8. **Mistral3 3B** — FP8 scalar params vs FSDP2 -### Infrastructure improvements: +#### Infrastructure improvements 10. **Shared resume tolerance profiles and numeric overrides** — DONE 11. **Memory thresholds** for remaining models (Llama, GPT-OSS, Nano V3 still missing) -## Commits on branch `adil-a/checkpoint-robustness-test` +### Historical implementation commits on branch `adil-a/checkpoint-robustness-test` - `7ef62d55` — Nemotron Nano V3 checkpoint robustness + vLLM smoke tests - `04620847` — Cross-cutting features (tokenizer, memory, phantom keys, fused QKV, resume) diff --git a/tests/functional_tests/checkpoint_robustness/parity_document.mdx b/tests/functional_tests/checkpoint_robustness/parity_document.mdx new file mode 100644 index 0000000000..e6ed83db4a --- /dev/null +++ b/tests/functional_tests/checkpoint_robustness/parity_document.mdx @@ -0,0 +1,1051 @@ +--- +title: "Supervised Fine-Tuning (SFT) and Parameter-Efficient Fine-Tuning (PEFT) with NeMo AutoModel" +description: "" +position: 2 +--- +## Introduction + +Pretrained language models are general-purpose: they know a lot about language but nothing about your particular domain, terminology, or task. Fine-tuning bridges that gap. You fine-tune the model on your own examples so it produces answers that are accurate and relevant for your use case, without the cost of training a model from scratch. The result is a model optimized for your data that you can evaluate, publish, and deploy. This guide walks you through that process end-to-end with NeMo AutoModel (from installation through training, evaluation, and deployment) using [Meta Llama 3.2 1B](https://huggingface.co/meta-llama/Llama-3.2-1B) and the [SQuAD v1.1](https://huggingface.co/datasets/rajpurkar/squad) dataset as a running example. + +NeMo AutoModel supports two fine-tuning modes: + +- **Supervised Fine-Tuning (SFT)** updates all model parameters. Use SFT when you need maximum accuracy and have sufficient compute. +- **Parameter-Efficient Fine-Tuning (PEFT)** using [LoRA](https://arxiv.org/abs/2106.09685) freezes the base model and trains small low-rank adapters. PEFT reduces trainable parameters to less than 1% of the original model, lowering memory and storage costs. + +### Workflow Overview + +```text +┌──────────────┐ ┌──────────────┐ ┌──────────────┐ ┌──────────────┐ ┌──────────────┐ ┌──────────────┐ ┌──────────────┐ +│ 1. Install │--->│ 2. Configure │--->│ 3. Train │--->│ 4. Inference │--->│ 5. Evaluate │--->│ 6. Publish │--->│ 7. Deploy │ +│ │ │ │ │ │ │ │ │ │ │ (optional) │ │ (optional) │ +│ uv venv + │ │ YAML config │ │ automodel CLI│ │ HF generate │ │ Val loss + │ │ HF Hub │ │ vLLM serving │ +│ uv pip │ │ Choose SFT │ │ or torchrun │ │ API │ │ lm-eval- │ │ upload │ │ │ +│ or Docker │ │ or PEFT │ │ │ │ │ │ harness │ │ │ │ │ +└──────────────┘ └──────────────┘ └──────────────┘ └──────────────┘ └──────────────┘ └──────────────┘ └──────────────┘ +``` + +The following table outlines each step in the fine-tuning workflow for SFT and PEFT modes. + +| Step | Section | SFT | PEFT | +|------|---------|-----|------| +| **1. Install** | [Install NeMo AutoModel](#install-nemo-automodel) | Same | Same | +| **2. Configure** | [Configure Your Training Recipe](#configure-your-training-recipe) | YAML without `peft:` section | YAML with `peft:` section | +| **3. Train** | [Fine-Tune the Model](#fine-tune-the-model) | Same command for both modes | Same command for both modes | +| **4. Inference** | [Run Inference](#run-inference) | Load consolidated checkpoint directly | Load base model + adapter | +| **5. Evaluate** | [Evaluate the Fine-Tuned Model](#evaluate-the-fine-tuned-model) | Validation loss during training; lm-eval-harness post-training | Same | +| **6. Publish** | [Publish to HF Hub](#publish-to-the-hugging-face-hub) | Upload `model/consolidated/` | Upload `model/` (adapter only) | +| **7. Deploy** | [Deploy with vLLM](#deploy-with-vllm) | `vllm.LLM(model=...)` | Base model + native vLLM `LoRARequest` | + +## Install NeMo AutoModel + +```bash +uv venv --python 3.12 --seed +source .venv/bin/activate +uv pip install "nemo-automodel" --torch-backend=auto +``` + +Alternatively, if you run into dependency or driver issues, use the pre-built Docker container: + +```bash +docker pull nvcr.io/nvidia/nemo-automodel:26.06.00 +docker run --gpus all -it --rm --shm-size=8g -v $(pwd)/checkpoints:/opt/Automodel/checkpoints nvcr.io/nvidia/nemo-automodel:26.06.00 +``` + + +Docker containers are ephemeral, and files written inside the container are lost when it stops. The `-v` flag in the `docker run` command above bind-mounts a local `checkpoints/` directory into the container so that saved checkpoints persist across runs. For more details, see [Save Checkpoints When Using Docker](/development/checkpointing#save-checkpoints-when-using-docker). + + + +For the full set of installation methods, see the [Installation Guide](/get-started/installation). + +## Configure Your Training Recipe + +Training is configured through a [YAML](https://en.wikipedia.org/wiki/YAML) config file. For the `automodel` CLI, the file must name a **recipe**. The fine-tuning runtime also requires `model`, `dataset`, `dataloader`, `loss_fn`, `optimizer`, and `distributed` sections. If you configure a `validation_dataset`, also configure a `validation_dataloader`. The `step_scheduler` section is optional because it has typed default values, but you should set its global and local batch sizes explicitly for a predictable effective batch. Add `peft` only for LoRA. For the complete copy-pastable file, see [Full Config YAML](#full-config-yaml). + +Under the hood, both SFT and PEFT are executed by a **recipe**, which is a self-contained Python class that integrates model loading, dataset preparation, training, checkpointing, and logging. The fine-tuning recipe is [`TrainFinetuneRecipeForNextTokenPrediction`](https://github.com/NVIDIA-NeMo/Automodel/blob/main/nemo_automodel/recipes/llm/train_ft.py). The config file tells the recipe *what* to build; the recipe decides *how* to build it. + + +NeMo AutoModel configs use a convention borrowed from [Hydra](https://hydra.cc/): the special `_target_` key tells the framework *which* Python class or function to call, and **every other key** in the same YAML block is passed as a keyword argument to that call. For example: + +```yaml +optimizer: + _target_: torch.optim.Adam + lr: 1.0e-5 + weight_decay: 0 +``` + +At the YAML boundary, `RecipeConfig` normalizes this target and its fields into an `OptimizerConfig` wrapper. The recipe then supplies the runtime model parameters when it builds the optimizer: + +```python +optimizer_config = recipe_config.optimizer +optimizers = optimizer_config.build(model, device_mesh=device_mesh) +``` + +The `_target_` value is a **dotted Python import path**: the same string you would use in an `import` statement. The framework resolves it at runtime by importing the module and looking up the attribute. This means you can point `_target_` at any class constructor or factory function, and the remaining keys become its arguments. + + +For raw model, dataset, dataloader, and PEFT sections, inspect the Python signature of `_target_`. For typed sections such as optimizer, loss, step scheduler, learning rate (LR) scheduler, and checkpoint, inspect the corresponding config dataclass fields because `RecipeConfig` validates and owns those settings. + + + +From YAML to running code, here is the path a config takes through the framework: + +```text +finetune_config.yaml + │ + ▼ + ┌──────────────┐ load_yaml_config() parses the file into + │ ConfigNode │◄─── a tree of ConfigNode objects, one per + └──────┬───────┘ YAML section. + │ + ▼ + ┌──────────────┐ The recipe wraps the raw tree in RecipeConfig. + │ RecipeConfig │◄─── Known component sections become typed configs; + └──────┬───────┘ raw construction sections remain ConfigNodes. + │ + ┌────┴───────────────────────────────┐ + ▼ ▼ +typed component configs raw ConfigNode sections +optimizer, loss_fn, model, dataset, +step_scheduler, lr_scheduler, dataloader, peft +checkpoint + │ │ + ▼ ▼ +config.build(runtime arguments) recipe builders call instantiate() + │ with runtime model/tokenizer/data + └───────────────┬────────────────────┘ + ▼ + runtime components +``` + +Raw `ConfigNode` sections call `.instantiate()` inside the relevant recipe builder. `.instantiate()` does two things: + +1. **Resolves `_target_`**: Imports the Python path and obtains the callable (class or function). +2. **Calls it**: Passes every other key in the section as a keyword argument. + +Nested `_target_` blocks (such as a configured `collate_fn` inside `dataloader`) are recursively instantiated the same way. In contrast, `RecipeConfig` converts recognized sections to typed configs that own construction through `.build(...)`: optimizer receives the model and device mesh, step scheduler receives the dataloader, data-parallel (DP) size, and local batch size, and checkpoint receives runtime ranks and meshes. + +**The `recipe` key.** Every config file includes a top-level `recipe` key that tells the CLI *which recipe class* to run. You can write it as a **short name** or as a **fully-qualified Python path**, and both resolve to the same class: + +```yaml +# Short name (the CLI looks up the class automatically) +recipe: TrainFinetuneRecipeForNextTokenPrediction + +# Fully-qualified path (used as-is) +recipe: nemo_automodel.recipes.llm.train_ft.TrainFinetuneRecipeForNextTokenPrediction +``` + +The short name form is a convenience. The CLI scans all recipe modules under `nemo_automodel.recipes` and matches the bare class name. If you invoke the recipe script directly with `torchrun` instead of the `automodel` CLI, the `recipe` key is not required because the script itself *is* the recipe. + +Not every section uses `_target_`. `step_scheduler` and `checkpoint` are converted into `StepSchedulerConfig` and `CheckpointingConfig`, whose `.build(...)` methods receive runtime-only values from the recipe. The `distributed` mapping is parsed into typed strategy, parallelism, pipeline, and MoE configs before `DistributedSetup.build(...)` creates the meshes. + + +### Model + +```yaml +model: + _target_: nemo_automodel.NeMoAutoModelForCausalLM.from_pretrained + pretrained_model_name_or_path: meta-llama/Llama-3.2-1B +``` + +| Key | Role | +|-----|------| +| `_target_` | Points to [`NeMoAutoModelForCausalLM.from_pretrained`](https://github.com/NVIDIA-NeMo/Automodel/blob/main/nemo_automodel/_transformers/auto_model.py), which is a factory method that downloads (or loads from cache) a pretrained Hugging Face model and wraps it with NeMo distributed-training support. | +| `pretrained_model_name_or_path` | A keyword argument to `from_pretrained`. Any argument that [`from_pretrained`](https://huggingface.co/docs/transformers/main_classes/model#transformers.PreTrainedModel.from_pretrained) accepts can be added here (e.g., `cache_dir`, `torch_dtype`). | + +This guide uses **Meta Llama 3.2 1B** as a running example. Replace `pretrained_model_name_or_path` with any supported [Hugging Face model ID](/model-coverage/large-language-models/overview). + + +Llama is a family of decoder-only transformer models developed by Meta. The 1B variant is a compact model suitable for research and edge deployment, featuring rotary position embeddings (RoPE), grouped-query attention (GQA), and SwiGLU activations. + + + +Some Hugging Face models are **gated**. If the model page shows a **Request access** button, complete the following steps: + +1. Log in with your Hugging Face account and accept the license. +2. Ensure the token you use (from `huggingface-cli login` or `HF_TOKEN`) belongs to the approved account. + +Pulling a gated model without an authorized token triggers a 403 error. + + +### Dataset + +```yaml +dataset: + _target_: nemo_automodel.components.datasets.llm.squad.make_squad_dataset + dataset_name: rajpurkar/squad # HF-Hub ID used to pull the dataset + split: train + +validation_dataset: + _target_: nemo_automodel.components.datasets.llm.squad.make_squad_dataset + dataset_name: rajpurkar/squad + split: validation +``` + +| Key | Role | +|-----|------| +| `_target_` | Points to [`make_squad_dataset`](https://github.com/NVIDIA-NeMo/Automodel/blob/main/nemo_automodel/components/datasets/llm/squad.py), which is a factory function that downloads the SQuAD dataset, tokenizes it, and returns a `torch.utils.data.Dataset`. To use a different dataset, change `_target_` to a different factory function (see [Integrate Your Own Text Dataset](/datasets/text-dataset)). | +| `dataset_name`, `split` | Keyword arguments passed to `make_squad_dataset`. Each dataset factory defines its own parameters. Check the function signature to see what is available. | + +This guide uses **SQuAD v1.1** as a running example. Swap the dataset by changing `_target_` and the dataset arguments. See [Integrate Your Own Text Dataset](/datasets/text-dataset) and [Dataset Overview: LLM, VLM, and Retrieval Datasets](/datasets/overview). + + +The Stanford Question Answering Dataset (SQuAD) is a reading comprehension dataset where each example consists of a Wikipedia passage, a question, and an answer span. SQuAD v1.1 guarantees all questions are answerable from the context, making it suitable for straightforward fine-tuning. + +Example: +```json +{ + "context": "Architecturally, the school has a Catholic character. ...", + "question": "To whom did the Virgin Mary allegedly appear in 1858 in Lourdes France?", + "answers": { "text": ["Saint Bernadette Soubirous"], "answer_start": [515] } +} +``` + + +### Optional PEFT + +```yaml +peft: + _target_: nemo_automodel.components._peft.lora.PeftConfig + target_modules: "*_proj" # glob pattern matching projection layer FQNs + dim: 8 # low-rank dimension of the adapters + alpha: 32 # scaling factor for learned weights +``` + +| Key | Role | +|-----|------| +| `_target_` | Points to [`PeftConfig`](https://github.com/NVIDIA-NeMo/Automodel/blob/main/nemo_automodel/components/_peft/lora.py), which is a dataclass that describes which layers to adapt and how. Unlike the model and dataset sections, this instantiation produces a *config object*, not the adapter itself. The recipe passes the resulting `PeftConfig` into `build_model`, which applies LoRA adapters to the model. | +| `target_modules` | A glob pattern matched against fully-qualified layer names (for example, `"*_proj"` matches Llama projection layers such as `q_proj`, `k_proj`, and `down_proj`). | +| `dim` | The low-rank dimension *r*, which controls adapter capacity. Larger values learn more but use more memory. | +| `alpha` | Scaling factor applied to the adapter output (`alpha / dim`). Higher values give adapters more influence during training. | + +Including a `peft:` section enables LoRA fine-tuning. Remove it entirely to run SFT instead. See [Switch Between SFT and PEFT](#switch-between-sft-and-peft). + +#### QLoRA for Quantized Low-Rank Adaptation + +If GPU memory is a constraint, [QLoRA](https://arxiv.org/abs/2305.14314) combines LoRA with 4-bit NormalFloat (NF4) quantization to reduce memory usage by up to 75% compared to full-parameter SFT in 16-bit precision, while maintaining comparable quality to standard LoRA. + +QLoRA requires `bitsandbytes`, which is provided by AutoModel's `cuda_source` +install profile. Add it to the active environment created above before running a +QLoRA recipe: + +```bash +uv pip install "nemo-automodel[cuda_source]" +``` + +To enable QLoRA, add a `quantization:` section alongside the `peft:` section in your config. The example keeps the `"*_proj"` projection-layer pattern and increases `dim` from 8 to 16 for additional adapter capacity. + +```yaml +model: + _target_: nemo_automodel.NeMoAutoModelForCausalLM.from_pretrained + pretrained_model_name_or_path: meta-llama/Llama-3.2-1B + +peft: + _target_: nemo_automodel.components._peft.lora.PeftConfig + target_modules: "*_proj" # match projection layers such as q_proj and down_proj + dim: 16 # LoRA rank (higher than default to offset quantization) + alpha: 32 # scaling factor + dropout: 0.1 # LoRA dropout rate + +quantization: + load_in_4bit: True # enable 4-bit quantization + load_in_8bit: False # use 4-bit, not 8-bit + bnb_4bit_compute_dtype: bfloat16 # compute dtype + bnb_4bit_use_double_quant: True # double quantization for extra savings + bnb_4bit_quant_type: nf4 # NormalFloat quantization type + bnb_4bit_quant_storage: bfloat16 # storage dtype for quantized weights +``` + +### Training Schedule + +```yaml +step_scheduler: + global_batch_size: 64 # samples per optimizer step across all DP ranks + local_batch_size: 8 # samples per forward/backward pass on each DP rank + num_epochs: 1 # train over the dataset once +``` + +`step_scheduler` has no `_target_`. `RecipeConfig` converts it to `StepSchedulerConfig`, and the recipe builds the runtime scheduler with the training dataloader, data-parallel group size, and `local_batch_size`. Gradient accumulation is derived rather than configured directly: + +```text +gradient accumulation steps = global_batch_size / (local_batch_size × DP size) +``` + +With the values above, a single DP rank accumulates 8 micro-batches; 8 DP ranks require no accumulation. `global_batch_size` must be divisible by `local_batch_size × DP size`. + + +Most example recipes use bf16 training by default for memory and throughput. If you are running long fine-tuning, especially full-parameter SFT, and you require a higher-precision optimizer state, configure it explicitly instead of assuming it from the mixed-precision compute policy. See the [Mixed-Precision Training Guide](/development/mixed-precision-training) for the recommended Transformer Engine (TE) and torch AdamW patterns. + + + +### Full Config YAML + + +Save as `finetune_config.yaml`. This config runs PEFT (LoRA). To run SFT instead, remove the `peft:` section. Use this complete config rather than the current checked-in Llama SFT YAML. That file still nests the seed under `rng:`, which this recipe ignores, and its commented W&B block uses unsupported `save_dir` instead of `dir`. + +```yaml +recipe: TrainFinetuneRecipeForNextTokenPrediction + +step_scheduler: + global_batch_size: 64 + local_batch_size: 8 + ckpt_every_steps: 10 + val_every_steps: 10 + num_epochs: 1 + +dist_env: + backend: nccl + timeout_minutes: 1 + +distributed: + strategy: fsdp2 + dp_size: null + tp_size: 1 + cp_size: 1 + +model: + _target_: nemo_automodel.NeMoAutoModelForCausalLM.from_pretrained + pretrained_model_name_or_path: meta-llama/Llama-3.2-1B + +peft: + _target_: nemo_automodel.components._peft.lora.PeftConfig + target_modules: "*_proj" + dim: 8 + alpha: 32 + +loss_fn: + _target_: nemo_automodel.components.loss.masked_ce.MaskedCrossEntropy + +dataset: + _target_: nemo_automodel.components.datasets.llm.squad.make_squad_dataset + dataset_name: rajpurkar/squad + split: train + +dataloader: + _target_: torchdata.stateful_dataloader.StatefulDataLoader + collate_fn: + _target_: nemo_automodel.components.datasets.utils.default_collater + shuffle: false + +validation_dataset: + _target_: nemo_automodel.components.datasets.llm.squad.make_squad_dataset + dataset_name: rajpurkar/squad + split: validation + +validation_dataloader: + _target_: torchdata.stateful_dataloader.StatefulDataLoader + collate_fn: + _target_: nemo_automodel.components.datasets.utils.default_collater + +optimizer: + _target_: torch.optim.Adam + lr: 1.0e-5 + betas: [0.9, 0.999] + eps: 1.0e-8 + weight_decay: 0 + +checkpoint: + enabled: true + checkpoint_dir: checkpoints/ + model_save_format: safetensors + save_consolidated: final +``` + + +## Fine-Tune the Model + +You can run the recipe using the AutoModel CLI or directly with `torchrun` (advanced). + +```bash +automodel --nproc-per-node=8 finetune_config.yaml +``` + +The `--nproc-per-node=8` flag specifies the number of GPUs per node. Adjust as needed. For a single GPU, omit the `--nproc-per-node` option. + +### Invoke the Recipe Script Directly (Advanced) + +Alternatively, you can invoke the [`train_ft.py`](https://github.com/NVIDIA-NeMo/Automodel/blob/main/nemo_automodel/recipes/llm/train_ft.py) recipe script directly by using [`torchrun`](https://docs.pytorch.org/docs/stable/elastic/run.html): + +```bash +torchrun --nproc-per-node=8 nemo_automodel/recipes/llm/train_ft.py -c finetune_config.yaml +``` + +### Sample Output +Running the recipe with the `automodel` application or by invoking the recipe script directly produces the following log output: + +``` +$ automodel finetune_config.yaml +INFO:nemo_automodel.cli.app:Config: finetune_config.yaml +INFO:nemo_automodel.cli.app:Recipe: nemo_automodel.recipes.llm.train_ft.TrainFinetuneRecipeForNextTokenPrediction +INFO:nemo_automodel.cli.app:Launching job interactively (local) +cfg-path: finetune_config.yaml +INFO:root:step 4 | epoch 0 | loss 1.5514 | grad_norm 102.0000 | mem: 11.66 GiB | tps 6924.50 +INFO:root:step 8 | epoch 0 | loss 0.7913 | grad_norm 46.2500 | mem: 14.58 GiB | tps 9328.79 +Saving checkpoint to checkpoints/epoch_0_step_10 +INFO:root:step 12 | epoch 0 | loss 0.4358 | grad_norm 23.8750 | mem: 15.48 GiB | tps 9068.99 +INFO:root:step 16 | epoch 0 | loss 0.2057 | grad_norm 12.9375 | mem: 16.47 GiB | tps 9148.28 +INFO:root:step 20 | epoch 0 | loss 0.2557 | grad_norm 13.4375 | mem: 12.35 GiB | tps 9196.97 +Saving checkpoint to checkpoints/epoch_0_step_20 +INFO:root:[val] step 20 | epoch 0 | loss 0.2469 +``` + +Each log line reports the current loss, gradient norm, peak GPU memory, and tokens per second (TPS). Small fluctuations between steps (for example, 0.2057 to 0.2557) are normal. Look at the overall downward trend rather than individual values. + +### Checkpoint Contents + +Checkpoints are saved as Hugging Face-compatible safetensors. For SFT, `save_consolidated: final` is the default. Intermediate checkpoints contain sharded model weights and a generated `model/consolidate.sh` helper. The final checkpoint also contains `model/consolidated/` for Transformers, vLLM, lm-eval-harness, and other Hugging Face ecosystem tools. Set `save_consolidated: false` to keep every checkpoint sharded, and then run the helper script for whichever checkpoint you want to export. Use `save_consolidated: every` (or the legacy value `true`) only when you intentionally want inline Hugging Face export at every checkpoint save. + +**PEFT checkpoints** contain only the adapter weights (megabytes instead of gigabytes) and are saved directly under `model/`. They do not use `model/consolidate.sh`. At inference time, you must load the original base model and apply the adapter on top. This distinction affects every downstream step, including inference, publishing, and deployment. + + +**SFT checkpoint:** +```bash +$ tree checkpoints/epoch_0_step_10/ +checkpoints/epoch_0_step_10/ +├── config.yaml +├── dataloader.pt +├── model +│ ├── consolidate.sh +│ ├── shard-00001-model-00001-of-00001.safetensors +│ └── shard-00002-model-00001-of-00001.safetensors +├── optim +│ ├── __0_0.distcp +│ └── __1_0.distcp +├── rng.pt +└── step_scheduler.pt + +3 directories, 9 files +``` + +**PEFT checkpoint:** +```bash +$ tree checkpoints/epoch_0_step_10/ +checkpoints/epoch_0_step_10/ +├── dataloader.pt +├── config.yaml +├── model +│ ├── adapter_config.json +│ ├── adapter_model.safetensors +│ └── automodel_peft_config.json +├── optim +│ ├── __0_0.distcp +│ └── __1_0.distcp +├── rng.pt +└── step_scheduler.pt + +2 directories, 8 files +``` + + +## Run Inference + +Inference uses the Hugging Face `generate` API. Because exported SFT checkpoints are self-contained while PEFT checkpoints store only adapter weights (see [Checkpoint Contents](#checkpoint-contents)), the loading procedure differs between the two modes. + +### SFT Inference + +For an intermediate checkpoint, or any run with `save_consolidated: false`, first run the generated helper for the checkpoint you want to load: + +```bash +bash checkpoints/epoch_0_step_10/model/consolidate.sh +``` + +The exported SFT checkpoint at `model/consolidated/` is a complete Hugging Face model and can be loaded directly: + +```python +import torch +from transformers import AutoModelForCausalLM, AutoTokenizer + +ckpt_path = "checkpoints/epoch_0_step_10/model/consolidated" +tokenizer = AutoTokenizer.from_pretrained(ckpt_path) +model = AutoModelForCausalLM.from_pretrained(ckpt_path) + +device = "cuda" if torch.cuda.is_available() else "cpu" +model.to(device) + +prompt = ( + "Context: Architecturally, the school has a Catholic character. " + "Atop the Main Building's gold dome is a golden statue of the Virgin Mary. " + "Immediately in front of the Main Building and facing it, is a copper statue of Christ " + "with arms upraised with the legend 'Venite Ad Me Omnes'.\n\n" + "Question: What is atop the Main Building?\n\n" + "Answer:" +) +inputs = tokenizer(prompt, return_tensors="pt").to(device) +output = model.generate(**inputs, max_new_tokens=50) +print(tokenizer.decode(output[0], skip_special_tokens=True)) +``` + +### PEFT Inference + +PEFT adapters must be loaded on top of the base model. Because the base NeMo AutoModel installation does not include the Hugging Face PEFT library, install a tested version in your active environment before running the inference code: + +```bash +uv pip install "peft>=0.18.1" +``` + +```python +import torch +from transformers import AutoModelForCausalLM, AutoTokenizer +from peft import PeftModel + +base_model_name = "meta-llama/Llama-3.2-1B" +tokenizer = AutoTokenizer.from_pretrained(base_model_name) +model = AutoModelForCausalLM.from_pretrained(base_model_name) + +adapter_path = "checkpoints/epoch_0_step_10/model/" +model = PeftModel.from_pretrained(model, adapter_path) + +device = "cuda" if torch.cuda.is_available() else "cpu" +model.to(device) + +prompt = ( + "Context: Architecturally, the school has a Catholic character. " + "Atop the Main Building's gold dome is a golden statue of the Virgin Mary. " + "Immediately in front of the Main Building and facing it, is a copper statue of Christ " + "with arms upraised with the legend 'Venite Ad Me Omnes'.\n\n" + "Question: What is atop the Main Building?\n\n" + "Answer:" +) +inputs = tokenizer(prompt, return_tensors="pt").to(device) +output = model.generate(**inputs, max_new_tokens=50) +print(tokenizer.decode(output[0], skip_special_tokens=True)) +``` + +## Evaluate the Fine-Tuned Model + +### Validation Loss During Training + +The recipe automatically computes validation loss at the interval set by `val_every_steps`. Look for `[val]` lines in the training log: + +```text +INFO:root:[val] step 20 | epoch 0 | loss 0.2469 +``` + +A decreasing validation loss across checkpoints indicates the model is learning. If validation loss plateaus or increases while training loss continues to drop, the model might be overfitting, so consider stopping earlier or reducing the learning rate. + +### Post-Training Evaluation Using lm-eval-harness + +For task-specific benchmarks (for example, MMLU, GSM8K, and HellaSwag accuracy), use [lm-eval-harness](https://github.com/EleutherAI/lm-evaluation-harness) with the fine-tuned checkpoint. For an intermediate SFT checkpoint, or any SFT run with `save_consolidated: false`, run `bash checkpoints/epoch_0_step_20/model/consolidate.sh` before pointing evaluation at `model/consolidated/`. Create and activate a clean evaluation environment, and then install both model backends before running either command: + +```bash +uv venv --python 3.12 --seed .venv-lm-eval +source .venv-lm-eval/bin/activate +uv pip install "lm_eval[hf,vllm]" --torch-backend=cu130 + +# SFT checkpoint (using vLLM backend for faster evaluation) +lm_eval --model vllm \ + --model_args pretrained=checkpoints/epoch_0_step_20/model/consolidated/ \ + --tasks hellaswag \ + --batch_size auto + +# PEFT adapter (using Hugging Face backend with built-in PEFT support) +lm_eval --model hf \ + --model_args pretrained=meta-llama/Llama-3.2-1B,peft=checkpoints/epoch_0_step_20/model/ \ + --tasks hellaswag \ + --batch_size auto +``` + + +The quoted `lm_eval[hf,vllm]` requirement installs both backends and their +dependencies. The SFT example uses the `vllm` backend for faster evaluation. +The PEFT example uses the `hf` backend with lm-eval's built-in PEFT support to +load the adapter on top of the base model. + + + + +Run `lm_eval` on the base model *before* fine-tuning to establish a baseline, and then compare it against the fine-tuned checkpoint. + + + +## Publish to the Hugging Face Hub + +Fine-tuned checkpoints and PEFT adapters are stored in Hugging Face-native format and can be uploaded directly to the Hub. For an intermediate SFT checkpoint, or any SFT run with `save_consolidated: false`, upload `model/consolidated/` after running the generated consolidation helper. + +1. Install the Hugging Face Hub library (if not already installed): + +```bash +uv pip install "huggingface_hub" +``` + +2. Log in to Hugging Face: + +```bash +huggingface-cli login +``` + +3. Upload the folder: + +**SFT checkpoint:** +```python +from huggingface_hub import HfApi + +api = HfApi() +api.upload_folder( + folder_path="checkpoints/epoch_0_step_10/model/consolidated", + repo_id="your-username/llama3.2_1b-finetuned-squad", + repo_type="model", +) +``` + +**PEFT adapter:** +```python +from huggingface_hub import HfApi + +api = HfApi() +api.upload_folder( + folder_path="checkpoints/epoch_0_step_10/model", + repo_id="your-username/llama3.2_1b-lora-squad", + repo_type="model", +) +``` + +After uploading, load the checkpoint or adapter directly from the Hugging Face Hub: + +**SFT:** +```python +from transformers import AutoModelForCausalLM + +model = AutoModelForCausalLM.from_pretrained("your-username/llama3.2_1b-finetuned-squad") +``` + +**PEFT:** +```python +from transformers import AutoModelForCausalLM +from peft import PeftModel + +model = AutoModelForCausalLM.from_pretrained("meta-llama/Llama-3.2-1B") +model = PeftModel.from_pretrained(model, "your-username/llama3.2_1b-lora-squad") +``` + +## Deploy with vLLM + +[vLLM](https://github.com/vllm-project/vllm) is an efficient inference engine for production deployment of LLMs. + + +Use a vLLM-compatible environment. For NVIDIA GPUs, create and activate a virtual environment before installing vLLM: + +```bash +uv venv --python 3.12 --seed .venv-vllm +source .venv-vllm/bin/activate +uv pip install "vllm" --torch-backend=cu130 +``` + +You can skip these commands if your active environment already includes vLLM. + + + +### SFT Checkpoint with vLLM + +For an intermediate checkpoint, or any run with `save_consolidated: false`, run the generated `model/consolidate.sh` helper before serving from `model/consolidated/`: + +```python +from vllm import LLM, SamplingParams + +llm = LLM(model="checkpoints/epoch_0_step_10/model/consolidated/", model_impl="transformers") +params = SamplingParams(max_tokens=20) +outputs = llm.generate("Toronto is a city in Canada.", sampling_params=params) +print(f"Generated text: {outputs[0].outputs[0].text}") +``` +```text +>>> Generated text: It is the capital of Ontario. Toronto is a global hub for cultural tourism. The City of Toronto +``` + +### PEFT Adapter with vLLM + +vLLM can load the NeMo AutoModel LoRA adapter directly on top of the original base model. Enable LoRA when constructing `LLM`, and then select the adapter for the request with `LoRARequest`: + +```python +from vllm import LLM, SamplingParams +from vllm.lora.request import LoRARequest + +base_model = "meta-llama/Llama-3.2-1B" +adapter_path = "checkpoints/epoch_0_step_10/model" + +llm = LLM(model=base_model, enable_lora=True, max_lora_rank=8) +params = SamplingParams(temperature=0, max_tokens=20) +outputs = llm.generate( + ["How are you doing?"], + params, + lora_request=LoRARequest("squad_adapter", 1, adapter_path), +) +print(outputs[0].outputs[0].text) +``` + +## Full Configuration Reference + +This section documents all available config fields for the fine-tuning recipe. For the quick-start config, see [Configure Your Training Recipe](#configure-your-training-recipe). + +### Switch Between SFT and PEFT + +The `peft:` section controls which mode runs, as described in the following table. + +| Mode | What to Do in the YAML | +|------|----------------------| +| **PEFT (LoRA)** | Include the `peft:` section as shown below. | +| **SFT (full-parameter)** | Remove/comment the `peft:` section entirely. | + +All other config sections remain the same for both modes. + +### Full Configuration + + +```yaml +# Recipe +# Selects which recipe class runs the training loop. +# Use a short name (auto-discovered) or a fully-qualified Python path: +# recipe: nemo_automodel.recipes.llm.train_ft.TrainFinetuneRecipeForNextTokenPrediction +recipe: TrainFinetuneRecipeForNextTokenPrediction + +# Training Schedule +# RecipeConfig converts this mapping to StepSchedulerConfig. The recipe passes +# runtime dataloader and DP values to its build() method. +step_scheduler: + global_batch_size: 64 # samples per optimizer step across all DP ranks + local_batch_size: 8 # samples per forward/backward pass on each DP rank + # accumulation = global / (local × DP size) + ckpt_every_steps: 10 # save a checkpoint every N optimizer steps + val_every_steps: 10 # run the validation loop every N optimizer steps + num_epochs: 1 # how many full passes over the training dataset + +# Process Group +# Initializes the PyTorch distributed process group. +# No _target_ (backend and timeout are passed to initialize_distributed()). +# You normally would not need to tune this. +dist_env: + backend: nccl # communication backend: "nccl" (GPU, recommended) or "gloo" (CPU) + timeout_minutes: 1 # timeout for collective operations; increase for large models + # that take longer to initialize + +# Distributed Strategy +# Determines how model weights, data, and compute are split across GPUs. +# Parsed into typed strategy and parallelism configs before DistributedSetup.build(). +# See "Distributed Training: TP, PP, CP, and EP" in Advanced Topics for details. +distributed: + strategy: fsdp2 # parallelism strategy: "fsdp2" (recommended), "megatron_fsdp", + # or "ddp". FSDP2 shards parameters and optimizer states across + # the data-parallel group. + dp_size: null # data-parallel group size. null = auto-detect from + # world_size ÷ (tp_size × cp_size × pp_size). + tp_size: 1 # tensor-parallel size: splits weight matrices across GPUs. + # Set to 2, 4, or 8 if the model does not fit on one GPU. + # Should divide evenly into the number of attention heads. + cp_size: 1 # context-parallel size: splits the input sequence across GPUs. + # Increase for very long contexts (for example, 32k+ tokens). + sequence_parallel: false # when true, extends TP to also shard activations along + # the sequence dimension for additional memory savings + +{/* docs-review-start: mamba-ssd-prewarm */} + +# Setup Prewarms (optional) +# Moves selected one-time CUDA and NCCL initialization work out of the first +# training step. Leave disabled unless the corresponding lazy initialization +# causes a first-step out-of-memory error. +prewarm: + cublas_backward: false # initialize cuBLAS backward workspaces during setup + fla_gdn_autotune: false # autotune FLA gated-delta-net Triton kernels during setup + mamba_ssd_autotune: false # autotune Mamba SSD Triton kernels during setup + comm_groups: false # initialize grad-norm communication groups during setup + +{/* docs-review-end: mamba-ssd-prewarm */} + +# Random Number Generator +# The recipe constructs its checkpointable, rank-aware StatefulRNG from this seed. +seed: 1111 # global random seed for reproducibility + +# Model +# _target_ → NeMoAutoModelForCausalLM.from_pretrained: downloads (or loads from +# cache) a pretrained Hugging Face model and wraps it with NeMo distributed-training +# support. Any from_pretrained kwarg is accepted (such as cache_dir and torch_dtype). +model: + _target_: nemo_automodel.NeMoAutoModelForCausalLM.from_pretrained + pretrained_model_name_or_path: meta-llama/Llama-3.2-1B + +# PEFT (remove / comment this entire section for full-parameter SFT) +# _target_ → PeftConfig: a dataclass describing which layers get LoRA adapters. +# The recipe passes this config into build_model(), which attaches adapters +# to the matching layers. +peft: + _target_: nemo_automodel.components._peft.lora.PeftConfig + target_modules: "*_proj" # glob pattern matched against fully-qualified layer names; + # matches Llama projection layers such as q_proj and down_proj + dim: 8 # low-rank dimension r, which controls adapter capacity. + # Larger values are more expressive but use more memory. + alpha: 32 # LoRA scaling factor: adapter output is scaled by alpha/dim. + # Higher values give adapters more influence during training. + use_triton: True # use an optimized Triton kernel for LoRA forward/backward + # (requires the triton package) + +# Checkpointing +# RecipeConfig converts this mapping to CheckpointingConfig; the recipe supplies +# runtime ranks and meshes to its build() method. +checkpoint: + enabled: true # set to false to skip saving checkpoints entirely + checkpoint_dir: checkpoints/ # output directory. Docker users: bind-mount this path + # (for example, -v $(pwd)/checkpoints:/opt/Automodel/checkpoints) + # to persist checkpoints across container restarts. + model_save_format: safetensors # "safetensors" (recommended, faster and safer) or + # "torch_save" (legacy pickle-based format) + save_consolidated: final # default: intermediate checkpoints are sharded; the final checkpoint + # also exports consolidated HF weights. Other modes: false (sharded only) + # or every/true (export every checkpoint). + +# Training Dataset +# _target_ → make_squad_dataset: a factory function that downloads the SQuAD +# dataset, tokenizes it, and returns a torch Dataset. To use a different dataset, +# change _target_ to another factory function (see the dataset guide). +dataset: + _target_: nemo_automodel.components.datasets.llm.squad.make_squad_dataset + dataset_name: rajpurkar/squad # HuggingFace Hub dataset ID + split: train # which split to use (train, validation, test) + +# Validation Dataset +validation_dataset: + _target_: nemo_automodel.components.datasets.llm.squad.make_squad_dataset + dataset_name: rajpurkar/squad + split: validation + limit_dataset_samples: 64 # cap validation set to 64 samples for faster eval loops; + # remove this line to use the full validation set + +# Training Dataloader +# _target_ → StatefulDataLoader: a checkpointable DataLoader from torchdata that +# saves and restores iteration state across training restarts, so resumed runs +# do not re-process already-seen batches. +dataloader: + _target_: torchdata.stateful_dataloader.StatefulDataLoader + collate_fn: + _target_: nemo_automodel.components.datasets.utils.default_collater + # pads and batches individual samples into tensors; + # can be swapped for custom collation + shuffle: true # whether to shuffle the dataset each epoch + +# The recipe sets the actual loader batch size from +# step_scheduler.local_batch_size, overriding any dataloader.batch_size key. + +# Validation Dataloader +validation_dataloader: + _target_: torchdata.stateful_dataloader.StatefulDataLoader + collate_fn: + _target_: nemo_automodel.components.datasets.utils.default_collater + +# Loss Function +# _target_ → MaskedCrossEntropy: standard cross-entropy loss that automatically +# ignores padding tokens so they do not affect the gradient. +# Other available loss functions (swap _target_ to use): +# - nemo_automodel.components.loss.chunked_ce.ChunkedCrossEntropy +# Computes CE in chunks along the sequence dimension to reduce peak memory. +# Useful for very long sequences. Accepts chunk_len (default 32). +# - nemo_automodel.components.loss.linear_ce.FusedLinearCrossEntropy +# Fuses the final linear projection (lm_head) with the CE computation, +# avoiding the full logit tensor. Significant **memory savings** for large vocabs. +# - nemo_automodel.components.loss.te_parallel_ce.TEParallelCrossEntropy +# TE-based parallel CE with a Triton kernel. Designed for tensor-parallel +# setups where logits are sharded across TP ranks. +loss_fn: + _target_: nemo_automodel.components.loss.masked_ce.MaskedCrossEntropy + +# Optimizer +# _target_ → torch.optim.Adam: RecipeConfig wraps this target and its fields in +# an OptimizerConfig, which receives model parameters when the recipe calls build(). +optimizer: + _target_: torch.optim.Adam + lr: 1.0e-5 # learning rate (the most important hyperparameter to tune) + betas: [0.9, 0.999] # Adam momentum coefficients (β₁ for mean, β₂ for variance) + eps: 1e-8 # small constant added to the denominator for numerical stability + weight_decay: 0 # L2 regularization strength (0 = no regularization) + +# Logging (optional) +# Uncomment to enable Weights & Biases experiment tracking. +# wandb: +# project: # W&B project name +# entity: # W&B team or username +# name: # display name for this run +# dir: # local directory for W&B artifacts +``` + + +### Config Field Reference + +The following table describes the config fields for the fine-tuning recipe. + +| Section | Required? | What to Change | +|---------|-----------|----------------| +| `recipe` | CLI only | Required by `automodel`; select `TrainFinetuneRecipeForNextTokenPrediction` by short name or fully-qualified path. | +| `model` | Yes | Set `pretrained_model_name_or_path` to your Hugging Face model ID. Source: [`auto_model.py`](https://github.com/NVIDIA-NeMo/Automodel/blob/main/nemo_automodel/_transformers/auto_model.py). | +| `peft` | PEFT only | Remove entirely for SFT. Adjust `dim` and `alpha` to tune adapter capacity. `use_triton: True` enables an optimized LoRA kernel (requires the `triton` package). For reduced memory usage, see [QLoRA](#qlora-quantized-low-rank-adaptation). Source: [`lora.py`](https://github.com/NVIDIA-NeMo/Automodel/blob/main/nemo_automodel/components/_peft/lora.py). | +| `dataset` | Yes | Change `_target_`, `dataset_name`, and `split` for your data. Source: [`squad.py`](https://github.com/NVIDIA-NeMo/Automodel/blob/main/nemo_automodel/components/datasets/llm/squad.py). | +| `dataloader` | Yes | Configure the loader target, collation, workers, and shuffling. Set its effective batch size with `step_scheduler.local_batch_size`; the recipe injects that value into the loader. Uses [`StatefulDataLoader`](https://meta-pytorch.org/data/main/torchdata.stateful_dataloader.html) for checkpointable iteration. Collation: [`utils.py`](https://github.com/NVIDIA-NeMo/Automodel/blob/main/nemo_automodel/components/datasets/utils.py). | +| `validation_dataset`, `validation_dataloader` | For validation | Configure both when running validation. The validation loader also receives `step_scheduler.local_batch_size`. | +| `loss_fn` | Yes | [`MaskedCrossEntropy`](https://github.com/NVIDIA-NeMo/Automodel/blob/main/nemo_automodel/components/loss/masked_ce.py) is the standard choice. Alternatives: [`ChunkedCrossEntropy`](https://github.com/NVIDIA-NeMo/Automodel/blob/main/nemo_automodel/components/loss/chunked_ce.py) (long sequences), [`FusedLinearCrossEntropy`](https://github.com/NVIDIA-NeMo/Automodel/blob/main/nemo_automodel/components/loss/linear_ce.py) (large vocabs), [`TEParallelCrossEntropy`](https://github.com/NVIDIA-NeMo/Automodel/blob/main/nemo_automodel/components/loss/te_parallel_ce.py) (tensor-parallel). | +| `seed` | Optional | Controls the checkpointable, rank-aware RNG that the recipe constructs. Source: [`rng.py`](https://github.com/NVIDIA-NeMo/Automodel/blob/main/nemo_automodel/components/training/rng.py). | +| `step_scheduler` | Optional (recommended) | Set `global_batch_size` and `local_batch_size`; accumulation is derived as `global / (local × DP size)`. `ckpt_every_steps` and `val_every_steps` are counted in optimizer steps. | +| `distributed` | Yes | `dp_size: null` means auto-detect from world size. Adjust `tp_size` for tensor parallelism across GPUs. | +| `prewarm` | Optional | Enable only the targeted setup warmup associated with a first-step out-of-memory error. See [Prewarm One-Time CUDA Initialization](/get-started/configuration#prewarm-one-time-cuda-initialization). Source: [`prewarm.py`](https://github.com/NVIDIA-NeMo/Automodel/blob/main/nemo_automodel/components/training/prewarm.py). | +| `checkpoint` | Recommended | Set `checkpoint_dir` to a persistent path, especially in Docker. `save_consolidated` defaults to `final`: intermediate SFT checkpoints are sharded and the final checkpoint also exports consolidated HF weights. | +| `optimizer` | Yes | Any supported `torch.optim` class can be selected using `_target_`; `RecipeConfig` normalizes it to an `OptimizerConfig` that owns construction. For long fine-tuning, especially full-parameter SFT, see the [mixed-precision training guide](/development/mixed-precision-training) before combining torch AdamW with bf16 resident parameters. | +| `wandb` | Optional | Uncomment and configure to enable Weights & Biases logging. | + +For the fine-tuning recipe itself, see [`train_ft.py`](https://github.com/NVIDIA-NeMo/Automodel/blob/main/nemo_automodel/recipes/llm/train_ft.py). + +### Per-Parameter-Group Learning Rate + +To give a subset of parameters a different learning rate or weight decay (for example, a smaller learning rate for MoE router parameters), add `param_group_overrides` to the `optimizer` block. Each entry matches parameters by name and applies a multiplier that the LR scheduler folds into that group every step: + +```yaml +optimizer: + _target_: torch.optim.AdamW + lr: 2.0e-5 + weight_decay: 0.01 + param_group_overrides: + - pattern: "router" # Python regex matched against each parameter name (re.search) + lr_mult: 0.1 # this group trains at 0.1 x the scheduled learning rate + wd_mult: 1.0 # (optional) multiplier for weight decay; defaults to 1.0 +``` + +Each parameter joins the group of the **first** override whose `pattern` matches its name; parameters matching no pattern keep the base learning rate. A pattern that matches nothing is skipped with a warning. This works for the standard `torch.optim` optimizers; Dion-family optimizers do their own parameter grouping and ignore `param_group_overrides`. + +## Distributed Training with TP, PP, CP, and EP + +The `distributed:` section controls how the model and data are split across GPUs. NeMo AutoModel supports five parallelism dimensions, each of which slices the workload differently, as described in the following table: + +| Dimension | Key | What It Shards | When to Use | +|-----------|-----|---------------|-------------| +| **Data Parallelism (DP)** | `dp_size` | Replicates the model on each group of GPUs; each replica trains on a different data batch. | Default. Scales batch size linearly with GPU count. | +| **Tensor Parallelism (TP)** | `tp_size` | Splits individual weight matrices (attention, MLP) across GPUs within a node. | Model is too large to fit on a single GPU, or you want to reduce per-GPU memory at the cost of extra communication. | +| **Pipeline Parallelism (PP)** | `pp_size` | Assigns different *layers* (stages) to different GPUs and pipelines micro-batches through them. | Very deep models that do not fit even with TP, or multi-node training where TP's all-reduce is too expensive across nodes. | +| **Context Parallelism (CP)** | `cp_size` | Splits the input *sequence* across GPUs so each GPU processes a portion of the context. | Very long sequences that exceed single-GPU memory. | +| **Expert Parallelism (EP)** | `ep_size` | Distributes MoE experts across GPUs so each GPU holds a subset of experts. | Mixture-of-Experts models only. | + +These dimensions compose with each other. The relationship between them and total GPU count is: + +```text +world_size = pp_size × dp_size × cp_size × tp_size +``` + +When `dp_size` is set to `null` (the default), it is inferred automatically: + +```text +dp_size = world_size ÷ (tp_size × cp_size × pp_size) +``` + +EP does not appear in this formula. Experts are distributed across the DP×CP rank groups, with the constraint that `(dp_size × cp_size)` must be divisible by `ep_size`. + +#### Data Parallelism as the Default + +Data parallelism is the default. With `strategy: fsdp2`, FSDP2 shards both model parameters and optimizer states across the DP group, so memory usage shrinks as you add GPUs: + +```yaml +distributed: + strategy: fsdp2 + dp_size: null # auto-detected from world_size ÷ (tp × cp × pp) + tp_size: 1 + cp_size: 1 +``` + +#### Tensor Parallelism + +TP splits weight matrices across GPUs within a single node. Set `tp_size` to the number of GPUs you want to shard over. The value is typically 2, 4, or 8 and should divide evenly into the number of attention heads: + +```yaml +distributed: + strategy: fsdp2 + dp_size: null + tp_size: 4 + cp_size: 1 + sequence_parallel: false # set to true for additional memory savings +``` + +`sequence_parallel: true` extends TP to also shard activation memory along the sequence dimension, further reducing per-GPU memory at the cost of additional communication. + +#### Pipeline Parallelism + +PP assigns groups of layers to different GPUs and streams micro-batches through the stages. It requires an additional nested `pipeline:` section: + +```yaml +distributed: + strategy: fsdp2 + dp_size: null + tp_size: 4 + pp_size: 4 + cp_size: 1 + activation_checkpointing: true + + pipeline: + pp_schedule: interleaved1f1b # pipeline schedule (1f1b or interleaved1f1b) + pp_microbatch_size: 1 # micro-batch size per pipeline step + layers_per_stage: 4 # how many layers each stage handles + scale_grads_in_schedule: false +``` + +| Key | Role | +|-----|------| +| `pp_schedule` | The micro-batch schedule. `1f1b` is simpler; `interleaved1f1b` overlaps compute and communication for better throughput. | +| `pp_microbatch_size` | Number of samples per micro-batch fed into the pipeline. Must satisfy: `local_batch_size ÷ pp_microbatch_size ≥ pp_size`. | +| `layers_per_stage` | Number of transformer layers each pipeline stage contains. If omitted, the framework splits layers evenly across `pp_size` stages. | + + +PP is available only when the model declares pipeline-parallel support and defines a `_pp_plan` that tells the framework how to split layers into stages. Do not assume every built-in model supports PP; check the model's declared capabilities before setting `pp_size > 1`. + + + +#### Context Parallelism + +CP splits the sequence across GPUs, which is useful for very long contexts that exceed single-GPU memory. Set `cp_size` to the target split factor: + +```yaml +distributed: + strategy: fsdp2 + dp_size: null + tp_size: 1 + cp_size: 2 +``` + + +When `cp_size > 1`, fused RoPE is automatically disabled. Some models also require the Transformer Engine (TE) attention backend for CP with packed sequences. The framework will raise an error with instructions if this applies. + + + +#### Expert Parallelism for MoE Models + +EP distributes MoE experts across GPUs. Set `ep_size` to the number of GPUs that share the full set of experts: + +```yaml +distributed: + strategy: fsdp2 + tp_size: 1 + cp_size: 1 + pp_size: 1 + ep_size: 8 + activation_checkpointing: true +``` + +EP only applies to Mixture-of-Experts models (such as Qwen3-MoE, Mixtral, and DeepSeek-V3). For dense models, leave `ep_size` at `1` or omit it. + +#### Combine Multiple Dimensions + +You can combine TP, PP, CP, and EP in a single config. For example, a large MoE model on a multi-node cluster might use: + +```yaml +distributed: + strategy: fsdp2 + dp_size: null + tp_size: 1 + cp_size: 2 + pp_size: 1 + ep_size: 4 + activation_checkpointing: true +``` + +When choosing a combination, keep these rules in mind: + +- **`world_size` must be evenly divisible by `pp_size × tp_size × cp_size`** so that the quotient is an integer `dp_size`. +- **`(dp_size × cp_size) % ep_size == 0`**, where EP shares the DP×CP groups. +- **TP within a node, PP across nodes** is the typical layout. TP requires fast NVLink bandwidth, whereas PP tolerates higher latency. +- **Start simple.** Use DP-only first. Add TP if the model does not fit on one GPU. Add PP for very large models. Add CP for long sequences. Add EP only for MoE architectures. + +## Next Steps + +- [Integrate Your Own Text Dataset](/datasets/text-dataset) to swap the SQuAD example for your own data. +- [Recipes and End-to-End Examples](/recipes-e2e-examples/overview) to browse the validated recipe guides available in NeMo AutoModel. +- [Dataset Overview: LLM, VLM, and Retrieval Datasets](/datasets/overview) to see all supported dataset types across LLM, VLM, and retrieval tasks. +- [Knowledge Distillation](/recipes-e2e-examples/knowledge-distillation) to distill a fine-tuned model into a smaller one. diff --git a/tests/functional_tests/checkpoint_robustness/parity_metrics.py b/tests/functional_tests/checkpoint_robustness/parity_metrics.py new file mode 100644 index 0000000000..a9d532cb3c --- /dev/null +++ b/tests/functional_tests/checkpoint_robustness/parity_metrics.py @@ -0,0 +1,323 @@ +# Copyright (c) 2026, NVIDIA CORPORATION. All rights reserved. +# +# 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. + +"""Shared full-logit metrics and numerical profiles for checkpoint parity.""" + +from __future__ import annotations + +import math +from dataclasses import asdict, dataclass +from typing import Literal + +import torch +import torch.nn.functional as F + +_ComparisonKind = Literal["same_implementation", "cross_framework", "cross_topology"] +_PARITY_COMPARISONS = {"source_load", "automodel_reload", "hf_reload", "cross_tp"} +_PARITY_OVERRIDE_METRICS = {"mean_kl", "p95_kl", "cosine_similarity"} + + +@dataclass(frozen=True) +class _ParityMetrics: + """Summary statistics for one full-logit comparison.""" + + token_count: int + vocab_size: int + mean_kl: float + p95_kl: float + max_kl: float + cosine_similarity: float + mean_absolute_logit_difference: float + max_absolute_logit_difference: float + + def to_dict(self) -> dict[str, int | float]: + """Return a JSON-serializable metric mapping.""" + return asdict(self) + + +@dataclass(frozen=True) +class _ParityThresholds: + """Mean KL, p95 KL, and cosine gates selected by one numerical profile.""" + + mean_kl: float + p95_kl: float + cosine_similarity: float + + def to_dict(self) -> dict[str, float]: + """Return a JSON-serializable threshold mapping.""" + return asdict(self) + + +_PARITY_PROFILES: dict[str, dict[_ComparisonKind, _ParityThresholds]] = { + "strict": { + "same_implementation": _ParityThresholds(mean_kl=1e-7, p95_kl=1e-6, cosine_similarity=0.999999), + "cross_framework": _ParityThresholds(mean_kl=1e-4, p95_kl=1e-3, cosine_similarity=0.9999), + "cross_topology": _ParityThresholds(mean_kl=1e-6, p95_kl=1e-5, cosine_similarity=0.99999), + }, + "standard": { + "same_implementation": _ParityThresholds(mean_kl=3e-3, p95_kl=1.2e-2, cosine_similarity=0.999), + "cross_framework": _ParityThresholds(mean_kl=6e-3, p95_kl=3e-2, cosine_similarity=0.998), + "cross_topology": _ParityThresholds(mean_kl=6e-3, p95_kl=3e-2, cosine_similarity=0.998), + }, + "relaxed": { + "same_implementation": _ParityThresholds(mean_kl=2e-2, p95_kl=5e-2, cosine_similarity=0.995), + "cross_framework": _ParityThresholds(mean_kl=2.5e-2, p95_kl=1e-1, cosine_similarity=0.99), + "cross_topology": _ParityThresholds(mean_kl=2e-2, p95_kl=5e-2, cosine_similarity=0.995), + }, +} + + +def _validate_logits(logits: torch.Tensor, *, chunk_tokens: int = 16) -> tuple[int, int]: + """Validate a complete logit tensor without allocating another full-size tensor. + + Args: + logits: Tensor of shape [..., vocab], with arbitrary leading token dimensions. + chunk_tokens: Number of flattened tokens checked together. + + Returns: + The flattened token count and vocabulary size. + """ + if logits.ndim < 2 or logits.shape[-1] <= 0: + raise ValueError(f"Expected logits of shape [..., vocab], got {tuple(logits.shape)}") + if chunk_tokens <= 0: + raise ValueError(f"chunk_tokens must be positive, got {chunk_tokens}") + + vocab_size = logits.shape[-1] + flattened_logits = logits.detach().reshape(-1, vocab_size) + token_count = flattened_logits.shape[0] + if token_count == 0: + raise ValueError("Logit tensor must contain at least one token") + for start in range(0, token_count, chunk_tokens): + end = min(start + chunk_tokens, token_count) + if not bool(torch.isfinite(flattened_logits[start:end]).all()): + raise ValueError(f"Logits contain non-finite values in flattened token range [{start}, {end})") + return token_count, vocab_size + + +def _compute_parity_metrics( + reference_logits: torch.Tensor, + candidate_logits: torch.Tensor, + *, + chunk_tokens: int = 16, +) -> _ParityMetrics: + """Compute bounded-memory statistics over two complete logit tensors. + + Args: + reference_logits: Reference tensor of shape [..., vocab], with arbitrary leading token dimensions. + candidate_logits: Candidate tensor of shape [..., vocab], matching ``reference_logits`` exactly. + chunk_tokens: Number of flattened tokens processed together. This bounds temporary softmax memory. + + Returns: + Scalar statistics over every token and vocabulary element. Returned values do not alias the inputs. + """ + if reference_logits.shape != candidate_logits.shape: + raise ValueError( + f"Logit shape mismatch: reference={tuple(reference_logits.shape)}, " + f"candidate={tuple(candidate_logits.shape)}" + ) + if reference_logits.ndim < 2 or reference_logits.shape[-1] <= 0: + raise ValueError(f"Expected logits of shape [..., vocab], got {tuple(reference_logits.shape)}") + if chunk_tokens <= 0: + raise ValueError(f"chunk_tokens must be positive, got {chunk_tokens}") + + vocab_size = reference_logits.shape[-1] + reference_tokens = reference_logits.detach().reshape(-1, vocab_size) + candidate_tokens = candidate_logits.detach().reshape(-1, vocab_size) + token_count = reference_tokens.shape[0] + if token_count == 0: + raise ValueError("Cannot compare empty logit tensors") + + kl_chunks: list[torch.Tensor] = [] + absolute_difference_sum = 0.0 + max_absolute_difference = 0.0 + dot_product = 0.0 + reference_squared_norm = 0.0 + candidate_squared_norm = 0.0 + + for start in range(0, token_count, chunk_tokens): + end = min(start + chunk_tokens, token_count) + reference_chunk = reference_tokens[start:end].float() + candidate_chunk = candidate_tokens[start:end].float() + if not bool(torch.isfinite(reference_chunk).all()): + raise ValueError(f"Reference logits contain non-finite values in flattened token range [{start}, {end})") + if not bool(torch.isfinite(candidate_chunk).all()): + raise ValueError(f"Candidate logits contain non-finite values in flattened token range [{start}, {end})") + + reference_log_probs = F.log_softmax(reference_chunk, dim=-1) + candidate_log_probs = F.log_softmax(candidate_chunk, dim=-1) + reference_probs = reference_log_probs.exp() + token_kl = (reference_probs * (reference_log_probs - candidate_log_probs)).sum(dim=-1) + kl_chunks.append(token_kl.cpu()) + + absolute_difference = (reference_chunk - candidate_chunk).abs() + absolute_difference_sum += absolute_difference.sum(dtype=torch.float64).item() + max_absolute_difference = max(max_absolute_difference, absolute_difference.max().item()) + dot_product += (reference_chunk * candidate_chunk).sum(dtype=torch.float64).item() + reference_squared_norm += reference_chunk.square().sum(dtype=torch.float64).item() + candidate_squared_norm += candidate_chunk.square().sum(dtype=torch.float64).item() + + per_token_kl = torch.cat(kl_chunks) + if not bool(torch.isfinite(per_token_kl).all()): + raise ValueError("KL divergence contains non-finite values") + + norm_product = math.sqrt(reference_squared_norm * candidate_squared_norm) + if norm_product == 0.0: + cosine_similarity = 1.0 if max_absolute_difference == 0.0 else 0.0 + else: + cosine_similarity = dot_product / norm_product + + return _ParityMetrics( + token_count=token_count, + vocab_size=vocab_size, + mean_kl=per_token_kl.mean().item(), + p95_kl=torch.quantile(per_token_kl, 0.95).item(), + max_kl=per_token_kl.max().item(), + cosine_similarity=cosine_similarity, + mean_absolute_logit_difference=absolute_difference_sum / reference_logits.numel(), + max_absolute_logit_difference=max_absolute_difference, + ) + + +def _resolve_parity_thresholds(profile: str, comparison_kind: _ComparisonKind) -> _ParityThresholds: + """Resolve one named profile for the requested comparison kind.""" + if profile not in _PARITY_PROFILES: + raise ValueError(f"Unknown parity tolerance profile {profile!r}; expected one of {sorted(_PARITY_PROFILES)}") + return _PARITY_PROFILES[profile][comparison_kind] + + +def _normalize_parity_profile_overrides(raw_overrides: object) -> dict[str, str]: + """Validate and normalize optional per-comparison profile overrides.""" + if raw_overrides is None: + return {} + if not isinstance(raw_overrides, dict): + raise ValueError("parity_tolerance_profile_overrides must be a mapping") + + non_string_comparisons = [repr(comparison) for comparison in raw_overrides if not isinstance(comparison, str)] + if non_string_comparisons: + raise ValueError( + "parity_tolerance_profile_overrides comparison names must be strings, got " + + ", ".join(non_string_comparisons) + ) + unknown_comparisons = sorted(set(raw_overrides) - _PARITY_COMPARISONS) + if unknown_comparisons: + raise ValueError( + "Unknown parity_tolerance_profile_overrides comparisons: " + f"{', '.join(unknown_comparisons)}; expected one of {sorted(_PARITY_COMPARISONS)}" + ) + + normalized: dict[str, str] = {} + for comparison, profile in raw_overrides.items(): + if not isinstance(profile, str): + raise ValueError(f"parity_tolerance_profile_overrides.{comparison} must be a profile name") + _resolve_parity_thresholds(profile, "same_implementation") + normalized[comparison] = profile + return normalized + + +def _select_parity_profile(default_profile: str, raw_overrides: object, comparison: str) -> str: + """Select one comparison profile, falling back to the global profile.""" + _resolve_parity_thresholds(default_profile, "same_implementation") + if comparison not in _PARITY_COMPARISONS: + raise ValueError(f"Unknown parity comparison {comparison!r}; expected one of {sorted(_PARITY_COMPARISONS)}") + return _normalize_parity_profile_overrides(raw_overrides).get(comparison, default_profile) + + +def _apply_parity_threshold_overrides( + thresholds: _ParityThresholds, + *, + mean_kl: float | None = None, + p95_kl: float | None = None, + cosine_similarity: float | None = None, +) -> _ParityThresholds: + """Replace selected profile gates with explicit model-specific values.""" + for name, value in (("mean_kl", mean_kl), ("p95_kl", p95_kl)): + if value is not None and (not math.isfinite(value) or value < 0): + raise ValueError(f"{name} threshold override must be finite and non-negative, got {value}") + if cosine_similarity is not None and (not math.isfinite(cosine_similarity) or not -1.0 <= cosine_similarity <= 1.0): + raise ValueError( + f"cosine_similarity threshold override must be finite and between -1 and 1, got {cosine_similarity}" + ) + return _ParityThresholds( + mean_kl=thresholds.mean_kl if mean_kl is None else mean_kl, + p95_kl=thresholds.p95_kl if p95_kl is None else p95_kl, + cosine_similarity=thresholds.cosine_similarity if cosine_similarity is None else cosine_similarity, + ) + + +def _normalize_parity_threshold_overrides(raw_overrides: object) -> dict[str, dict[str, float]]: + """Validate and normalize optional per-comparison profile threshold overrides.""" + if raw_overrides is None: + return {} + if not isinstance(raw_overrides, dict): + raise ValueError("parity_threshold_overrides must be a mapping") + + non_string_comparisons = [repr(comparison) for comparison in raw_overrides if not isinstance(comparison, str)] + if non_string_comparisons: + raise ValueError( + "parity_threshold_overrides comparison names must be strings, got " + ", ".join(non_string_comparisons) + ) + unknown_comparisons = sorted(set(raw_overrides) - _PARITY_COMPARISONS) + if unknown_comparisons: + raise ValueError( + "Unknown parity_threshold_overrides comparisons: " + f"{', '.join(unknown_comparisons)}; expected one of {sorted(_PARITY_COMPARISONS)}" + ) + + normalized: dict[str, dict[str, float]] = {} + for comparison, raw_metrics in raw_overrides.items(): + if not isinstance(raw_metrics, dict): + raise ValueError(f"parity_threshold_overrides.{comparison} must be a mapping") + non_string_metrics = [repr(metric) for metric in raw_metrics if not isinstance(metric, str)] + if non_string_metrics: + raise ValueError( + f"parity_threshold_overrides.{comparison} metric names must be strings, got " + + ", ".join(non_string_metrics) + ) + unknown_metrics = sorted(set(raw_metrics) - _PARITY_OVERRIDE_METRICS) + if unknown_metrics: + raise ValueError( + f"Unknown parity_threshold_overrides.{comparison} metrics: {', '.join(unknown_metrics)}; " + f"expected one of {sorted(_PARITY_OVERRIDE_METRICS)}" + ) + + metrics: dict[str, float] = {} + for metric, raw_value in raw_metrics.items(): + if isinstance(raw_value, bool) or not isinstance(raw_value, (int, float)): + raise ValueError(f"parity_threshold_overrides.{comparison}.{metric} must be numeric") + metrics[metric] = float(raw_value) + _apply_parity_threshold_overrides( + _ParityThresholds(mean_kl=0.0, p95_kl=0.0, cosine_similarity=0.0), + mean_kl=metrics.get("mean_kl"), + p95_kl=metrics.get("p95_kl"), + cosine_similarity=metrics.get("cosine_similarity"), + ) + normalized[comparison] = metrics + return normalized + + +def _parity_failures( + metrics: _ParityMetrics, + thresholds: _ParityThresholds, +) -> tuple[str, ...]: + """Return failed mean-KL, p95-KL, and cosine gates for a named profile.""" + failures: list[str] = [] + if metrics.mean_kl > thresholds.mean_kl: + failures.append(f"mean KL {metrics.mean_kl:.6e} > profile threshold {thresholds.mean_kl:.6e}") + if metrics.p95_kl > thresholds.p95_kl: + failures.append(f"p95 KL {metrics.p95_kl:.6e} > profile threshold {thresholds.p95_kl:.6e}") + if metrics.cosine_similarity < thresholds.cosine_similarity: + failures.append( + f"cosine similarity {metrics.cosine_similarity:.8f} < profile threshold {thresholds.cosine_similarity:.8f}" + ) + return tuple(failures) diff --git a/tests/functional_tests/checkpoint_robustness/resume_trajectory.py b/tests/functional_tests/checkpoint_robustness/resume_trajectory.py index 532ee87434..55d59298ae 100644 --- a/tests/functional_tests/checkpoint_robustness/resume_trajectory.py +++ b/tests/functional_tests/checkpoint_robustness/resume_trajectory.py @@ -310,7 +310,12 @@ def _resume_plan_from_config(cfg: object, *, continuation_steps: int = 3) -> _Re def _configure_uninterrupted_run(cfg: object, plan: _ResumePlan) -> None: """Extend Phase 1 while preserving its original LR schedule and checkpoint boundary.""" cfg.step_scheduler.max_steps = plan.final_max_steps + # ``max_steps`` is only a cap: a finite dataloader can stop earlier when the + # configured epochs are exhausted. Allow one epoch per requested step so a + # non-empty dataloader always reaches the shared checkpoint and continuation. + cfg.step_scheduler.num_epochs = plan.final_max_steps cfg.step_scheduler.ckpt_every_steps = plan.boundary_step + cfg.step_scheduler.save_checkpoint_every_epoch = False cfg.checkpoint.save_consolidated = "final" if hasattr(cfg, "lr_scheduler") and cfg.lr_scheduler is not None: cfg.lr_scheduler.lr_decay_steps = plan.boundary_step @@ -319,7 +324,9 @@ def _configure_uninterrupted_run(cfg: object, plan: _ResumePlan) -> None: def _configure_resumed_run(cfg: object, plan: _ResumePlan, checkpoint_path: Path) -> None: """Restore the boundary checkpoint into an output directory separate from the reference branch.""" cfg.step_scheduler.max_steps = plan.final_max_steps + cfg.step_scheduler.num_epochs = plan.final_max_steps cfg.step_scheduler.ckpt_every_steps = plan.boundary_step + cfg.step_scheduler.save_checkpoint_every_epoch = False if hasattr(cfg, "lr_scheduler") and cfg.lr_scheduler is not None: cfg.lr_scheduler.lr_decay_steps = plan.boundary_step cfg.checkpoint.restore_from = str(checkpoint_path) @@ -327,6 +334,11 @@ def _configure_resumed_run(cfg: object, plan: _ResumePlan, checkpoint_path: Path cfg.checkpoint.save_consolidated = False +def _disable_checkpoint_saves_after_restore(trainer: object) -> None: + """Disable new checkpoint writes after the resume checkpoint has loaded.""" + trainer.checkpointer.config.enabled = False + + def _checkpoint_for_completed_steps(plan: _ResumePlan, completed_steps: int) -> Path: """Locate the checkpoint written after exactly ``completed_steps`` optimizer steps.""" checkpoint_step = completed_steps - 1 diff --git a/tests/functional_tests/checkpoint_robustness/test_checkpoint_robustness_biencoder.py b/tests/functional_tests/checkpoint_robustness/test_checkpoint_robustness_biencoder.py index f084c7c437..ef807a89e2 100644 --- a/tests/functional_tests/checkpoint_robustness/test_checkpoint_robustness_biencoder.py +++ b/tests/functional_tests/checkpoint_robustness/test_checkpoint_robustness_biencoder.py @@ -19,8 +19,10 @@ KL divergence. Launch: torchrun --nproc-per-node= -m --config - [--cosine_threshold ] [--hf_cosine_threshold ] - [--check_hf_reload] [--check_resume] + [--parity_tolerance_profile ] + [--parity_tolerance_profile_overrides ] + [--parity_threshold_overrides ] + [--skip_hf_reload] [--skip_resume] [--resume_tolerance_profile ] [--resume_first_loss_threshold ] [--resume_loss_threshold ] """ @@ -38,11 +40,19 @@ from nemo_automodel.components.config._arg_parser import parse_args_and_load_config from nemo_automodel.recipes.retrieval.train_bi_encoder import TrainBiEncoderRecipe +from tests.functional_tests.checkpoint_robustness.parity_metrics import ( + _apply_parity_threshold_overrides, + _normalize_parity_profile_overrides, + _normalize_parity_threshold_overrides, + _resolve_parity_thresholds, + _select_parity_profile, +) from tests.functional_tests.checkpoint_robustness.resume_trajectory import ( _checkpoint_for_completed_steps, _checkpoint_state_snapshot, _configure_resumed_run, _configure_uninterrupted_run, + _disable_checkpoint_saves_after_restore, _gather_rank_failures, _load_reference_trajectory, _persist_reference_trajectory, @@ -63,20 +73,30 @@ # Default test sentence for embedding extraction _DEFAULT_PROMPT = "The quick brown fox jumps over the lazy dog" - - -def _extract_custom_args(argv: list[str]) -> tuple[dict[str, str | bool], list[str]]: +_REMOVED_BIENCODER_FIELDS = { + "automodel_reload_cosine_threshold", + "automodel_reload_mean_kl_threshold", + "automodel_reload_p95_kl_threshold", + "check_hf_reload", + "check_resume", + "cosine_threshold", + "hf_cosine_threshold", +} + + +def _extract_custom_args(argv: list[str]) -> tuple[dict[str, object], list[str]]: """Separate test-specific CLI flags from config parser arguments.""" custom_keys = { - "--cosine_threshold", - "--hf_cosine_threshold", + "--parity_threshold_overrides", + "--parity_tolerance_profile", + "--parity_tolerance_profile_overrides", "--training_reproducibility_loss_threshold", "--resume_first_loss_threshold", "--resume_loss_threshold", "--resume_tolerance_profile", } - boolean_keys = {"--check_hf_reload", "--check_resume"} - custom: dict[str, str | bool] = {} + boolean_keys = {"--skip_hf_reload", "--skip_resume"} + custom: dict[str, object] = {} remaining: list[str] = [] i = 0 while i < len(argv): @@ -99,18 +119,61 @@ def _extract_custom_args(argv: list[str]) -> tuple[dict[str, str | bool], list[s import yaml with open(config_path) as f: - raw_cfg = yaml.safe_load(f) + raw_cfg = yaml.safe_load(f) or {} ci_robustness = raw_cfg.get("ci", {}).get("checkpoint_robustness") or {} + removed_fields = sorted(_REMOVED_BIENCODER_FIELDS & ci_robustness.keys()) + if removed_fields: + raise ValueError( + "Removed retrieval checkpoint-robustness fields are not supported: " + ", ".join(removed_fields) + ) for cli_key in custom_keys | boolean_keys: key = cli_key.lstrip("-") if key in custom or key not in ci_robustness: continue value = ci_robustness[key] + if key in {"parity_threshold_overrides", "parity_tolerance_profile_overrides"}: + custom[key] = value + continue if isinstance(value, bool): - if value: - custom[key] = True + custom[key] = value else: custom[key] = str(value) + raw_threshold_overrides = custom.get("parity_threshold_overrides") + if isinstance(raw_threshold_overrides, str): + import yaml + + raw_threshold_overrides = yaml.safe_load(raw_threshold_overrides) + if raw_threshold_overrides is not None: + threshold_overrides = _normalize_parity_threshold_overrides(raw_threshold_overrides) + unsupported_comparisons = sorted(set(threshold_overrides) - {"automodel_reload", "hf_reload"}) + if unsupported_comparisons: + raise ValueError( + "Retrieval parity_threshold_overrides supports only automodel_reload and hf_reload, got " + + ", ".join(unsupported_comparisons) + ) + for comparison, overrides in threshold_overrides.items(): + unsupported_metrics = sorted(set(overrides) - {"cosine_similarity"}) + if unsupported_metrics: + raise ValueError( + f"Retrieval parity_threshold_overrides.{comparison} supports only cosine_similarity, got " + + ", ".join(unsupported_metrics) + ) + custom["parity_threshold_overrides"] = threshold_overrides + raw_profile_overrides = custom.get("parity_tolerance_profile_overrides") + if isinstance(raw_profile_overrides, str): + import yaml + + raw_profile_overrides = yaml.safe_load(raw_profile_overrides) + if raw_profile_overrides is not None: + profile_overrides = _normalize_parity_profile_overrides(raw_profile_overrides) + unsupported_comparisons = sorted(set(profile_overrides) - {"automodel_reload", "hf_reload"}) + if unsupported_comparisons: + raise ValueError( + "Retrieval parity_tolerance_profile_overrides supports only automodel_reload and hf_reload, got " + + ", ".join(unsupported_comparisons) + ) + custom["parity_tolerance_profile_overrides"] = profile_overrides + _resolve_parity_thresholds(str(custom.get("parity_tolerance_profile", "standard")), "same_implementation") return custom, remaining @@ -168,10 +231,32 @@ def test_checkpoint_robustness_biencoder(): """Train biencoder -> checkpoint -> reload from consolidated, compare embeddings.""" custom_args, config_argv = _extract_custom_args(sys.argv[1:]) sys.argv = [sys.argv[0]] + config_argv - cosine_threshold = float(custom_args.get("cosine_threshold", "0.999")) - hf_cosine_threshold = float(custom_args.get("hf_cosine_threshold", "0.999")) - check_hf_reload = bool(custom_args.get("check_hf_reload", False)) - check_resume = bool(custom_args.get("check_resume", False)) + parity_profile = str(custom_args.get("parity_tolerance_profile", "standard")) + threshold_overrides = _normalize_parity_threshold_overrides(custom_args.get("parity_threshold_overrides")) + automodel_overrides = threshold_overrides.get("automodel_reload", {}) + hf_overrides = threshold_overrides.get("hf_reload", {}) + automodel_thresholds = _resolve_parity_thresholds( + _select_parity_profile( + parity_profile, + custom_args.get("parity_tolerance_profile_overrides"), + "automodel_reload", + ), + "same_implementation", + ) + automodel_thresholds = _apply_parity_threshold_overrides( + automodel_thresholds, + cosine_similarity=automodel_overrides.get("cosine_similarity"), + ) + hf_thresholds = _resolve_parity_thresholds( + _select_parity_profile(parity_profile, custom_args.get("parity_tolerance_profile_overrides"), "hf_reload"), + "cross_framework", + ) + hf_thresholds = _apply_parity_threshold_overrides( + hf_thresholds, + cosine_similarity=hf_overrides.get("cosine_similarity"), + ) + hf_reload_enabled = not bool(custom_args.get("skip_hf_reload", False)) + resume_enabled = not bool(custom_args.get("skip_resume", False)) training_reproducibility_loss_threshold = float(custom_args.get("training_reproducibility_loss_threshold", "5e-2")) resume_tolerance = _resolve_resume_loss_tolerance( str(custom_args.get("resume_tolerance_profile", "standard")), @@ -180,11 +265,11 @@ def test_checkpoint_robustness_biencoder(): ) # ------------------------------------------------------------------ - # Phase 1: Train biencoder and checkpoint + # Phase 1: Train, save, and capture reference artifacts # ------------------------------------------------------------------ torch.cuda.reset_peak_memory_stats() cfg = parse_args_and_load_config() - resume_plan = _resume_plan_from_config(cfg) if check_resume else None + resume_plan = _resume_plan_from_config(cfg) if resume_enabled else None if resume_plan is not None: _configure_uninterrupted_run(cfg, resume_plan) trainer = TrainBiEncoderRecipe(cfg) @@ -223,22 +308,20 @@ def test_checkpoint_robustness_biencoder(): if _rank0(): print(f"\n[Memory] Peak VRAM: {peak_vram_gb:.2f} GB, Peak CPU RSS: {peak_cpu_gb:.2f} GB") - # ------------------------------------------------------------------ - # Phase 2: Capture reference embeddings before teardown - # ------------------------------------------------------------------ + # Capture Phase 1 reference embeddings before teardown. device = next(trainer.model_parts[0].parameters()).device tokenizer = trainer.tokenizer reference_embeddings = _get_embeddings(trainer.model_parts[0], tokenizer, _DEFAULT_PROMPT, device) hf_reference_query = None hf_reference_document = None - if check_hf_reload: + if hf_reload_enabled: hf_reference_query, hf_reference_document = _get_hf_style_embeddings(trainer.model_parts[0], _DEFAULT_PROMPT) if _rank0(): - print(f"\n[Phase 2] Reference embedding shape: {reference_embeddings.shape}") - print(f"[Phase 2] Reference embedding norm: {reference_embeddings.norm().item():.6f}") + print(f"\n[Phase 1] Reference embedding shape: {reference_embeddings.shape}") + print(f"[Phase 1] Reference embedding norm: {reference_embeddings.norm().item():.6f}") # ------------------------------------------------------------------ - # Phase 3: Reload from consolidated checkpoint, compare embeddings + # Phase 2: Reload the consolidated checkpoint with AutoModel and compare embeddings # ------------------------------------------------------------------ checkpoint_dir = Path(cfg.checkpoint.checkpoint_dir) if resume_plan is not None: @@ -268,12 +351,12 @@ def test_checkpoint_robustness_biencoder(): cosine_sim = _cosine_similarity(reference_embeddings, restored_embeddings) if _rank0(): print( - f"\n[Phase 3] Cosine similarity (original vs consolidated): " - f"{cosine_sim:.6f} (threshold: {cosine_threshold})" + f"\n[Phase 2] AutoModel reload cosine similarity: {cosine_sim:.6f} " + f"(profile: {parity_profile}, threshold: {automodel_thresholds.cosine_similarity})" ) - assert cosine_sim >= cosine_threshold, ( + assert cosine_sim >= automodel_thresholds.cosine_similarity, ( f"Cosine similarity between original and consolidated embeddings too low: " - f"{cosine_sim:.6f} < threshold {cosine_threshold}" + f"{cosine_sim:.6f} < threshold {automodel_thresholds.cosine_similarity}" ) del restored_trainer @@ -281,9 +364,9 @@ def test_checkpoint_robustness_biencoder(): _barrier() # ------------------------------------------------------------------ - # Phase 4 (optional): Reload with vanilla Hugging Face AutoModel + # Phase 3: Reload the consolidated checkpoint with vanilla Hugging Face AutoModel # ------------------------------------------------------------------ - if check_hf_reload: + if hf_reload_enabled: hf_reload_sync_paths = _prepare_hf_reload_sync(cfg) hf_reload_error = None if _rank0(): @@ -302,19 +385,19 @@ def test_checkpoint_robustness_biencoder(): query_cosine_sim = _cosine_similarity(hf_reference_query, hf_query) document_cosine_sim = _cosine_similarity(hf_reference_document, hf_document) print( - f"\n[Phase 4] HF reload query cosine similarity: {query_cosine_sim:.6f}; " + f"\n[Phase 3] HF reload query cosine similarity: {query_cosine_sim:.6f}; " f"image-document cosine similarity: {document_cosine_sim:.6f} " - f"(threshold: {hf_cosine_threshold})" + f"(profile: {parity_profile}, threshold: {hf_thresholds.cosine_similarity})" ) - if query_cosine_sim < hf_cosine_threshold: + if query_cosine_sim < hf_thresholds.cosine_similarity: hf_reload_error = ( f"HF-reloaded query embedding cosine similarity too low: " - f"{query_cosine_sim:.6f} < threshold {hf_cosine_threshold}" + f"{query_cosine_sim:.6f} < threshold {hf_thresholds.cosine_similarity}" ) - if document_cosine_sim < hf_cosine_threshold: + if document_cosine_sim < hf_thresholds.cosine_similarity: document_error = ( f"HF-reloaded image-document embedding cosine similarity too low: " - f"{document_cosine_sim:.6f} < threshold {hf_cosine_threshold}" + f"{document_cosine_sim:.6f} < threshold {hf_thresholds.cosine_similarity}" ) hf_reload_error = "\n".join(filter(None, (hf_reload_error, document_error))) del hf_model @@ -325,9 +408,9 @@ def test_checkpoint_robustness_biencoder(): assert hf_reload_error is None, hf_reload_error # ------------------------------------------------------------------ - # Phase 5 (optional): restore the exact Phase 1 boundary and replay its continuation. + # Phase 4: restore the exact Phase 1 boundary and replay its continuation. # ------------------------------------------------------------------ - if check_resume: + if resume_enabled: assert resume_plan is not None reference_trajectory = _load_reference_trajectory(resume_plan) checkpoint_path = _checkpoint_for_completed_steps(resume_plan, resume_plan.boundary_step) @@ -335,6 +418,7 @@ def test_checkpoint_robustness_biencoder(): _configure_resumed_run(cfg, resume_plan, checkpoint_path) resume_trainer = TrainBiEncoderRecipe(cfg) resume_trainer.setup() + _disable_checkpoint_saves_after_restore(resume_trainer) restored_state = _checkpoint_state_snapshot(resume_trainer, state_is_being_saved=False) local_failure = _restored_state_mismatch(reference_trajectory["boundary_state"], restored_state) failure_message = _gather_rank_failures(local_failure, check="restored_state") diff --git a/tests/functional_tests/checkpoint_robustness/test_checkpoint_robustness_llm.py b/tests/functional_tests/checkpoint_robustness/test_checkpoint_robustness_llm.py index 8cacfa77a9..a18753bb8b 100644 --- a/tests/functional_tests/checkpoint_robustness/test_checkpoint_robustness_llm.py +++ b/tests/functional_tests/checkpoint_robustness/test_checkpoint_robustness_llm.py @@ -14,18 +14,15 @@ """Train, checkpoint, and validate AutoModel and vanilla-HF reloads. -Launch: torchrun --nproc-per-node= -m --config - [--isolated_phase ] - [--kl_threshold ] [--hf_kl_threshold ] - [--cross_tp_size ] [--cross_tp_kl_threshold ] - [--tokenizer_name ] - [--source_load_kl_threshold ] [--source_load_mean_kl_threshold ] - [--check_source_load_parity] [--check_fused_qkv_keys] [--check_phantom_keys] [--check_resume] - [--resume_tolerance_profile ] - [--resume_first_loss_threshold ] [--resume_loss_threshold ] - [--skip_automodel_logit_parity] [--skip_hf_logit_parity] [--hf_adapter_ignored_key_prefix ] - [--hf_source_post_load_dequantize] - [--max_vram_gb ] [--max_cpu_gb ] +Launch with ``torchrun --nproc-per-node= -m --config ``. + +The CI launcher runs phases in isolated processes by default through ``--isolated_phase``. Accepted phase names are +``source_load_reference``, ``source_load_parity``, ``train_and_save``, ``automodel_reload``, ``hf_reload``, ``resume``, +and ``cross_tp_reload``. Direct invocation without ``--isolated_phase`` retains the compatibility single-process +lifecycle. + +See ``tests/ci_tests/README.md#checkpoint-robustness`` for the public phase contract, tolerance profiles, and supported +recipe controls. """ from __future__ import annotations @@ -34,17 +31,17 @@ import hashlib import inspect import json -import math import os import sys import time import traceback from collections.abc import Callable from contextlib import AbstractContextManager, contextmanager, nullcontext +from dataclasses import dataclass from datetime import timedelta from functools import wraps from pathlib import Path -from typing import TYPE_CHECKING +from typing import TYPE_CHECKING, Literal if TYPE_CHECKING: from nemo_automodel.recipes.base_recipe import BaseRecipe @@ -52,7 +49,6 @@ import datasets import torch import torch.distributed as dist -import torch.nn.functional as F from torch.distributed.tensor import DTensor from nemo_automodel.components.checkpoint.checkpointing import ( @@ -62,11 +58,22 @@ from nemo_automodel.components.config._arg_parser import parse_args_and_load_config from nemo_automodel.components.config.loader import ConfigNode from nemo_automodel.shared.utils import dtype_from_str +from tests.functional_tests.checkpoint_robustness.parity_metrics import ( + _apply_parity_threshold_overrides, + _compute_parity_metrics, + _normalize_parity_profile_overrides, + _normalize_parity_threshold_overrides, + _parity_failures, + _resolve_parity_thresholds, + _select_parity_profile, + _validate_logits, +) from tests.functional_tests.checkpoint_robustness.resume_trajectory import ( _checkpoint_for_completed_steps, _checkpoint_state_snapshot, _configure_resumed_run, _configure_uninterrupted_run, + _disable_checkpoint_saves_after_restore, _gather_rank_failures, _load_reference_trajectory, _persist_reference_trajectory, @@ -82,47 +89,79 @@ datasets.disable_caching() -# Llama token IDs for "The quick brown fox jumps over the lazy dog" -_DEFAULT_INPUT_IDS = [791, 4996, 14198, 39935, 35308, 927, 279, 16053, 5679] -_DEFAULT_PROMPT = "The quick brown fox jumps over the lazy dog" - - -def _extract_custom_args(argv): +_PARITY_DOCUMENT_PATH = Path(__file__).with_name("parity_document.mdx") +_PARITY_DOCUMENT_SHA256 = "8f734b2ee925ab82afb56dfa3a512108b70d3c54a2489f7978a036420da34cdb" # pragma: allowlist secret +_REMOVED_CHECKPOINT_ROBUSTNESS_FIELDS = { + "automodel_reload_cosine_threshold", + "automodel_reload_mean_kl_threshold", + "automodel_reload_p95_kl_threshold", + "check_hf_reload", + "check_resume", + "check_source_load_parity", + "cosine_threshold", + "cross_tp_kl_threshold", + "hf_cosine_threshold", + "hf_kl_threshold", + "kl_threshold", + "no_check_resume", + "skip_automodel_logit_parity", + "skip_hf_logit_parity", + "source_load_cosine_threshold", + "source_load_kl_threshold", + "source_load_mean_kl_threshold", +} + + +@dataclass(frozen=True) +class _LogitParityPolicy: + """Configuration and enforcement state for one full-logit comparison.""" + + phase: str + comparison: str + comparison_kind: Literal["same_implementation", "cross_framework", "cross_topology"] + profile: str + enforce: bool = True + mean_kl_threshold_override: float | None = None + p95_kl_threshold_override: float | None = None + cosine_threshold_override: float | None = None + + +def _extract_custom_args(argv: list[str]) -> tuple[dict[str, object], list[str]]: """Separate test-specific CLI flags from config parser arguments.""" custom_keys = { - "--kl_threshold", - "--hf_kl_threshold", "--isolated_phase", "--cross_tp_size", - "--cross_tp_kl_threshold", "--experts_implementation", "--hf_adapter_ignored_key_prefix", + "--hf_device_map_cpu_max_memory_gib", "--hf_device_map_max_memory_gib", + "--hf_reload_timeout_seconds", "--tokenizer_name", "--max_vram_gb", "--max_cpu_gb", "--training_reproducibility_loss_threshold", + "--parity_sequence_length", + "--parity_threshold_overrides", + "--parity_tolerance_profile", + "--parity_tolerance_profile_overrides", "--resume_first_loss_threshold", "--resume_loss_threshold", "--resume_tolerance_profile", - "--source_load_cosine_threshold", - "--source_load_kl_threshold", - "--source_load_mean_kl_threshold", } boolean_keys = { "--trust_remote_code", - "--check_source_load_parity", "--check_fused_qkv_keys", "--check_phantom_keys", - "--check_resume", "--hf_device_map_auto", "--hf_source_post_load_dequantize", - "--no_check_resume", + "--skip_resume", + "--skip_source_load_parity", + "--skip_source_load_logit_parity", "--skip_hf_reload", - "--skip_automodel_logit_parity", - "--skip_hf_logit_parity", + "--skip_automodel_reload_logit_parity", + "--skip_hf_reload_logit_parity", } - custom = {} + custom: dict[str, object] = {} remaining = [] i = 0 while i < len(argv): @@ -135,6 +174,7 @@ def _extract_custom_args(argv): else: remaining.append(argv[i]) i += 1 + cli_custom_keys = set(custom) # Read ci.checkpoint_robustness from the YAML config as defaults. # CLI args take precedence over YAML values. @@ -143,36 +183,91 @@ def _extract_custom_args(argv): if arg == "--config" and j + 1 < len(remaining): config_path = remaining[j + 1] break + ci_robustness: dict = {} if config_path: import yaml with open(config_path) as f: - raw_cfg = yaml.safe_load(f) + raw_cfg = yaml.safe_load(f) or {} ci_robustness = raw_cfg.get("ci", {}).get("checkpoint_robustness") or {} - no_check_resume = ci_robustness.pop("no_check_resume", False) - if no_check_resume: - custom["no_check_resume"] = True + removed_fields = sorted(_REMOVED_CHECKPOINT_ROBUSTNESS_FIELDS & ci_robustness.keys()) + if removed_fields: + raise ValueError("Removed checkpoint-robustness fields are not supported: " + ", ".join(removed_fields)) + default_on_control_keys = { + "parity_threshold_overrides", + "parity_tolerance_profile_overrides", + "skip_resume", + "skip_source_load_parity", + } for k, v in ci_robustness.items(): + if k in default_on_control_keys: + continue if k not in custom: if "." in k: # Dotted keys are config overrides (e.g. distributed.tp_size), # route them to the config parser instead of the custom dict. remaining.extend([f"--{k}", str(v)]) - elif isinstance(v, bool) and v: - custom[k] = True + elif isinstance(v, bool) and (v or k == "trust_remote_code"): + # ``false`` is meaningful for trust_remote_code: it must be + # able to override a recipe model that normally uses remote code. + custom[k] = v elif not isinstance(v, bool): custom[k] = str(v) - # Enable check_resume by default unless no_check_resume is set - if not no_check_resume and "check_resume" not in custom: - custom["check_resume"] = True + + raw_threshold_overrides = custom.get("parity_threshold_overrides") + if raw_threshold_overrides is None: + raw_threshold_overrides = ci_robustness.get("parity_threshold_overrides") + if isinstance(raw_threshold_overrides, str): + import yaml + + raw_threshold_overrides = yaml.safe_load(raw_threshold_overrides) + if raw_threshold_overrides is not None: + custom["parity_threshold_overrides"] = _normalize_parity_threshold_overrides(raw_threshold_overrides) + + raw_profile_overrides = custom.get("parity_tolerance_profile_overrides") + if raw_profile_overrides is None: + raw_profile_overrides = ci_robustness.get("parity_tolerance_profile_overrides") + if isinstance(raw_profile_overrides, str): + import yaml + + raw_profile_overrides = yaml.safe_load(raw_profile_overrides) + if raw_profile_overrides is not None: + custom["parity_tolerance_profile_overrides"] = _normalize_parity_profile_overrides(raw_profile_overrides) + + if "skip_source_load_parity" in cli_custom_keys: + source_load_parity_enabled = False + elif "skip_source_load_parity" in ci_robustness: + source_load_parity_enabled = not bool(ci_robustness["skip_source_load_parity"]) + else: + source_load_parity_enabled = True + custom["source_load_parity_enabled"] = source_load_parity_enabled + if not source_load_parity_enabled: + custom["skip_source_load_parity"] = True + + if "skip_resume" in cli_custom_keys: + resume_enabled = False + elif "skip_resume" in ci_robustness: + resume_enabled = not bool(ci_robustness["skip_resume"]) + else: + resume_enabled = True + custom["resume_enabled"] = resume_enabled + if not resume_enabled: + custom["skip_resume"] = True + + parity_sequence_length = int(custom.get("parity_sequence_length", "2048")) + if parity_sequence_length <= 0: + raise ValueError(f"parity_sequence_length must be positive, got {parity_sequence_length}") + if "hf_reload_timeout_seconds" in custom and int(custom["hf_reload_timeout_seconds"]) <= 0: + raise ValueError("hf_reload_timeout_seconds must be positive") + _resolve_parity_thresholds(str(custom.get("parity_tolerance_profile", "standard")), "same_implementation") return custom, remaining def _get_input_ids(tokenizer_name: str | None) -> list[int]: - """Return input IDs for the test prompt, using dynamic tokenization if tokenizer_name is set.""" + """Tokenize the repository's long-form finetuning guide for parity testing.""" if tokenizer_name is None: - return _DEFAULT_INPUT_IDS + raise ValueError("tokenizer_name is required to tokenize the checkpoint parity document") from nemo_automodel import NeMoAutoTokenizer tokenizer = NeMoAutoTokenizer.from_pretrained( @@ -180,7 +275,22 @@ def _get_input_ids(tokenizer_name: str | None) -> list[int]: trust_remote_code=True, local_files_only=os.environ.get("HF_HUB_OFFLINE", "0") == "1", ) - return tokenizer.encode(_DEFAULT_PROMPT, add_special_tokens=False) + return tokenizer.encode(_get_parity_document(), add_special_tokens=False) + + +def _get_parity_document() -> str: + """Load and validate the fixed long-form document shared by LLM and VLM parity tests.""" + try: + document_bytes = _PARITY_DOCUMENT_PATH.read_bytes() + except OSError as exc: + raise RuntimeError(f"Unable to load checkpoint parity document: {_PARITY_DOCUMENT_PATH}") from exc + document_sha256 = hashlib.sha256(document_bytes).hexdigest() + if document_sha256 != _PARITY_DOCUMENT_SHA256: + raise RuntimeError( + "Checkpoint parity document changed unexpectedly: " + f"expected sha256={_PARITY_DOCUMENT_SHA256}, got sha256={document_sha256}" + ) + return document_bytes.decode("utf-8") def _load_hf_config( @@ -233,6 +343,42 @@ def _load_hf_fp8_dequantized_config( return config +def _repair_legacy_partial_rotary_config(config) -> bool: + """Restore a legacy partial-rotary spec dropped by newer Transformers configs. + + Checkpoints such as MiniMax-M2.* express partial RoPE only through the + legacy ``rotary_dim`` config field. Transformers 5.x in-tree configs keep + ``rotary_dim`` as a plain attribute while their models read only + ``rope_parameters["partial_rotary_factor"]``, so the vanilla reference + silently rotates the full head dimension with the wrong frequency ladder + and becomes a deterministic but invalid reference (AMINT-286). Derive the + missing factor as ``rotary_dim / head_dim``. + + Args: + config: Loaded HF config for the vanilla reference model. + + Returns: + True when the config's rope parameters were repaired; False when the + config has no legacy spec or already carries a partial factor. + """ + rotary_dim = getattr(config, "rotary_dim", None) + head_dim = getattr(config, "head_dim", None) + if not rotary_dim or not head_dim or rotary_dim == head_dim: + return False + rope_parameters = getattr(config, "rope_parameters", None) + if isinstance(rope_parameters, dict): + if rope_parameters.get("partial_rotary_factor"): + return False + rope_parameters["partial_rotary_factor"] = rotary_dim / head_dim + return True + if rope_parameters is not None and hasattr(rope_parameters, "partial_rotary_factor"): + if rope_parameters.partial_rotary_factor: + return False + rope_parameters.partial_rotary_factor = rotary_dim / head_dim + return True + return False + + def _dequantize_hf_fp8_weights_in_place(model, output_dtype: torch.dtype) -> int: """Dequantize native per-tensor HF FP8 modules without their runtime kernel. @@ -369,18 +515,184 @@ def _rss_gb() -> float: return rss_pages * page_size / 1024**3 -def _kl_divergence_from_logits(reference_logits: torch.Tensor, candidate_logits: torch.Tensor) -> torch.Tensor: - """Per-token KL(reference || candidate) for full [B, T, V] logits.""" - assert reference_logits.shape == candidate_logits.shape - vocab_size = reference_logits.shape[-1] - ref_log_probs = F.log_softmax(reference_logits.float(), dim=-1).reshape(-1, vocab_size) - cand_log_probs = F.log_softmax(candidate_logits.float(), dim=-1).reshape(-1, vocab_size) - return F.kl_div(cand_log_probs, ref_log_probs, reduction="none", log_target=True).sum(-1) +def _fit_input_ids_to_sequence_length(input_ids: list[int], sequence_length: int) -> list[int]: + """Truncate a tokenized document to the requested parity length.""" + if not input_ids: + raise ValueError("Tokenized parity document must not be empty") + if sequence_length <= 0: + raise ValueError(f"parity_sequence_length must be positive, got {sequence_length}") + if len(input_ids) < sequence_length: + raise ValueError( + f"Tokenized parity document contains {len(input_ids)} tokens, but parity_sequence_length requires " + f"{sequence_length}; choose a shorter sequence length or a longer parity document" + ) + return input_ids[:sequence_length] -def _cosine_similarity_from_logits(reference_logits: torch.Tensor, candidate_logits: torch.Tensor) -> float: - """Cosine similarity over flattened float32 logits.""" - return F.cosine_similarity(reference_logits.flatten().float(), candidate_logits.flatten().float(), dim=0).item() +def _compare_logits( + artifact_dir: Path, + reference_logits: torch.Tensor, + candidate_logits: torch.Tensor, + policy: _LogitParityPolicy, +) -> str | None: + """Compute, persist, and optionally enforce one full-logit comparison. + + Args: + artifact_dir: Directory that owns checkpoint-robustness artifacts. + reference_logits: Reference tensor of shape [..., vocab], with arbitrary leading token dimensions. + candidate_logits: Candidate tensor of shape [..., vocab], matching ``reference_logits`` exactly. + policy: Comparison identity, numerical profile, targeted overrides, and enforcement state. + + Returns: + A failure message when an enforced gate fails, otherwise ``None``. + """ + metrics = _compute_parity_metrics(reference_logits, candidate_logits) + profile_thresholds = _resolve_parity_thresholds(policy.profile, policy.comparison_kind) + threshold_overrides = { + "mean_kl": policy.mean_kl_threshold_override, + "p95_kl": policy.p95_kl_threshold_override, + "cosine_similarity": policy.cosine_threshold_override, + } + uses_threshold_overrides = any(value is not None for value in threshold_overrides.values()) + active_profile_thresholds = _apply_parity_threshold_overrides( + profile_thresholds, + mean_kl=policy.mean_kl_threshold_override, + p95_kl=policy.p95_kl_threshold_override, + cosine_similarity=policy.cosine_threshold_override, + ) + profile_failures = _parity_failures(metrics, profile_thresholds) + active_failures = _parity_failures(metrics, active_profile_thresholds) + threshold_mode = "profile_with_numeric_overrides" if uses_threshold_overrides else "profile" + payload = { + "schema_version": 1, + "parity_document_sha256": _PARITY_DOCUMENT_SHA256, + "phase": policy.phase, + "comparison": policy.comparison, + "comparison_kind": policy.comparison_kind, + "profile": policy.profile, + "profile_thresholds": profile_thresholds.to_dict(), + "threshold_overrides": threshold_overrides, + "active_thresholds": active_profile_thresholds.to_dict(), + "threshold_mode": threshold_mode, + "enforced": policy.enforce, + "passed": not policy.enforce or not active_failures, + "within_active_thresholds": not active_failures, + "would_pass_profile": not profile_failures, + "failures": list(active_failures) if policy.enforce else [], + "threshold_failures": list(active_failures), + "profile_failures": list(profile_failures), + "reference_logits": { + "dtype": str(reference_logits.dtype), + "shape": list(reference_logits.shape), + }, + "candidate_logits": { + "dtype": str(candidate_logits.dtype), + "shape": list(candidate_logits.shape), + }, + "metrics": metrics.to_dict(), + } + report_dir = artifact_dir / "parity_metrics" + report_dir.mkdir(parents=True, exist_ok=True) + report_path = report_dir / f"{policy.phase}_{policy.comparison}.json" + temporary_report_path = report_path.with_suffix(".tmp") + temporary_report_path.write_text(json.dumps(payload, indent=2, sort_keys=True) + "\n") + temporary_report_path.replace(report_path) + print(f"CHECKPOINT_PARITY_METRICS {json.dumps(payload, sort_keys=True)}") + + if not policy.enforce: + print( + f"[{policy.phase}] {policy.comparison} metrics are informational; " + f"would_pass_active_thresholds={not active_failures}, would_pass_profile={not profile_failures}" + ) + return None + if not active_failures: + return None + return f"{policy.comparison} parity failed: " + "; ".join(active_failures) + + +def _comparison_threshold_overrides(custom_args: dict[str, object], comparison: str) -> dict[str, float]: + """Return normalized overrides for one comparison.""" + all_overrides = _normalize_parity_threshold_overrides(custom_args.get("parity_threshold_overrides")) + return all_overrides.get(comparison, {}) + + +def _comparison_profile(custom_args: dict[str, object], comparison: str) -> str: + """Return the comparison profile, falling back to the global profile.""" + return _select_parity_profile( + str(custom_args.get("parity_tolerance_profile", "standard")), + custom_args.get("parity_tolerance_profile_overrides"), + comparison, + ) + + +def _source_load_parity_policy(custom_args: dict[str, object], *, enforce: bool = True) -> _LogitParityPolicy: + """Build the Phase 0 source-load policy.""" + overrides = _comparison_threshold_overrides(custom_args, "source_load") + return _LogitParityPolicy( + phase="phase_0", + comparison="source_load", + comparison_kind="cross_framework", + profile=_comparison_profile(custom_args, "source_load"), + enforce=enforce and not bool(custom_args.get("skip_source_load_logit_parity", False)), + mean_kl_threshold_override=overrides.get("mean_kl"), + p95_kl_threshold_override=overrides.get("p95_kl"), + cosine_threshold_override=overrides.get("cosine_similarity"), + ) + + +def _repeatability_policy(*, phase: str, comparison: str, profile: str) -> _LogitParityPolicy: + """Build an informational policy for two forwards through one loaded model.""" + return _LogitParityPolicy( + phase=phase, + comparison=comparison, + comparison_kind="same_implementation", + profile=profile, + enforce=False, + ) + + +def _automodel_reload_parity_policy(custom_args: dict[str, object]) -> _LogitParityPolicy: + """Build the Phase 2 AutoModel model-reload policy.""" + overrides = _comparison_threshold_overrides(custom_args, "automodel_reload") + return _LogitParityPolicy( + phase="phase_2", + comparison="automodel_model_reload", + comparison_kind="same_implementation", + profile=_comparison_profile(custom_args, "automodel_reload"), + enforce=not bool(custom_args.get("skip_automodel_reload_logit_parity", False)), + mean_kl_threshold_override=overrides.get("mean_kl"), + p95_kl_threshold_override=overrides.get("p95_kl"), + cosine_threshold_override=overrides.get("cosine_similarity"), + ) + + +def _hf_reload_parity_policy(custom_args: dict[str, object]) -> _LogitParityPolicy: + """Build the Phase 3 vanilla-HF export-reload policy.""" + overrides = _comparison_threshold_overrides(custom_args, "hf_reload") + return _LogitParityPolicy( + phase="phase_3", + comparison="hf_export_reload", + comparison_kind="cross_framework", + profile=_comparison_profile(custom_args, "hf_reload"), + enforce=not bool(custom_args.get("skip_hf_reload_logit_parity", False)), + mean_kl_threshold_override=overrides.get("mean_kl"), + p95_kl_threshold_override=overrides.get("p95_kl"), + cosine_threshold_override=overrides.get("cosine_similarity"), + ) + + +def _cross_tp_parity_policy(custom_args: dict[str, object]) -> _LogitParityPolicy: + """Build the optional Phase 5 cross-topology policy.""" + overrides = _comparison_threshold_overrides(custom_args, "cross_tp") + return _LogitParityPolicy( + phase="phase_5", + comparison="cross_tp_reload", + comparison_kind="cross_topology", + profile=_comparison_profile(custom_args, "cross_tp"), + mean_kl_threshold_override=overrides.get("mean_kl"), + p95_kl_threshold_override=overrides.get("p95_kl"), + cosine_threshold_override=overrides.get("cosine_similarity"), + ) def _tensor_digest(tensor: torch.Tensor) -> dict[str, object]: @@ -518,6 +830,48 @@ def _model_kwargs_from_config(model_cfg: ConfigNode) -> dict: } +def _model_pretrained_path(model_cfg: ConfigNode, model_kwargs: dict | None = None) -> str | Path: + """Resolve the source checkpoint for from-pretrained and config-based recipes.""" + direct_path = getattr(model_cfg, "pretrained_model_name_or_path", None) + if direct_path: + return direct_path + + nested_config = getattr(model_cfg, "config", None) + nested_path = getattr(nested_config, "pretrained_model_name_or_path", None) + if nested_path: + return nested_path + nested_name_or_path = getattr(nested_config, "name_or_path", None) + if nested_name_or_path: + return nested_name_or_path + + if model_kwargs is not None: + materialized_config = model_kwargs.get("config") + for attribute in ("pretrained_model_name_or_path", "name_or_path", "_name_or_path"): + materialized_path = getattr(materialized_config, attribute, None) + if materialized_path: + return materialized_path + + raise ValueError( + "Checkpoint robustness requires model.pretrained_model_name_or_path or " + "model.config.pretrained_model_name_or_path" + ) + + +def _set_model_pretrained_path(model_cfg: ConfigNode, pretrained_model_name_or_path: str | Path) -> None: + """Retarget both from-pretrained and config-based recipes to an exported checkpoint.""" + path = str(pretrained_model_name_or_path) + nested_config = getattr(model_cfg, "config", None) + if nested_config is not None and ( + hasattr(nested_config, "pretrained_model_name_or_path") or hasattr(nested_config, "name_or_path") + ): + if hasattr(nested_config, "pretrained_model_name_or_path"): + nested_config.pretrained_model_name_or_path = path + if hasattr(nested_config, "name_or_path"): + nested_config.name_or_path = path + return + model_cfg.pretrained_model_name_or_path = path + + def _resolve_hf_model_class( pretrained_model_name_or_path: str | Path, default_model_cls: type, @@ -525,8 +879,12 @@ def _resolve_hf_model_class( revision: str | None = None, token: str | bool | None = None, ) -> type: - """Honor a checkpoint's advertised HF auto-model class when the VLM default is absent.""" + """Select the vanilla-HF auto-model class supported by the checkpoint.""" from transformers import AutoModelForCausalLM, AutoModelForImageTextToText, PretrainedConfig + from transformers.models.auto.modeling_auto import ( + MODEL_FOR_CAUSAL_LM_MAPPING_NAMES, + MODEL_FOR_IMAGE_TEXT_TO_TEXT_MAPPING_NAMES, + ) config_kwargs: dict[str, str | bool] = { "local_files_only": os.environ.get("HF_HUB_OFFLINE", "0") == "1", @@ -537,15 +895,26 @@ def _resolve_hf_model_class( config_kwargs["token"] = token config_dict, _ = PretrainedConfig.get_config_dict(pretrained_model_name_or_path, **config_kwargs) auto_map = config_dict.get("auto_map") or {} - if not auto_map or default_model_cls.__name__ in auto_map: - return default_model_cls - supported_classes = { model_cls.__name__: model_cls for model_cls in (AutoModelForImageTextToText, AutoModelForCausalLM) } - advertised_classes = [model_cls for name, model_cls in supported_classes.items() if name in auto_map] - if len(advertised_classes) == 1: - return advertised_classes[0] + + if auto_map: + if default_model_cls.__name__ in auto_map: + return default_model_cls + advertised_classes = [model_cls for name, model_cls in supported_classes.items() if name in auto_map] + if len(advertised_classes) == 1: + return advertised_classes[0] + return default_model_cls + + model_type = config_dict.get("model_type") + native_mappings = { + AutoModelForCausalLM: MODEL_FOR_CAUSAL_LM_MAPPING_NAMES, + AutoModelForImageTextToText: MODEL_FOR_IMAGE_TEXT_TO_TEXT_MAPPING_NAMES, + } + native_classes = [model_cls for model_cls, mapping in native_mappings.items() if model_type in mapping] + if len(native_classes) == 1: + return native_classes[0] return default_model_cls @@ -566,20 +935,69 @@ def _get_trust_remote_code_attn_implementation( token: str | bool | None = None, ) -> str: """Select the vanilla-HF attention implementation for a remote-code model.""" + from transformers import PretrainedConfig + + config_kwargs: dict[str, str | bool] = { + "local_files_only": os.environ.get("HF_HUB_OFFLINE", "0") == "1", + } + if revision is not None: + config_kwargs["revision"] = revision + if token is not None: + config_kwargs["token"] = token + config_dict, _ = PretrainedConfig.get_config_dict(pretrained_model_name_or_path, **config_kwargs) + + # Remote-code checkpoints do not share optimized attention backend support: + # these models reject the recipe backend under the pinned Transformers + # version. Eager is their common vanilla-HF reference path. + eager_model_types = {"deepseek_v4", "nemotron-nas", "nemotron_flash", "nemotron_h", "step3p7"} + return "eager" if config_dict.get("model_type") in eager_model_types else "flash_attention_2" + + +def _resolve_hf_attn_implementation( + pretrained_model_name_or_path: str | Path, + requested_implementation: str | None, + *, + hf_model_cls: type, + trust_remote_code: bool, + revision: str | None = None, + token: str | bool | None = None, +) -> str | None: + """Use the recipe backend when vanilla HF supports it, otherwise use eager.""" + if trust_remote_code: + compatible_implementation = _get_trust_remote_code_attn_implementation( + pretrained_model_name_or_path, + revision=revision, + token=token, + ) + if compatible_implementation == "eager" or requested_implementation is None: + return compatible_implementation + return requested_implementation + + if requested_implementation not in {"sdpa", "flash_attention_2"}: + return requested_implementation + from transformers import AutoConfig - config_kwargs: dict[str, str | bool] = {"trust_remote_code": True} + config_kwargs: dict[str, str | bool] = { + "local_files_only": os.environ.get("HF_HUB_OFFLINE", "0") == "1", + } if revision is not None: config_kwargs["revision"] = revision if token is not None: config_kwargs["token"] = token config = AutoConfig.from_pretrained(pretrained_model_name_or_path, **config_kwargs) + try: + concrete_model_cls = hf_model_cls._model_mapping[type(config)] + except (AttributeError, KeyError): + return requested_implementation - # Remote-code checkpoints do not share optimized attention backend support: - # Nemotron-H has incompatible FA2/SDPA paths, and Step-3.7 explicitly rejects - # FA2. Eager is their common HF reference path. Other remote-code models - # (notably Nemotron-Flash) still require FA2. - return "eager" if config.model_type in {"nemotron_h", "step3p7"} else "flash_attention_2" + support_attribute = { + "sdpa": "_supports_sdpa", + "flash_attention_2": "_supports_flash_attn", + }[requested_implementation] + if not bool(getattr(concrete_model_cls, support_attribute, False)): + return "eager" + return requested_implementation def _hf_source_load_kwargs( @@ -589,6 +1007,7 @@ def _hf_source_load_kwargs( source_dtype: torch.dtype, trust_remote_code: bool, experts_implementation: str | None, + hf_model_cls: type, device: torch.device, hf_device_map_auto: bool, ) -> dict: @@ -603,14 +1022,18 @@ def _hf_source_load_kwargs( } hf_kwargs = {k: v for k, v in model_kwargs.items() if k in hf_allowed_keys} hf_kwargs["torch_dtype"] = source_dtype - hf_kwargs["trust_remote_code"] = trust_remote_code or bool(hf_kwargs.get("trust_remote_code", False)) + hf_kwargs["trust_remote_code"] = trust_remote_code hf_kwargs["local_files_only"] = os.environ.get("HF_HUB_OFFLINE", "0") == "1" - if hf_kwargs["trust_remote_code"] and "attn_implementation" not in hf_kwargs: - hf_kwargs["attn_implementation"] = _get_trust_remote_code_attn_implementation( - pretrained_model_name_or_path, - revision=hf_kwargs.get("revision"), - token=hf_kwargs.get("token"), - ) + attn_implementation = _resolve_hf_attn_implementation( + pretrained_model_name_or_path, + hf_kwargs.get("attn_implementation"), + hf_model_cls=hf_model_cls, + trust_remote_code=hf_kwargs["trust_remote_code"], + revision=hf_kwargs.get("revision"), + token=hf_kwargs.get("token"), + ) + if attn_implementation is not None: + hf_kwargs["attn_implementation"] = attn_implementation if experts_implementation and not trust_remote_code: hf_kwargs["experts_implementation"] = experts_implementation hf_kwargs["trust_remote_code"] = False @@ -723,6 +1146,17 @@ def _hf_fp32_module_names(hf_config: object) -> tuple[str, ...]: for name in getattr(model_cls, "_keep_in_fp32_modules_strict", None) or (): if name not in module_names: module_names.append(name) + # AutoModel strict names use AutoModel module paths, but vanilla HF + # layouts can hang the same tensor off a different parent (in-tree + # MiniMax-M2 keeps e_score_correction_bias on ``mlp``, not + # ``mlp.gate``), so the AutoModel-path entry silently fails to match + # and the reference's router bias was cast to bf16 — scrambling 30-73% + # of knife-edge top-k selections per layer (AMINT-286). Also register + # the distinctive leaf so any layout keeps the tensor in fp32; generic + # ``weight``/``bias`` leaves are excluded to avoid pinning everything. + leaf = name.rsplit(".", 1)[-1] + if leaf not in ("weight", "bias") and leaf not in module_names: + module_names.append(leaf) return tuple(module_names) @@ -821,15 +1255,17 @@ def _load_input_ids_once( cfg, input_ids_loader: Callable[[str | None], list[int]], tokenizer_name: str | None, + *, + sequence_length: int, ) -> list[int]: - """Load dynamic input IDs once before distributed initialization. + """Load and expand dynamic input IDs once before distributed initialization. The tokenizer and processor imports are I/O-heavy on shared filesystems. Loading on every worker can turn a cold import into a multi-node import storm, so rank 0 writes the small result for the other ranks to read. """ if tokenizer_name is None or _preinit_world_size() == 1: - return input_ids_loader(tokenizer_name) + return _fit_input_ids_to_sequence_length(input_ids_loader(tokenizer_name), sequence_length) sync_dir, payload_path, done_path, fail_path = _input_ids_sync_paths(cfg) if _preinit_global_rank() != 0: @@ -843,7 +1279,7 @@ def _load_input_ids_once( done_path.unlink(missing_ok=True) fail_path.unlink(missing_ok=True) try: - input_ids = input_ids_loader(tokenizer_name) + input_ids = _fit_input_ids_to_sequence_length(input_ids_loader(tokenizer_name), sequence_length) temporary_payload_path = payload_path.with_suffix(".tmp") temporary_payload_path.write_text(json.dumps(input_ids)) temporary_payload_path.replace(payload_path) @@ -901,9 +1337,10 @@ def _hf_reload_sync_paths(cfg) -> tuple[Path, Path]: return sync_dir, sync_dir / "done" -def _wait_for_hf_reload_rank0(done_path: Path) -> None: +def _wait_for_hf_reload_rank0(done_path: Path, *, timeout_s: int | None = None) -> None: """Wait without an active collective for rank 0 to finish the vanilla-HF reload.""" - timeout_s = int(os.environ.get("HF_RELOAD_TIMEOUT_SECONDS", "1800")) + if timeout_s is None: + timeout_s = int(os.environ.get("HF_RELOAD_TIMEOUT_SECONDS", "1800")) deadline = time.monotonic() + timeout_s while time.monotonic() < deadline: if done_path.exists(): @@ -912,7 +1349,7 @@ def _wait_for_hf_reload_rank0(done_path: Path) -> None: raise TimeoutError(f"Timed out waiting {timeout_s}s for rank 0 vanilla-HF reload") -def _prepare_hf_reload_sync(cfg) -> tuple[Path, Path] | None: +def _prepare_hf_reload_sync(cfg, *, timeout_s: int | None = None) -> tuple[Path, Path] | None: """Prepare ranks for a long rank-0-only HF reload without starting an NCCL wait.""" if not dist.is_initialized() or dist.get_world_size() == 1: return None @@ -923,7 +1360,7 @@ def _prepare_hf_reload_sync(cfg) -> tuple[Path, Path] | None: done_path.unlink(missing_ok=True) _barrier() # ensure all ranks released recipe memory and rank 0 reset the marker if not _rank0(): - _wait_for_hf_reload_rank0(done_path) + _wait_for_hf_reload_rank0(done_path, timeout_s=timeout_s) return sync_dir, done_path @@ -966,15 +1403,25 @@ def _record_deferred_failure( print(f"[{phase}] Comparison failed; deferring failure until later checkpoint phases complete.") +def _broadcast_rank0_failure(failure_message: str | None) -> str | None: + """Broadcast one rank-0 comparison result so every worker follows the same path.""" + if not dist.is_initialized(): + return failure_message + payload = [failure_message] + dist.broadcast_object_list(payload, src=0) + return payload[0] + + def _prepare_source_load_reference( cfg, input_ids: list[int], *, hf_model_cls: type, - trust_remote_code: bool, + trust_remote_code: bool | None, experts_implementation: str | None, hf_device_map_auto: bool, hf_source_post_load_dequantize: bool, + parity_tolerance_profile: str = "standard", ) -> tuple[torch.Tensor, bool | None, bool | None] | None: """Compute vanilla HF source-load reference logits before trainer construction.""" if _preinit_world_size() > 1: @@ -1001,6 +1448,7 @@ def _prepare_source_load_reference( experts_implementation=experts_implementation, hf_device_map_auto=hf_device_map_auto, hf_source_post_load_dequantize=hf_source_post_load_dequantize, + parity_tolerance_profile=parity_tolerance_profile, ) except Exception: if fail_path is not None: @@ -1017,10 +1465,11 @@ def _prepare_source_load_reference_rank0( input_ids: list[int], *, hf_model_cls: type, - trust_remote_code: bool, + trust_remote_code: bool | None, experts_implementation: str | None, hf_device_map_auto: bool, hf_source_post_load_dequantize: bool, + parity_tolerance_profile: str = "standard", ) -> tuple[torch.Tensor, bool | None, bool | None]: """Rank-0 implementation of vanilla HF source-load reference capture.""" from nemo_automodel._transformers.utils import apply_cache_compatibility_patches @@ -1029,8 +1478,7 @@ def _prepare_source_load_reference_rank0( _patch_remote_masking_api_compatibility() model_kwargs = _model_kwargs_from_config(cfg.model) - original_pretrained_path = model_kwargs.get("pretrained_model_name_or_path") - assert original_pretrained_path is not None, "source-load parity requires model.pretrained_model_name_or_path" + original_pretrained_path = _model_pretrained_path(cfg.model, model_kwargs) hf_model_cls = _resolve_hf_model_class( original_pretrained_path, hf_model_cls, @@ -1038,7 +1486,8 @@ def _prepare_source_load_reference_rank0( token=model_kwargs.get("token"), ) source_dtype = _resolve_source_load_dtype(model_kwargs) - trust_remote_code = trust_remote_code or bool(model_kwargs.get("trust_remote_code", False)) + if trust_remote_code is None: + trust_remote_code = bool(model_kwargs.get("trust_remote_code", False)) device = torch.device("cuda", torch.cuda.current_device()) if torch.cuda.is_available() else torch.device("cpu") hf_kwargs = _hf_source_load_kwargs( @@ -1047,9 +1496,16 @@ def _prepare_source_load_reference_rank0( source_dtype=source_dtype, trust_remote_code=trust_remote_code, experts_implementation=experts_implementation, + hf_model_cls=hf_model_cls, device=device, hf_device_map_auto=hf_device_map_auto, ) + requested_attn_implementation = model_kwargs.get("attn_implementation") + if hf_kwargs.get("attn_implementation") != requested_attn_implementation: + print( + "[Phase 0] Vanilla-HF attention compatibility fallback: " + f"requested={requested_attn_implementation!r}, selected={hf_kwargs.get('attn_implementation')!r}" + ) if hf_source_post_load_dequantize and hf_kwargs.get("device_map") == "auto" and torch.cuda.is_available(): # Accelerate sizes the automatic map for the on-disk FP8 tensors. The # post-load BF16 representation needs roughly twice that memory, so cap @@ -1080,6 +1536,10 @@ def _prepare_source_load_reference_rank0( revision=hf_kwargs.get("revision"), token=hf_kwargs.get("token"), ) + if _repair_legacy_partial_rotary_config(hf_config): + # The repaired spec only reaches the model when the config object is + # passed explicitly; from_pretrained otherwise re-reads config.json. + hf_kwargs["config"] = hf_config model_load_context = _hf_model_load_context( trust_remote_code=trust_remote_code, @@ -1105,6 +1565,18 @@ def _prepare_source_load_reference_rank0( fix_rotary_embeddings([hf_model]) hf_logits = _get_logits(hf_model, input_ids, device) + repeated_hf_logits = _get_logits(hf_model, input_ids, device) + _compare_logits( + _robustness_artifact_dir(cfg), + hf_logits, + repeated_hf_logits, + _repeatability_policy( + phase="phase_0", + comparison="hf_source_self_repeat", + profile=parity_tolerance_profile, + ), + ) + del repeated_hf_logits hf_aliased = _lm_head_embedding_aliased(hf_model) explicit_tie_word_embeddings = _explicit_tie_word_embeddings(hf_model.config) del hf_model @@ -1117,9 +1589,8 @@ def _compare_source_load_parity( candidate_logits: torch.Tensor, candidate_aliased: bool | None, *, - source_load_kl_threshold: float, - source_load_mean_kl_threshold: float, - source_load_cosine_threshold: float, + artifact_dir: Path, + policy: _LogitParityPolicy, ) -> str | None: """Compare the vanilla HF source-load reference against the constructed trainer model. @@ -1128,9 +1599,8 @@ def _compare_source_load_parity( embedding alias state, and the explicit tie-word-embeddings setting. Other ranks pass ``None``. candidate_logits: Constructed trainer logits of shape [batch, sequence, vocab]. candidate_aliased: Constructed trainer input/output embedding alias state. - source_load_kl_threshold: Maximum allowed per-token KL divergence. - source_load_mean_kl_threshold: Maximum allowed mean per-token KL divergence. - source_load_cosine_threshold: Minimum allowed cosine similarity over flattened logits. + artifact_dir: Directory that owns checkpoint-robustness artifacts. + policy: Source-load metric profile, legacy overrides, and enforcement state. Returns: Synchronized failure traceback when source-load parity fails, otherwise ``None``. The caller may defer this @@ -1145,32 +1615,12 @@ def _compare_source_load_parity( f"Source-load parity shape mismatch: HF logits {hf_logits.shape} vs trainer logits " f"{candidate_logits.shape}" ) - kl_source = _kl_divergence_from_logits(hf_logits, candidate_logits) - max_kl_source = kl_source.max().item() - mean_kl_source = kl_source.mean().item() - p95_kl_source = torch.quantile(kl_source, 0.95).item() - cosine_source = _cosine_similarity_from_logits(hf_logits, candidate_logits) + parity_failure = _compare_logits(artifact_dir, hf_logits, candidate_logits, policy) + if parity_failure is not None: + raise AssertionError(parity_failure) print( - f"[Phase 0] Source-load vs constructed-trainer max KL: {max_kl_source:.6e} " - f"(threshold: {source_load_kl_threshold:.6e}); mean KL: {mean_kl_source:.6e} " - f"(threshold: {source_load_mean_kl_threshold:.6e}); p95 KL: {p95_kl_source:.6e}; " - f"cosine={cosine_source:.8f} " - f"(threshold: {source_load_cosine_threshold:.8f}); " - f"hf_aliased={hf_aliased}; trainer_aliased={candidate_aliased}; " - f"tie_word_embeddings={explicit_tie_word_embeddings}" - ) - - assert max_kl_source <= source_load_kl_threshold, ( - f"KL divergence between original HF source load and constructed trainer model too large: " - f"max per-token KL = {max_kl_source:.6e} > threshold {source_load_kl_threshold:.6e}" - ) - assert mean_kl_source <= source_load_mean_kl_threshold, ( - f"Mean KL divergence between original HF source load and constructed trainer model too large: " - f"mean per-token KL = {mean_kl_source:.6e} > threshold {source_load_mean_kl_threshold:.6e}" - ) - assert cosine_source >= source_load_cosine_threshold, ( - f"Cosine similarity between original HF source load and constructed trainer model too low: " - f"cosine = {cosine_source:.8f} < threshold {source_load_cosine_threshold:.8f}" + f"[Phase 0] Source-load aliases: hf_aliased={hf_aliased}; " + f"trainer_aliased={candidate_aliased}; tie_word_embeddings={explicit_tie_word_embeddings}" ) if hf_aliased is not None and candidate_aliased is not None: assert hf_aliased == candidate_aliased, ( @@ -1208,37 +1658,14 @@ def _get_logits_pp(trainer, input_ids, device) -> torch.Tensor: pp_batch_size = trainer.pipeline_config.pp_batch_size orig_seq_len = len(input_ids) - # PP recv buffer shapes are locked at first forward. r0.4.0 lacks - # AutoPipeline.update_seq_len (added in #1689) to resize on the fly, so - # discover the locked seq_len from the stages and pad input_ids to match - # for the forward pass. Captured logits are sliced back to orig_seq_len. - def _discover_pp_seq_len() -> int: - pp_seq_len = getattr(trainer.pp, "pp_seq_len", None) - if pp_seq_len: - return pp_seq_len - for stage in getattr(trainer.pp.info, "stages", None) or (): - inputs_meta = getattr(stage, "inputs_meta", None) - if not inputs_meta: - inputs_meta = getattr(getattr(stage, "_user_meta", None), "inputs", None) - for meta in inputs_meta or (): - shape = getattr(meta, "shape", ()) - if len(shape) >= 2 and shape[1] > 0: - return shape[1] - ds_seq_length = trainer.cfg.get("dataset.seq_length", None) - return ds_seq_length or orig_seq_len - - pp_seq_len = _discover_pp_seq_len() - if orig_seq_len < pp_seq_len: - input_ids = list(input_ids) + [0] * (pp_seq_len - orig_seq_len) + # PyTorch pipeline stages preallocate activation buffers for one sequence + # shape. Resize those buffers before this parity-only forward just as the + # training recipes do before every schedule step. + trainer.pp.update_seq_len(orig_seq_len) # Replicate the prompt to pp_batch_size so the schedule's batch split is valid. ids = torch.tensor([input_ids] * pp_batch_size, device=device, dtype=torch.long) - # The PP schedule requires the static stage sequence length, but the parity - # prompt is usually much shorter. Keep synthetic tail tokens out of both - # attention and MoE dispatch so this forward represents the same prompt as - # the unpadded HF reference. - attention_mask = torch.zeros_like(ids) - attention_mask[:, :orig_seq_len] = 1 + attention_mask = torch.ones_like(ids) targets = torch.zeros_like(ids) if trainer.pp.info.has_last_stage else None captured = [None] @@ -1396,34 +1823,62 @@ def _prepopulate_hf_dynamic_modules_cache(local_dir: Path | str) -> None: shutil.copy2(src_py, dst_py) -def _tp_size_from_argv(argv) -> int: - """Peek at --distributed.tp_size / --config YAML without constructing the cfg. +def _prepare_consolidated_hf_cache_once(cfg, consolidated_dir: Path) -> None: + """Prepare remote-code files once before an isolated distributed setup. - Returns 1 if no TP setting is found. Used before cfg parsing to pick a - reasonable default kl_threshold. + Every worker reaches this function before the recipe initializes a process + group, so ``dist.get_rank()`` cannot select the writer. A small shared-file + marker lets pre-init global rank 0 finish all cache writes before the other + workers import the consolidated checkpoint's dynamic modules. """ - for i, a in enumerate(argv): - if a == "--distributed.tp_size" and i + 1 < len(argv): - try: - return int(argv[i + 1]) - except (TypeError, ValueError): - return 1 - config_path = None - for i, a in enumerate(argv): - if a == "--config" and i + 1 < len(argv): - config_path = argv[i + 1] - break - if config_path: + expected_payload = str(consolidated_dir.resolve()) + sync_dir = _robustness_artifact_dir(cfg) / "hf_dynamic_modules_cache" + done_path = sync_dir / "done" + fail_path = sync_dir / "fail" + + def is_ready() -> bool: try: - import yaml + return done_path.read_text() == expected_payload + except FileNotFoundError: + return False - with open(config_path) as f: - raw_cfg = yaml.safe_load(f) or {} - tp = (raw_cfg.get("distributed") or {}).get("tp_size", 1) - return int(tp) if tp is not None else 1 + def prepare_cache() -> None: + from transformers import AutoConfig + + _prepopulate_hf_dynamic_modules_cache(consolidated_dir) + try: + AutoConfig.from_pretrained(str(consolidated_dir), trust_remote_code=True) except Exception: pass - return 1 + + if is_ready(): + return + if _preinit_world_size() == 1: + prepare_cache() + return + if _preinit_global_rank() == 0: + sync_dir.mkdir(parents=True, exist_ok=True) + done_path.unlink(missing_ok=True) + fail_path.unlink(missing_ok=True) + try: + prepare_cache() + temporary_done_path = done_path.with_suffix(".tmp") + temporary_done_path.write_text(expected_payload) + temporary_done_path.replace(done_path) + except Exception: + fail_path.write_text(traceback.format_exc()) + raise + return + + timeout_s = int(os.environ.get("HF_DYNAMIC_MODULE_CACHE_TIMEOUT_SECONDS", "1800")) + deadline = time.monotonic() + timeout_s + while time.monotonic() < deadline: + if fail_path.exists(): + raise RuntimeError(f"Rank 0 dynamic-module cache preparation failed:\n{fail_path.read_text()}") + if is_ready(): + return + time.sleep(5) + raise TimeoutError(f"Timed out waiting {timeout_s}s for rank 0 dynamic-module cache preparation") def _rank0() -> bool: @@ -1497,25 +1952,13 @@ def _materialize_hf_quantization_config(cfg): return raw_quantization_config -def _hf_reload_kl_error(max_kl_hf: float, hf_kl_threshold: float) -> str | None: - """Return an actionable HF reload parity error, including non-finite results.""" - if not math.isfinite(max_kl_hf): - return f"HF-loaded model produced non-finite KL divergence: {max_kl_hf}" - if max_kl_hf > hf_kl_threshold: - return ( - "KL divergence between original and HF-loaded model too large: " - f"max per-token KL = {max_kl_hf:.6e} > threshold {hf_kl_threshold:.6e}" - ) - return None - - def _run_vanilla_hf_reload( cfg, input_ids: list[int], reference_logits: torch.Tensor, *, hf_model_cls: type, - custom_args: dict, + custom_args: dict[str, object], ) -> str | None: """Load the saved model with vanilla HF and validate its adapter and forward pass. @@ -1530,19 +1973,27 @@ def _run_vanilla_hf_reload( An error message when loading or parity fails, otherwise ``None``. """ try: + from nemo_automodel._transformers.utils import apply_cache_compatibility_patches + + # Match Phase 0's vanilla-HF setup. Exported trust-remote-code models can + # still carry Transformers-v4 list-form ``_tied_weights_keys``. + apply_cache_compatibility_patches() _patch_remote_masking_api_compatibility() _, ckpt_step_dir, consolidated_dir = _checkpoint_paths(cfg) is_peft = hasattr(cfg, "peft") - original_pretrained_path = cfg.model.pretrained_model_name_or_path model_kwargs = _model_kwargs_from_config(cfg.model) + original_pretrained_path = _model_pretrained_path(cfg.model, model_kwargs) original_quantization_config = _materialize_hf_quantization_config(cfg) - trust_remote_code = bool(custom_args.get("trust_remote_code", False)) + configured_trust_remote_code = custom_args.get("trust_remote_code") + trust_remote_code = ( + bool(model_kwargs.get("trust_remote_code", False)) + if configured_trust_remote_code is None + else bool(configured_trust_remote_code) + ) experts_implementation = custom_args.get("experts_implementation", None) hf_device_map_auto = bool(custom_args.get("hf_device_map_auto", False)) check_fused_qkv_keys = bool(custom_args.get("check_fused_qkv_keys", False)) - skip_hf_logit_parity = bool(custom_args.get("skip_hf_logit_parity", False)) hf_adapter_ignored_key_prefix = custom_args.get("hf_adapter_ignored_key_prefix") - hf_kl_threshold = float(custom_args.get("hf_kl_threshold", "5e-3")) device = torch.device("cuda", torch.cuda.current_device()) config_path = original_pretrained_path if is_peft else consolidated_dir hf_model_cls = _resolve_hf_model_class( @@ -1560,17 +2011,24 @@ def _run_vanilla_hf_reload( for key in ("revision", "token"): if model_kwargs.get(key) is not None: hf_kwargs[key] = model_kwargs[key] - # Load HF with the attention backend the recipe pins, the same way the - # source-load phase does. Attention backends are not bit-identical in bf16, - # so without this the two sides can run different backends and the reload - # reports a logit gap that the checkpoint did not cause. - if model_kwargs.get("attn_implementation") is not None: - hf_kwargs["attn_implementation"] = model_kwargs["attn_implementation"] - # Remote-code models can ship attention names that transformers 5.x - # rejects. Select a supported implementation while keeping Nemotron-H - # off HF's incompatible FlashAttention varlen path. - if trust_remote_code and "attn_implementation" not in hf_kwargs: - hf_kwargs["attn_implementation"] = _get_trust_remote_code_attn_implementation(config_path) + # Keep the recipe backend when vanilla HF supports it. Some model + # implementations reject that backend under the pinned Transformers + # version, so their independent HF reference uses eager instead. + attn_implementation = _resolve_hf_attn_implementation( + config_path, + model_kwargs.get("attn_implementation"), + hf_model_cls=hf_model_cls, + trust_remote_code=trust_remote_code, + revision=model_kwargs.get("revision"), + token=model_kwargs.get("token"), + ) + if attn_implementation is not None: + hf_kwargs["attn_implementation"] = attn_implementation + if attn_implementation != model_kwargs.get("attn_implementation") and _rank0(): + print( + "[Phase 3] Vanilla-HF attention compatibility fallback: " + f"requested={model_kwargs.get('attn_implementation')!r}, selected={attn_implementation!r}" + ) if experts_implementation and not trust_remote_code: hf_kwargs["experts_implementation"] = experts_implementation hf_kwargs["trust_remote_code"] = False @@ -1600,6 +2058,12 @@ def _run_vanilla_hf_reload( revision=model_kwargs.get("revision"), token=model_kwargs.get("token"), ) + if _repair_legacy_partial_rotary_config(hf_config): + # The repaired spec only reaches the model when the config object + # is passed explicitly; from_pretrained otherwise re-reads + # config.json (the consolidated export copies the source config, + # so it carries the same legacy rotary_dim field). + hf_kwargs["config"] = hf_config # Load the reference model straight onto the target GPU. Materialising a # 14B checkpoint on CPU and then ``.to(device)`` costs ~50-225s, and that # rank-0-only stall trips the NCCL watchdog while the other ranks idle at @@ -1656,6 +2120,18 @@ def _run_vanilla_hf_reload( f"{hf_adapter_ignored_key_prefix!r} ({ignored_adapter_tensors} tensors)" ) hf_logits = _get_logits(peft_model, input_ids, device) + repeated_hf_logits = _get_logits(peft_model, input_ids, device) + _compare_logits( + _robustness_artifact_dir(cfg), + hf_logits, + repeated_hf_logits, + _repeatability_policy( + phase="phase_3", + comparison="hf_export_self_repeat", + profile=_comparison_profile(custom_args, "hf_reload"), + ), + ) + del repeated_hf_logits if check_fused_qkv_keys: from safetensors import safe_open @@ -1689,15 +2165,26 @@ def _run_vanilla_hf_reload( if should_fix_rotary_embeddings([hf_model]): fix_rotary_embeddings([hf_model]) hf_logits = _get_logits(hf_model, input_ids, device) + repeated_hf_logits = _get_logits(hf_model, input_ids, device) + _compare_logits( + _robustness_artifact_dir(cfg), + hf_logits, + repeated_hf_logits, + _repeatability_policy( + phase="phase_3", + comparison="hf_export_self_repeat", + profile=_comparison_profile(custom_args, "hf_reload"), + ), + ) + del repeated_hf_logits del hf_model - hf_reload_error = None - if skip_hf_logit_parity: - print("[HF reload] Forward smoke passed; cross-implementation logit KL comparison skipped by config") - else: - max_kl_hf = _kl_divergence_from_logits(reference_logits, hf_logits).max().item() - print(f"[HF reload] HF-loaded max KL: {max_kl_hf:.6e} (threshold: {hf_kl_threshold:.6e})") - hf_reload_error = _hf_reload_kl_error(max_kl_hf, hf_kl_threshold) + hf_reload_error = _compare_logits( + _robustness_artifact_dir(cfg), + reference_logits, + hf_logits, + _hf_reload_parity_policy(custom_args), + ) del hf_logits _release_model_memory() return hf_reload_error @@ -1755,7 +2242,7 @@ def _raise_distributed_failure(failure_message: str | None) -> None: def _run_process_isolated_checkpoint_phase( phase: str, *, - custom_args: dict, + custom_args: dict[str, object], recipe_cls: type[BaseRecipe], hf_model_cls: type, input_ids_loader: Callable[[str | None], list[int]], @@ -1780,21 +2267,23 @@ def _run_process_isolated_checkpoint_phase( "automodel_reload", "hf_reload", "resume", + "cross_tp_reload", } if phase not in supported_phases: raise ValueError(f"Unsupported isolated checkpoint phase {phase!r}; expected one of {sorted(supported_phases)}") - if int(custom_args.get("cross_tp_size", "0")) > 0: - raise ValueError("Process-isolated checkpoint mode does not yet support cross_tp_size") - if custom_args.get("no_check_resume", False) and phase == "resume": - raise ValueError(f"Process-isolated phase {phase!r} conflicts with no_check_resume=true") + if custom_args.get("skip_resume", False) and phase == "resume": + raise ValueError(f"Process-isolated phase {phase!r} conflicts with skip_resume=true") + if phase == "cross_tp_reload" and int(custom_args.get("cross_tp_size", "0")) <= 0: + raise ValueError("Process-isolated cross_tp_reload requires cross_tp_size > 0") _disable_distributed_atexit_teardown() cfg = parse_args_and_load_config() tokenizer_name = custom_args.get("tokenizer_name", None) + parity_sequence_length = int(custom_args.get("parity_sequence_length", "2048")) if phase == "source_load_reference": - if not custom_args.get("check_source_load_parity", False): - raise ValueError("Isolated source_load_reference requires check_source_load_parity=true") + if not custom_args.get("source_load_parity_enabled", False): + raise ValueError("Isolated source_load_reference requires Phase 0 to be enabled") reference_path, metadata_path = _source_load_artifact_paths(cfg) source_load_fail_path = _source_load_sync_paths(cfg)[2] if _preinit_world_size() > 1 else None @@ -1808,20 +2297,31 @@ def _run_process_isolated_checkpoint_phase( path.unlink(missing_ok=True) _report_phase("Isolated Phase 0a source load: loading prompt input IDs") - input_ids = _load_input_ids_once(cfg, input_ids_loader, tokenizer_name) + input_ids = _load_input_ids_once( + cfg, + input_ids_loader, + tokenizer_name, + sequence_length=parity_sequence_length, + ) _report_phase("Isolated Phase 0a source load: starting vanilla-HF reference load") source_load_reference = _prepare_source_load_reference( cfg, input_ids, hf_model_cls=hf_model_cls, - trust_remote_code=bool(custom_args.get("trust_remote_code", False)), + trust_remote_code=custom_args.get("trust_remote_code"), experts_implementation=custom_args.get("experts_implementation", None), hf_device_map_auto=bool(custom_args.get("hf_device_map_auto", False)), hf_source_post_load_dequantize=bool(custom_args.get("hf_source_post_load_dequantize", False)), + parity_tolerance_profile=_comparison_profile(custom_args, "source_load"), ) if _preinit_global_rank() == 0: assert source_load_reference is not None, "rank 0 source-load reference was not captured" reference_logits, hf_aliased, explicit_tie_word_embeddings = source_load_reference + token_count, vocab_size = _validate_logits(reference_logits) + print( + f"[Phase 0] Vanilla-HF source forward produced finite logits for " + f"{token_count} tokens and vocab_size={vocab_size}" + ) reference_path.parent.mkdir(parents=True, exist_ok=True) temporary_reference_path = reference_path.with_suffix(".tmp") temporary_metadata_path = metadata_path.with_suffix(".tmp") @@ -1849,11 +2349,16 @@ def _run_process_isolated_checkpoint_phase( return if phase == "source_load_parity": - if not custom_args.get("check_source_load_parity", False): - raise ValueError("Isolated source_load_parity requires check_source_load_parity=true") + if not custom_args.get("source_load_parity_enabled", False): + raise ValueError("Isolated source_load_parity requires Phase 0 to be enabled") _report_phase("Isolated Phase 0b source parity: loading prompt input IDs") - input_ids = _load_input_ids_once(cfg, input_ids_loader, tokenizer_name) + input_ids = _load_input_ids_once( + cfg, + input_ids_loader, + tokenizer_name, + sequence_length=parity_sequence_length, + ) reference_path, metadata_path = _source_load_artifact_paths(cfg) assert reference_path.exists(), f"Source-load reference logits not found at {reference_path}" assert metadata_path.exists(), f"Source-load reference metadata not found at {metadata_path}" @@ -1886,9 +2391,8 @@ def _run_process_isolated_checkpoint_phase( source_load_reference, trainer_source_logits, _lm_head_embedding_aliased(source_trainer.model_parts[0]), - source_load_kl_threshold=float(custom_args.get("source_load_kl_threshold", "5e-3")), - source_load_mean_kl_threshold=float(custom_args.get("source_load_mean_kl_threshold", "1e-3")), - source_load_cosine_threshold=float(custom_args.get("source_load_cosine_threshold", "0.9999")), + artifact_dir=_robustness_artifact_dir(cfg), + policy=_source_load_parity_policy(custom_args), ) _barrier() if _rank0(): @@ -1907,8 +2411,13 @@ def _run_process_isolated_checkpoint_phase( if custom_args.get("skip_hf_reload", False): raise ValueError("Process-isolated hf_reload conflicts with skip_hf_reload=true") - _report_phase("Isolated vanilla-HF reload: loading prompt input IDs") - input_ids = _load_input_ids_once(cfg, input_ids_loader, tokenizer_name) + _report_phase("Isolated Phase 3 vanilla-HF export reload: loading prompt input IDs") + input_ids = _load_input_ids_once( + cfg, + input_ids_loader, + tokenizer_name, + sequence_length=parity_sequence_length, + ) # The HF model is sharded by one rank-0 process over all GPUs on its # node. A CPU process group keeps the remaining workers synchronized @@ -1924,7 +2433,10 @@ def _run_process_isolated_checkpoint_phase( _barrier() reference_path = _robustness_artifact_dir(cfg) / "reference_logits.pt" - hf_reload_sync_paths = _prepare_hf_reload_sync(cfg) + hf_reload_timeout_s = ( + int(custom_args["hf_reload_timeout_seconds"]) if "hf_reload_timeout_seconds" in custom_args else None + ) + hf_reload_sync_paths = _prepare_hf_reload_sync(cfg, timeout_s=hf_reload_timeout_s) hf_reload_error = None if _rank0(): if not reference_path.exists(): @@ -1944,19 +2456,24 @@ def _run_process_isolated_checkpoint_phase( f"CHECKPOINT_ROBUSTNESS_PHASE_FAILURE phase=hf_reload check=hf_reload_parity\n{hf_reload_error}" ) _raise_distributed_failure(hf_reload_error) - _report_phase("Isolated vanilla-HF reload: parity complete; exiting phase") + _report_phase("Isolated Phase 3 vanilla-HF export reload: parity complete; exiting phase") return if phase == "train_and_save": - _report_phase("Isolated train/save: loading prompt input IDs") - input_ids = _load_input_ids_once(cfg, input_ids_loader, tokenizer_name) + _report_phase("Isolated Phase 1 train/save/reference: loading prompt input IDs") + input_ids = _load_input_ids_once( + cfg, + input_ids_loader, + tokenizer_name, + sequence_length=parity_sequence_length, + ) resume_plan = None - if custom_args.get("check_resume", False): + if custom_args.get("resume_enabled", False): resume_plan = _resume_plan_from_config(cfg) _configure_uninterrupted_run(cfg, resume_plan) torch.cuda.reset_peak_memory_stats() - _report_phase("Isolated train/save: starting trainer setup") + _report_phase("Isolated Phase 1 train/save/reference: starting trainer setup") trainer = recipe_cls(cfg) trainer.setup() resume_recorder = None @@ -1968,16 +2485,16 @@ def _run_process_isolated_checkpoint_phase( if reproducibility_dir is not None: reproducibility_recorder = _TrainingReproducibilityRecorder(trainer) reproducibility_recorder.attach() - _report_phase("Isolated train/save: trainer setup complete") + _report_phase("Isolated Phase 1 train/save/reference: trainer setup complete") if tokenizer_name is not None and dist.is_initialized() and dist.get_world_size() > 1: _barrier() if _rank0(): _cleanup_input_ids_sync(cfg) _barrier() - _report_phase("Isolated train/save: starting training and checkpoint") + _report_phase("Isolated Phase 1 train/save/reference: starting training and checkpoint") trainer.run_train_validation_loop() - _report_phase("Isolated train/save: training and checkpoint complete") + _report_phase("Isolated Phase 1 train/save/reference: training and checkpoint complete") if resume_recorder is not None: _persist_reference_trajectory(resume_recorder) @@ -2011,9 +2528,27 @@ def _run_process_isolated_checkpoint_phase( if max_cpu_gb > 0: assert peak_cpu_gb <= max_cpu_gb, f"Peak CPU RSS {peak_cpu_gb:.2f} GB exceeds threshold {max_cpu_gb:.2f} GB" - _report_phase("Isolated train/save: capturing reference logits") + _report_phase("Isolated Phase 1 train/save/reference: capturing reference logits") device = next(trainer.model_parts[0].parameters()).device reference_logits = _get_logits(trainer.model_parts[0], input_ids, device, trainer=trainer) + token_count, vocab_size = _validate_logits(reference_logits) + repeated_reference_logits = _get_logits(trainer.model_parts[0], input_ids, device, trainer=trainer) + if _rank0(): + _compare_logits( + _robustness_artifact_dir(cfg), + reference_logits, + repeated_reference_logits, + _repeatability_policy( + phase="phase_1", + comparison="automodel_reference_self_repeat", + profile=str(custom_args.get("parity_tolerance_profile", "standard")), + ), + ) + print( + f"[Phase 1] Reference forward produced finite logits for " + f"{token_count} tokens and vocab_size={vocab_size}" + ) + del repeated_reference_logits _checkpoint_paths(cfg) artifact_dir = _robustness_artifact_dir(cfg) if _rank0(): @@ -2027,18 +2562,23 @@ def _run_process_isolated_checkpoint_phase( json.dumps(trainable_digests, sort_keys=True) ) _barrier() - _report_phase("Isolated train/save: reference artifacts persisted; exiting phase") + _report_phase("Isolated Phase 1 train/save/reference: reference artifacts persisted; exiting phase") return if phase == "automodel_reload": - _report_phase("Isolated AutoModel reload: loading prompt input IDs") - input_ids = _load_input_ids_once(cfg, input_ids_loader, tokenizer_name) + _report_phase("Isolated Phase 2 AutoModel model reload: loading prompt input IDs") + input_ids = _load_input_ids_once( + cfg, + input_ids_loader, + tokenizer_name, + sequence_length=parity_sequence_length, + ) checkpoint_dir, ckpt_step_dir, consolidated_dir = _checkpoint_paths(cfg) reference_path = _robustness_artifact_dir(cfg) / "reference_logits.pt" assert reference_path.exists(), f"Reference logits not found at {reference_path}" is_peft = hasattr(cfg, "peft") - if custom_args.get("check_phantom_keys", False) and _rank0(): + if custom_args.get("check_phantom_keys", False) and _preinit_global_rank() == 0: from safetensors import safe_open assert consolidated_dir.exists(), f"Phantom key check: {consolidated_dir} does not exist" @@ -2051,21 +2591,14 @@ def _run_process_isolated_checkpoint_phase( assert "_scales" not in key, f"Phantom mxfp4 key leaked: {key} in {sf_path.name}" if not is_peft: - if _rank0(): - from transformers import AutoConfig - - _prepopulate_hf_dynamic_modules_cache(consolidated_dir) - try: - AutoConfig.from_pretrained(str(consolidated_dir), trust_remote_code=True) - except Exception: - pass - cfg.model.pretrained_model_name_or_path = str(consolidated_dir) + _prepare_consolidated_hf_cache_once(cfg, consolidated_dir) + _set_model_pretrained_path(cfg.model, consolidated_dir) cfg.checkpoint.enabled = False - _report_phase("Isolated AutoModel reload: starting trainer setup") + _report_phase("Isolated Phase 2 AutoModel model reload: starting trainer setup") restored_trainer = recipe_cls(cfg) restored_trainer.setup() - _report_phase("Isolated AutoModel reload: trainer setup complete") + _report_phase("Isolated Phase 2 AutoModel model reload: trainer setup complete") if tokenizer_name is not None and dist.is_initialized() and dist.get_world_size() > 1: _barrier() if _rank0(): @@ -2126,28 +2659,103 @@ def _run_process_isolated_checkpoint_phase( ) _raise_distributed_failure(failure_message) + reload_policy = _automodel_reload_parity_policy(custom_args) failure_message = None if _rank0(): reference_logits = torch.load(reference_path, map_location="cpu", weights_only=True) - max_kl_restored = _kl_divergence_from_logits(reference_logits, restored_logits).max().item() - tp_size = _tp_size_from_argv(sys.argv[1:]) - default_threshold = "1e-5" if tp_size > 1 else "0" - kl_threshold = float(custom_args.get("kl_threshold", default_threshold)) - print(f"\n[Isolated AutoModel reload] max KL: {max_kl_restored:.6e} (threshold: {kl_threshold:.6e})") - if custom_args.get("skip_automodel_logit_parity", False): - print( - "[Isolated AutoModel reload] Cross-process logit KL is informational; " - "exact trainable-parameter fingerprints are the checkpoint-integrity gate" + failure_message = _compare_logits( + _robustness_artifact_dir(cfg), + reference_logits, + restored_logits, + reload_policy, + ) + failure_message = _broadcast_rank0_failure(failure_message) + if reload_policy.profile == "relaxed" or not reload_policy.enforce or failure_message is not None: + repeated_restored_logits = _get_logits( + restored_trainer.model_parts[0], + input_ids, + device, + trainer=restored_trainer, + ) + if _rank0(): + _compare_logits( + _robustness_artifact_dir(cfg), + restored_logits, + repeated_restored_logits, + _repeatability_policy( + phase="phase_2", + comparison="automodel_reload_self_repeat", + profile=reload_policy.profile, + ), ) - elif max_kl_restored > kl_threshold: + del repeated_restored_logits + if failure_message is not None: + failure_message = ( + "CHECKPOINT_ROBUSTNESS_PHASE_FAILURE phase=automodel_reload check=full_logit_parity\n" + failure_message + ) + _raise_distributed_failure(failure_message) + _report_phase( + f"Isolated AutoModel reload: parity complete for {ckpt_step_dir.relative_to(checkpoint_dir)}; exiting phase" + ) + return + + if phase == "cross_tp_reload": + if hasattr(cfg, "peft"): + raise ValueError("Process-isolated cross_tp_reload does not support PEFT checkpoints") + _report_phase("Isolated Phase 5 cross-TP reload: loading prompt input IDs") + input_ids = _load_input_ids_once( + cfg, + input_ids_loader, + tokenizer_name, + sequence_length=parity_sequence_length, + ) + checkpoint_dir, ckpt_step_dir, consolidated_dir = _checkpoint_paths(cfg) + reference_path = _robustness_artifact_dir(cfg) / "reference_logits.pt" + if not reference_path.exists(): + raise FileNotFoundError(f"Reference logits not found at {reference_path}") + + _prepare_consolidated_hf_cache_once(cfg, consolidated_dir) + _set_model_pretrained_path(cfg.model, consolidated_dir) + cfg.checkpoint.enabled = False + cfg.distributed.tp_size = int(custom_args["cross_tp_size"]) + cfg.distributed.dp_size = None + + _report_phase("Isolated Phase 5 cross-TP reload: starting trainer setup") + cross_tp_trainer = recipe_cls(cfg) + cross_tp_trainer.setup() + _report_phase("Isolated Phase 5 cross-TP reload: trainer setup complete") + if tokenizer_name is not None and dist.is_initialized() and dist.get_world_size() > 1: + _barrier() + if _rank0(): + _cleanup_input_ids_sync(cfg) + _barrier() + + device = next(cross_tp_trainer.model_parts[0].parameters()).device + cross_tp_logits = _get_logits( + cross_tp_trainer.model_parts[0], + input_ids, + device, + trainer=cross_tp_trainer, + ) + failure_message = None + if _rank0(): + reference_logits = torch.load(reference_path, map_location="cpu", weights_only=True) + failure_message = _compare_logits( + _robustness_artifact_dir(cfg), + reference_logits, + cross_tp_logits, + _cross_tp_parity_policy(custom_args), + ) + if failure_message is not None: failure_message = ( - "CHECKPOINT_ROBUSTNESS_PHASE_FAILURE phase=automodel_reload check=logit_kl\n" - "KL divergence between original and AutoModel checkpoint reload too large: " - f"max per-token KL = {max_kl_restored:.6e} > threshold {kl_threshold:.6e}" + "CHECKPOINT_ROBUSTNESS_PHASE_FAILURE phase=cross_tp_reload check=full_logit_parity\n" + + failure_message ) _raise_distributed_failure(failure_message) + _release_recipe_memory(cross_tp_trainer) _report_phase( - f"Isolated AutoModel reload: parity complete for {ckpt_step_dir.relative_to(checkpoint_dir)}; exiting phase" + f"Isolated Phase 5 cross-TP reload: parity complete for " + f"{ckpt_step_dir.relative_to(checkpoint_dir)}; exiting phase" ) return @@ -2156,9 +2764,12 @@ def _run_process_isolated_checkpoint_phase( checkpoint_path = _checkpoint_for_completed_steps(resume_plan, resume_plan.boundary_step) _configure_resumed_run(cfg, resume_plan, checkpoint_path) - _report_phase("Isolated resume: starting setup and optimizer checkpoint load") + _report_phase("Isolated Phase 4 native resume: starting setup and optimizer checkpoint load") resume_trainer = recipe_cls(cfg) resume_trainer.setup() + # The restore is complete. Phase 4 validates the continuation but does not + # need to publish another large distributed checkpoint at its final step. + _disable_checkpoint_saves_after_restore(resume_trainer) restored_state = _checkpoint_state_snapshot(resume_trainer, state_is_being_saved=False) local_failure = _restored_state_mismatch(reference_trajectory["boundary_state"], restored_state) failure_message = _gather_rank_failures(local_failure, check="restored_state") @@ -2172,9 +2783,9 @@ def _run_process_isolated_checkpoint_phase( resume_recorder = _TrajectoryRecorder(resume_plan, capture_boundary_state=False) resume_recorder.attach(resume_trainer) - _report_phase("Isolated resume: checkpoint state verified; starting shared-trajectory continuation") + _report_phase("Isolated Phase 4 native resume: checkpoint state verified; starting shared-trajectory continuation") resume_trainer.run_train_validation_loop() - _report_phase("Isolated resume: training complete") + _report_phase("Isolated Phase 4 native resume: training complete") resumed_trajectory = resume_recorder.to_dict() resume_tolerance = _resolve_resume_loss_tolerance( @@ -2191,7 +2802,7 @@ def _run_process_isolated_checkpoint_phase( local_failure = comparison_report["blocking_failure"] failure_message = _gather_rank_failures(local_failure, check="shared_trajectory") _raise_distributed_failure(failure_message) - _report_phase("Isolated resume: shared-trajectory checkpoint continuation verified; exiting phase") + _report_phase("Isolated Phase 4 native resume: shared-trajectory checkpoint continuation verified; exiting phase") def run_checkpoint_robustness( @@ -2219,24 +2830,15 @@ def run_checkpoint_robustness( input_ids_loader=input_ids_loader, ) return - # When tensor parallelism is active the forward pass uses row-parallel - # all-reduces and cuBLASLt plan caches whose order of accumulation is - # process-dependent; this produces ULP-level bf16 drift between the - # trainer's and restored model's logits even with bit-identical weights. - # Use a small tolerance when TP>1; keep strict 0 otherwise so real - # save/load regressions in non-TP setups still fail. - _tp_size = _tp_size_from_argv(config_argv) - _default_kl_threshold = "1e-5" if _tp_size > 1 else "0" - kl_threshold = float(custom_args.get("kl_threshold", _default_kl_threshold)) cross_tp_size = int(custom_args.get("cross_tp_size", "0")) - cross_tp_kl_threshold = float(custom_args.get("cross_tp_kl_threshold", "5e-3")) - trust_remote_code = bool(custom_args.get("trust_remote_code", False)) + trust_remote_code = custom_args.get("trust_remote_code") experts_implementation = custom_args.get("experts_implementation", None) tokenizer_name = custom_args.get("tokenizer_name", None) + parity_sequence_length = int(custom_args.get("parity_sequence_length", "2048")) max_vram_gb = float(custom_args.get("max_vram_gb", "0")) max_cpu_gb = float(custom_args.get("max_cpu_gb", "0")) check_phantom_keys = bool(custom_args.get("check_phantom_keys", False)) - check_resume = bool(custom_args.get("check_resume", False)) + resume_enabled = bool(custom_args.get("resume_enabled", False)) resume_tolerance = _resolve_resume_loss_tolerance( custom_args.get("resume_tolerance_profile", "standard"), first_step_override=custom_args.get("resume_first_loss_threshold"), @@ -2246,20 +2848,22 @@ def run_checkpoint_robustness( hf_device_map_auto = bool(custom_args.get("hf_device_map_auto", False)) hf_source_post_load_dequantize = bool(custom_args.get("hf_source_post_load_dequantize", False)) skip_hf_reload = bool(custom_args.get("skip_hf_reload", False)) - check_source_load_parity = bool(custom_args.get("check_source_load_parity", False)) - source_load_kl_threshold = float(custom_args.get("source_load_kl_threshold", "5e-3")) - source_load_mean_kl_threshold = float(custom_args.get("source_load_mean_kl_threshold", "1e-3")) - source_load_cosine_threshold = float(custom_args.get("source_load_cosine_threshold", "0.9999")) + source_load_parity_enabled = bool(custom_args.get("source_load_parity_enabled", False)) deferred_failures: list[str] = [] cfg = parse_args_and_load_config() - resume_plan = _resume_plan_from_config(cfg) if check_resume else None + resume_plan = _resume_plan_from_config(cfg) if resume_enabled else None if resume_plan is not None: _configure_uninterrupted_run(cfg, resume_plan) - input_ids = _load_input_ids_once(cfg, input_ids_loader, tokenizer_name) + input_ids = _load_input_ids_once( + cfg, + input_ids_loader, + tokenizer_name, + sequence_length=parity_sequence_length, + ) source_load_reference = None - if check_source_load_parity: + if source_load_parity_enabled: _report_phase("Phase 0: starting vanilla-HF source-load reference") source_load_reference = _prepare_source_load_reference( cfg, @@ -2269,6 +2873,7 @@ def run_checkpoint_robustness( experts_implementation=experts_implementation, hf_device_map_auto=hf_device_map_auto, hf_source_post_load_dequantize=hf_source_post_load_dequantize, + parity_tolerance_profile=_comparison_profile(custom_args, "source_load"), ) _barrier() _report_phase("Phase 0: vanilla-HF source-load reference complete") @@ -2286,7 +2891,7 @@ def run_checkpoint_robustness( _cleanup_input_ids_sync(cfg) _barrier() - if check_source_load_parity: + if source_load_parity_enabled: _report_phase("Phase 0: starting constructed-trainer parity forward") device = next(trainer.model_parts[0].parameters()).device trainer_source_logits = _get_logits(trainer.model_parts[0], input_ids, device, trainer=trainer) @@ -2294,9 +2899,8 @@ def run_checkpoint_robustness( source_load_reference, trainer_source_logits, _lm_head_embedding_aliased(trainer.model_parts[0]), - source_load_kl_threshold=source_load_kl_threshold, - source_load_mean_kl_threshold=source_load_mean_kl_threshold, - source_load_cosine_threshold=source_load_cosine_threshold, + artifact_dir=_robustness_artifact_dir(cfg), + policy=_source_load_parity_policy(custom_args), ) _record_deferred_failure(deferred_failures, "Phase 0 source-load parity", source_load_failure) del trainer_source_logits, source_load_reference @@ -2367,11 +2971,29 @@ def run_checkpoint_robustness( if max_cpu_gb > 0: assert peak_cpu_gb <= max_cpu_gb, f"Peak CPU RSS {peak_cpu_gb:.2f} GB exceeds threshold {max_cpu_gb:.2f} GB" - # Phase 2: Capture reference logits before teardown - _report_phase("Phase 2: starting reference-logits capture") + # Phase 1 also captures the reference distribution before teardown. It is + # persisted by the isolated runner for the independent reload processes. + _report_phase("Phase 1: starting reference-logits capture") device = next(trainer.model_parts[0].parameters()).device reference_logits = _get_logits(trainer.model_parts[0], input_ids, device, trainer=trainer) - _report_phase("Phase 2: reference-logits capture complete") + token_count, vocab_size = _validate_logits(reference_logits) + repeated_reference_logits = _get_logits(trainer.model_parts[0], input_ids, device, trainer=trainer) + if _rank0(): + _compare_logits( + _robustness_artifact_dir(cfg), + reference_logits, + repeated_reference_logits, + _repeatability_policy( + phase="phase_1", + comparison="automodel_reference_self_repeat", + profile=str(custom_args.get("parity_tolerance_profile", "standard")), + ), + ) + print( + f"[Phase 1] Reference forward produced finite logits for {token_count} tokens and vocab_size={vocab_size}" + ) + del repeated_reference_logits + _report_phase("Phase 1: reference-logits capture complete") # Locate the Phase 1 checkpoint used by the reload and resume checks. if resume_plan is not None: @@ -2385,7 +3007,7 @@ def run_checkpoint_robustness( _release_recipe_memory(trainer) del trainer - # Phase 3: Reload AutoModel from the consolidated checkpoint. + # Phase 2: Reload AutoModel from the exported HF-format consolidated weights. # Phantom key check: scan consolidated safetensors for leaked quantization keys if check_phantom_keys and _rank0(): from safetensors import safe_open @@ -2420,40 +3042,62 @@ def run_checkpoint_robustness( cfg = parse_args_and_load_config() if not is_peft: - cfg.model.pretrained_model_name_or_path = str(consolidated_dir) + _set_model_pretrained_path(cfg.model, consolidated_dir) cfg.checkpoint.enabled = False - _report_phase("Phase 3: starting AutoModel reload setup") + _report_phase("Phase 2: starting AutoModel model reload setup") restored_trainer = recipe_cls(cfg) restored_trainer.setup() - _report_phase("Phase 3: AutoModel reload setup complete") + _report_phase("Phase 2: AutoModel model reload setup complete") - _report_phase("Phase 3: starting restored-logits capture") + _report_phase("Phase 2: starting restored-logits capture") restored_logits = _get_logits(restored_trainer.model_parts[0], input_ids, device, trainer=restored_trainer) - _report_phase("Phase 3: restored-logits capture complete") + _report_phase("Phase 2: restored-logits capture complete") - kl_restored = _kl_divergence_from_logits(reference_logits, restored_logits) - max_kl_restored = kl_restored.max().item() - if _rank0(): - print(f"\n[Phase 3] Automodel-from-consolidated max KL: {max_kl_restored:.6e} (threshold: {kl_threshold:.6e})") + reload_policy = _automodel_reload_parity_policy(custom_args) automodel_reload_error = None - if max_kl_restored > kl_threshold: - automodel_reload_error = ( - "KL divergence between original and automodel-from-consolidated too large: " - f"max per-token KL = {max_kl_restored:.6e} > threshold {kl_threshold:.6e}" + if _rank0(): + automodel_reload_error = _compare_logits( + _robustness_artifact_dir(cfg), + reference_logits, + restored_logits, + reload_policy, ) - _record_deferred_failure(deferred_failures, "Phase 3 AutoModel reload parity", automodel_reload_error) + automodel_reload_error = _broadcast_rank0_failure(automodel_reload_error) + if reload_policy.profile == "relaxed" or not reload_policy.enforce or automodel_reload_error is not None: + repeated_restored_logits = _get_logits( + restored_trainer.model_parts[0], + input_ids, + device, + trainer=restored_trainer, + ) + if _rank0(): + _compare_logits( + _robustness_artifact_dir(cfg), + restored_logits, + repeated_restored_logits, + _repeatability_policy( + phase="phase_2", + comparison="automodel_reload_self_repeat", + profile=reload_policy.profile, + ), + ) + del repeated_restored_logits + _record_deferred_failure(deferred_failures, "Phase 2 AutoModel model reload parity", automodel_reload_error) _release_recipe_memory(restored_trainer) del restored_trainer - # Phase 4: Load into vanilla HF (rank 0 only) - _report_phase("Phase 4: starting vanilla-HF reload") - hf_reload_sync_paths = _prepare_hf_reload_sync(cfg) + # Phase 3: Load the same exported weights into vanilla HF (rank 0 only). + _report_phase("Phase 3: starting vanilla-HF export reload") + hf_reload_timeout_s = ( + int(custom_args["hf_reload_timeout_seconds"]) if "hf_reload_timeout_seconds" in custom_args else None + ) + hf_reload_sync_paths = _prepare_hf_reload_sync(cfg, timeout_s=hf_reload_timeout_s) hf_reload_error = None if skip_hf_reload: if _rank0(): - print("[Phase 4] Skipped (ci.checkpoint_robustness.skip_hf_reload=true).") + print("[Phase 3] Skipped (ci.checkpoint_robustness.skip_hf_reload=true).") elif _rank0(): hf_reload_error = _run_vanilla_hf_reload( cfg, @@ -2464,52 +3108,22 @@ def run_checkpoint_robustness( ) hf_reload_error = _finish_hf_reload_sync(hf_reload_sync_paths, hf_reload_error) - _record_deferred_failure(deferred_failures, "Phase 4 HF reload parity", hf_reload_error) - _report_phase("Phase 4: vanilla-HF reload complete") - - # Phase 5 (optional): Cross-TP — reload consolidated with a different TP size - if cross_tp_size > 0 and not is_peft: - _report_phase("Phase 5: starting cross-TP reload") - cfg = parse_args_and_load_config() - cfg.model.pretrained_model_name_or_path = str(consolidated_dir) - cfg.checkpoint.enabled = False - cfg.distributed.tp_size = cross_tp_size - cfg.distributed.dp_size = None - cross_tp_trainer = recipe_cls(cfg) - cross_tp_trainer.setup() - - cross_tp_logits = _get_logits(cross_tp_trainer.model_parts[0], input_ids, device, trainer=cross_tp_trainer) - - kl_cross_tp = _kl_divergence_from_logits(reference_logits, cross_tp_logits) - max_kl_cross_tp = kl_cross_tp.max().item() - if _rank0(): - print( - f"[Phase 5] Cross-TP (tp_size={cross_tp_size}) max KL: " - f"{max_kl_cross_tp:.6e} (threshold: {cross_tp_kl_threshold:.6e})" - ) - cross_tp_error = None - if max_kl_cross_tp > cross_tp_kl_threshold: - cross_tp_error = ( - "KL divergence between original and cross-TP model too large: " - f"max per-token KL = {max_kl_cross_tp:.6e} > threshold {cross_tp_kl_threshold:.6e}" - ) - _record_deferred_failure(deferred_failures, "Phase 5 cross-TP reload parity", cross_tp_error) - - _release_recipe_memory(cross_tp_trainer) - del cross_tp_trainer - _barrier() - _report_phase("Phase 5: cross-TP reload complete") + _record_deferred_failure(deferred_failures, "Phase 3 vanilla-HF export reload parity", hf_reload_error) + _report_phase("Phase 3: vanilla-HF export reload complete") - # Phase 6 (optional): restore the exact Phase 1 boundary and replay its continuation. - if check_resume: + # Phase 4: restore the exact Phase 1 boundary and replay its continuation. + if resume_enabled: assert resume_plan is not None reference_trajectory = _load_reference_trajectory(resume_plan) checkpoint_path = _checkpoint_for_completed_steps(resume_plan, resume_plan.boundary_step) cfg = parse_args_and_load_config() _configure_resumed_run(cfg, resume_plan, checkpoint_path) - _report_phase("Phase 6: starting resume setup and checkpoint load") + _report_phase("Phase 4: starting native-checkpoint resume setup and load") resume_trainer = recipe_cls(cfg) resume_trainer.setup() + # The restore is complete. Phase 4 validates the continuation but does not + # need to publish another large distributed checkpoint at its final step. + _disable_checkpoint_saves_after_restore(resume_trainer) restored_state = _checkpoint_state_snapshot(resume_trainer, state_is_being_saved=False) local_failure = _restored_state_mismatch(reference_trajectory["boundary_state"], restored_state) failure_message = _gather_rank_failures(local_failure, check="restored_state") @@ -2523,9 +3137,9 @@ def run_checkpoint_robustness( resumed_recorder = _TrajectoryRecorder(resume_plan, capture_boundary_state=False) resumed_recorder.attach(resume_trainer) - _report_phase("Phase 6: checkpoint state verified; starting shared-trajectory continuation") + _report_phase("Phase 4: checkpoint state verified; starting shared-trajectory continuation") resume_trainer.run_train_validation_loop() - _report_phase("Phase 6: resumed training complete") + _report_phase("Phase 4: resumed training complete") resumed_trajectory = resumed_recorder.to_dict() comparison_report = _report_resume_comparison( @@ -2549,7 +3163,35 @@ def run_checkpoint_robustness( _release_recipe_memory(resume_trainer) del resume_trainer _barrier() - _report_phase("Phase 6: resume comparison complete") + _report_phase("Phase 4: resume comparison complete") + + # Phase 5 (optional): reload the exported weights with a different TP size. + if cross_tp_size > 0 and not is_peft: + _report_phase("Phase 5: starting optional cross-TP consolidated reload") + cfg = parse_args_and_load_config() + _set_model_pretrained_path(cfg.model, consolidated_dir) + cfg.checkpoint.enabled = False + cfg.distributed.tp_size = cross_tp_size + cfg.distributed.dp_size = None + cross_tp_trainer = recipe_cls(cfg) + cross_tp_trainer.setup() + + cross_tp_logits = _get_logits(cross_tp_trainer.model_parts[0], input_ids, device, trainer=cross_tp_trainer) + cross_tp_error = None + if _rank0(): + cross_tp_error = _compare_logits( + _robustness_artifact_dir(cfg), + reference_logits, + cross_tp_logits, + _cross_tp_parity_policy(custom_args), + ) + cross_tp_error = _broadcast_rank0_failure(cross_tp_error) + _record_deferred_failure(deferred_failures, "Phase 5 cross-TP consolidated reload parity", cross_tp_error) + + _release_recipe_memory(cross_tp_trainer) + del cross_tp_trainer + _barrier() + _report_phase("Phase 5: optional cross-TP consolidated reload complete") # Skip the atexit-registered destroy_process_group() call. MoE models with expert # parallelism create NCCL sub-groups (DeepEP) that leave pending collective state, diff --git a/tests/functional_tests/checkpoint_robustness/test_checkpoint_robustness_vlm.py b/tests/functional_tests/checkpoint_robustness/test_checkpoint_robustness_vlm.py index 28b09a1b36..b2ff39caa7 100644 --- a/tests/functional_tests/checkpoint_robustness/test_checkpoint_robustness_vlm.py +++ b/tests/functional_tests/checkpoint_robustness/test_checkpoint_robustness_vlm.py @@ -17,8 +17,7 @@ import os from tests.functional_tests.checkpoint_robustness.test_checkpoint_robustness_llm import ( - _DEFAULT_INPUT_IDS, - _DEFAULT_PROMPT, + _get_parity_document, run_checkpoint_robustness, ) @@ -26,7 +25,7 @@ def _get_vlm_input_ids(processor_name: str | None) -> list[int]: """Encode the parity prompt through the same processor family used by VLM data loading.""" if processor_name is None: - return _DEFAULT_INPUT_IDS + raise ValueError("tokenizer_name is required to tokenize the checkpoint parity document") from transformers import AutoProcessor @@ -36,7 +35,7 @@ def _get_vlm_input_ids(processor_name: str | None) -> list[int]: local_files_only=os.environ.get("HF_HUB_OFFLINE", "0") == "1", ) tokenizer = getattr(processor, "tokenizer", processor) - return tokenizer.encode(_DEFAULT_PROMPT, add_special_tokens=False) + return tokenizer.encode(_get_parity_document(), add_special_tokens=False) def test_checkpoint_robustness_vlm() -> None: diff --git a/tests/functional_tests/conftest.py b/tests/functional_tests/conftest.py index 7760a5fdd5..e7bca7c11a 100644 --- a/tests/functional_tests/conftest.py +++ b/tests/functional_tests/conftest.py @@ -119,28 +119,29 @@ "config_path", "deploy_mode", "max_new_tokens", - "hf_adapter_ignored_key_prefix", - "resume_tolerance_profile", -] - -_BOOLEAN_OVERRIDES = [ - "vllm_smoke_test", - "kl_threshold", - "hf_kl_threshold", + "isolated_phase", "cross_tp_size", - "cross_tp_kl_threshold", - "tokenizer", "experts_implementation", "tokenizer_name", "max_vram_gb", "max_cpu_gb", + "hf_adapter_ignored_key_prefix", + "hf_device_map_max_memory_gib", + "hf_device_map_cpu_max_memory_gib", + "hf_reload_timeout_seconds", "training_reproducibility_loss_threshold", + "parity_sequence_length", + "parity_threshold_overrides", + "parity_tolerance_profile", + "parity_tolerance_profile_overrides", + "resume_tolerance_profile", "resume_first_loss_threshold", "resume_loss_threshold", - "source_load_cosine_threshold", - "source_load_kl_threshold", - "source_load_mean_kl_threshold", - "cosine_threshold", +] + +_BOOLEAN_OVERRIDES = [ + "vllm_smoke_test", + "tokenizer", "dataset.data_dir_list", "dataloader.dataset.data_dir_list", "tokenizer._target_", @@ -149,12 +150,14 @@ "trust_remote_code", "check_fused_qkv_keys", "check_phantom_keys", - "check_source_load_parity", - "check_resume", "hf_device_map_auto", "hf_source_post_load_dequantize", - "skip_automodel_logit_parity", - "skip_hf_logit_parity", + "skip_resume", + "skip_source_load_parity", + "skip_source_load_logit_parity", + "skip_hf_reload", + "skip_automodel_reload_logit_parity", + "skip_hf_reload_logit_parity", ] diff --git a/tests/unit_tests/ci_tests/test_checkpoint_parity_metrics.py b/tests/unit_tests/ci_tests/test_checkpoint_parity_metrics.py new file mode 100644 index 0000000000..089bc87d12 --- /dev/null +++ b/tests/unit_tests/ci_tests/test_checkpoint_parity_metrics.py @@ -0,0 +1,241 @@ +# Copyright (c) 2026, NVIDIA CORPORATION. All rights reserved. +# +# 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. + +import pytest +import torch +import torch.nn.functional as F + +from tests.functional_tests.checkpoint_robustness.parity_metrics import ( + _apply_parity_threshold_overrides, + _compute_parity_metrics, + _normalize_parity_profile_overrides, + _normalize_parity_threshold_overrides, + _parity_failures, + _resolve_parity_thresholds, + _select_parity_profile, + _validate_logits, +) + + +def test_identical_logits_have_zero_divergence_and_unit_cosine(): + logits = torch.tensor([[[2.0, 0.0, -1.0], [0.5, 0.25, -0.5]]]) + + metrics = _compute_parity_metrics(logits, logits.clone(), chunk_tokens=1) + + assert metrics.token_count == 2 + assert metrics.vocab_size == 3 + assert metrics.mean_kl == pytest.approx(0.0, abs=1e-8) + assert metrics.p95_kl == pytest.approx(0.0, abs=1e-8) + assert metrics.max_kl == pytest.approx(0.0, abs=1e-8) + assert metrics.cosine_similarity == pytest.approx(1.0) + assert metrics.mean_absolute_logit_difference == 0.0 + assert metrics.max_absolute_logit_difference == 0.0 + + +def test_full_logit_metrics_match_direct_reference_computation(): + reference = torch.tensor([[[1.0, -0.5, 0.25], [0.0, 2.0, -1.0]]]) + candidate = torch.tensor([[[0.75, -0.25, 0.0], [0.5, 1.25, -0.5]]]) + reference_flat = reference.reshape(-1, 3).float() + candidate_flat = candidate.reshape(-1, 3).float() + reference_log_probs = F.log_softmax(reference_flat, dim=-1) + expected_kl = (reference_log_probs.exp() * (reference_log_probs - F.log_softmax(candidate_flat, dim=-1))).sum(-1) + + metrics = _compute_parity_metrics(reference, candidate, chunk_tokens=1) + + assert metrics.mean_kl == pytest.approx(expected_kl.mean().item(), rel=1e-6, abs=1e-8) + assert metrics.p95_kl == pytest.approx(torch.quantile(expected_kl, 0.95).item(), rel=1e-6, abs=1e-8) + assert metrics.max_kl == pytest.approx(expected_kl.max().item(), rel=1e-6, abs=1e-8) + assert metrics.cosine_similarity == pytest.approx( + F.cosine_similarity(reference.flatten(), candidate.flatten(), dim=0).item(), rel=1e-6 + ) + absolute_difference = (reference - candidate).abs() + assert metrics.mean_absolute_logit_difference == pytest.approx(absolute_difference.mean().item()) + assert metrics.max_absolute_logit_difference == pytest.approx(absolute_difference.max().item()) + + +def test_p95_is_stable_against_a_single_token_outlier_while_max_remains_diagnostic(): + reference = torch.zeros(1, 100, 2) + candidate = reference.clone() + candidate[0, -1] = torch.tensor([20.0, -20.0]) + + metrics = _compute_parity_metrics(reference, candidate) + + assert metrics.mean_kl > 0.0 + assert metrics.p95_kl == pytest.approx(0.0, abs=1e-8) + assert metrics.max_kl > 1.0 + + +def test_metric_results_do_not_depend_on_chunk_size(): + generator = torch.Generator().manual_seed(1234) + reference = torch.randn(2, 7, 11, generator=generator) + candidate = reference + 0.01 * torch.randn(2, 7, 11, generator=generator) + + single_token_chunks = _compute_parity_metrics(reference, candidate, chunk_tokens=1) + all_token_chunk = _compute_parity_metrics(reference, candidate, chunk_tokens=14) + + assert single_token_chunks.to_dict() == pytest.approx(all_token_chunk.to_dict(), rel=1e-6, abs=1e-8) + + +@pytest.mark.parametrize("bad_value", [float("nan"), float("inf"), float("-inf")]) +def test_non_finite_logits_are_rejected(bad_value): + logits = torch.zeros(1, 2, 3) + logits[0, 1, 2] = bad_value + + with pytest.raises(ValueError, match="non-finite"): + _validate_logits(logits) + with pytest.raises(ValueError, match="non-finite"): + _compute_parity_metrics(torch.zeros_like(logits), logits) + + +def test_named_profiles_are_ordered_and_gate_mean_p95_and_cosine(): + strict = _resolve_parity_thresholds("strict", "cross_framework") + standard = _resolve_parity_thresholds("standard", "cross_framework") + relaxed = _resolve_parity_thresholds("relaxed", "cross_framework") + reference = torch.zeros(1, 100, 2) + candidate = reference.clone() + candidate[:, :10, 0] = 0.25 + metrics = _compute_parity_metrics(reference, candidate) + + assert strict.mean_kl < standard.mean_kl < relaxed.mean_kl + assert strict.p95_kl < standard.p95_kl < relaxed.p95_kl + assert strict.cosine_similarity > standard.cosine_similarity > relaxed.cosine_similarity + failures = _parity_failures(metrics, strict) + assert any("mean KL" in failure for failure in failures) + assert any("p95 KL" in failure for failure in failures) + assert any("cosine similarity" in failure for failure in failures) + + +@pytest.mark.parametrize("profile", ["strict", "standard", "relaxed"]) +def test_comparison_kinds_never_make_cross_topology_stricter_than_same_implementation(profile): + same_implementation = _resolve_parity_thresholds(profile, "same_implementation") + cross_topology = _resolve_parity_thresholds(profile, "cross_topology") + cross_framework = _resolve_parity_thresholds(profile, "cross_framework") + + assert same_implementation.mean_kl <= cross_topology.mean_kl <= cross_framework.mean_kl + assert same_implementation.p95_kl <= cross_topology.p95_kl <= cross_framework.p95_kl + assert ( + same_implementation.cosine_similarity >= cross_topology.cosine_similarity >= cross_framework.cosine_similarity + ) + + +@pytest.mark.parametrize( + ("profile", "comparison_kind", "expected_mean_kl", "expected_p95_kl", "expected_cosine"), + [ + ("standard", "same_implementation", 3e-3, 1.2e-2, 0.999), + ("standard", "cross_framework", 6e-3, 3e-2, 0.998), + ("standard", "cross_topology", 6e-3, 3e-2, 0.998), + ("relaxed", "same_implementation", 2e-2, 5e-2, 0.995), + ("relaxed", "cross_framework", 2.5e-2, 1e-1, 0.99), + ("relaxed", "cross_topology", 2e-2, 5e-2, 0.995), + ], +) +def test_calibrated_profile_thresholds(profile, comparison_kind, expected_mean_kl, expected_p95_kl, expected_cosine): + thresholds = _resolve_parity_thresholds(profile, comparison_kind) + + assert thresholds.mean_kl == expected_mean_kl + assert thresholds.p95_kl == expected_p95_kl + assert thresholds.cosine_similarity == expected_cosine + + +def test_selected_numeric_overrides_preserve_other_profile_gates(): + relaxed = _resolve_parity_thresholds("relaxed", "same_implementation") + + overridden = _apply_parity_threshold_overrides(relaxed, mean_kl=4e-2, cosine_similarity=0.99) + + assert overridden.mean_kl == 4e-2 + assert overridden.p95_kl == relaxed.p95_kl + assert overridden.cosine_similarity == 0.99 + + +def test_structured_profile_overrides_accept_every_comparison(): + overrides = _normalize_parity_profile_overrides( + { + "source_load": "strict", + "automodel_reload": "standard", + "hf_reload": "relaxed", + "cross_tp": "standard", + } + ) + + assert overrides == { + "source_load": "strict", + "automodel_reload": "standard", + "hf_reload": "relaxed", + "cross_tp": "standard", + } + + +def test_comparison_profile_override_falls_back_to_global_profile(): + overrides = {"hf_reload": "relaxed"} + + assert _select_parity_profile("standard", overrides, "hf_reload") == "relaxed" + assert _select_parity_profile("standard", overrides, "source_load") == "standard" + + +@pytest.mark.parametrize( + ("overrides", "error"), + [ + ({"unknown": "relaxed"}, "Unknown parity_tolerance_profile_overrides comparisons"), + ({"hf_reload": 1}, "hf_reload must be a profile name"), + ({"hf_reload": "custom"}, "Unknown parity tolerance profile"), + ], +) +def test_structured_profile_overrides_reject_invalid_schema(overrides, error): + with pytest.raises(ValueError, match=error): + _normalize_parity_profile_overrides(overrides) + + +def test_structured_threshold_overrides_accept_partial_gates_for_every_comparison(): + overrides = _normalize_parity_threshold_overrides( + { + "source_load": {"mean_kl": 0.01}, + "automodel_reload": {"p95_kl": 0.02}, + "hf_reload": {"cosine_similarity": 0.995}, + "cross_tp": {"mean_kl": 0.03, "p95_kl": 0.04}, + } + ) + + assert overrides == { + "source_load": {"mean_kl": 0.01}, + "automodel_reload": {"p95_kl": 0.02}, + "hf_reload": {"cosine_similarity": 0.995}, + "cross_tp": {"mean_kl": 0.03, "p95_kl": 0.04}, + } + + +@pytest.mark.parametrize( + ("overrides", "error"), + [ + ({"unknown": {"mean_kl": 0.01}}, "Unknown parity_threshold_overrides comparisons"), + ({"source_load": {"max_kl": 0.01}}, "Unknown parity_threshold_overrides.source_load metrics"), + ({"hf_reload": {"mean_kl": "0.01"}}, "hf_reload.mean_kl must be numeric"), + ({"cross_tp": {"cosine_similarity": 2.0}}, "cosine_similarity threshold override"), + ], +) +def test_structured_threshold_overrides_reject_invalid_schema(overrides, error): + with pytest.raises(ValueError, match=error): + _normalize_parity_threshold_overrides(overrides) + + +@pytest.mark.parametrize("bad_threshold", [float("nan"), float("inf"), -1.0]) +def test_invalid_profile_threshold_override_is_rejected(bad_threshold): + thresholds = _resolve_parity_thresholds("relaxed", "same_implementation") + + with pytest.raises(ValueError, match="finite and non-negative"): + _apply_parity_threshold_overrides(thresholds, mean_kl=bad_threshold) + + +def test_unknown_profile_is_rejected(): + with pytest.raises(ValueError, match="Unknown parity tolerance profile"): + _resolve_parity_thresholds("custom", "cross_framework") diff --git a/tests/unit_tests/ci_tests/test_checkpoint_robustness_hf_kwargs.py b/tests/unit_tests/ci_tests/test_checkpoint_robustness_hf_kwargs.py index 91fdac3889..9f4cfd6130 100644 --- a/tests/unit_tests/ci_tests/test_checkpoint_robustness_hf_kwargs.py +++ b/tests/unit_tests/ci_tests/test_checkpoint_robustness_hf_kwargs.py @@ -12,41 +12,58 @@ # See the License for the specific language governing permissions and # limitations under the License. +import json from contextlib import nullcontext +from copy import deepcopy from types import SimpleNamespace from unittest.mock import Mock, patch import pytest import torch +from transformers import AutoModelForCausalLM from tests.functional_tests.checkpoint_robustness.test_checkpoint_robustness_biencoder import ( _extract_custom_args as _extract_biencoder_custom_args, ) from tests.functional_tests.checkpoint_robustness.test_checkpoint_robustness_llm import ( + _PARITY_DOCUMENT_SHA256, _assert_peft_adapter_matches_checkpoint, + _automodel_reload_parity_policy, + _compare_logits, _compare_source_load_parity, + _cross_tp_parity_policy, _dequantize_hf_fp8_weights_in_place, _extract_custom_args, _finish_hf_reload_sync, _get_input_ids, _get_logits_pp, + _get_parity_document, _hf_device_map_max_memory, _hf_fp32_module_names, _hf_model_load_context, - _hf_reload_kl_error, + _hf_reload_parity_policy, _hf_source_load_kwargs, _keep_hf_modules_in_fp32, _lm_head_embedding_aliased, _load_hf_fp8_dequantized_config, _load_input_ids_once, + _LogitParityPolicy, + _model_pretrained_path, _normalize_peft_no_split_modules, _patch_remote_masking_api_compatibility, _peft_adapter_load_kwargs, _post_load_dequant_max_memory, + _prepare_consolidated_hf_cache_once, _raise_distributed_failure, _record_deferred_failure, + _repair_legacy_partial_rotary_config, + _repeatability_policy, + _resolve_hf_attn_implementation, _resolve_hf_model_class, _run_process_isolated_checkpoint_phase, + _run_vanilla_hf_reload, + _set_model_pretrained_path, + _source_load_parity_policy, _trainable_parameter_digests, _wait_for_hf_reload_rank0, _wait_for_source_load_artifacts, @@ -55,6 +72,42 @@ from tests.functional_tests.checkpoint_robustness.test_checkpoint_vllm_deploy import _tokenize_for_generation +def test_model_pretrained_path_resolves_direct_recipe_path(): + model_cfg = SimpleNamespace(pretrained_model_name_or_path="org/direct-model") + + assert _model_pretrained_path(model_cfg) == "org/direct-model" + + +def test_model_pretrained_path_resolves_config_based_recipe_path(): + model_cfg = SimpleNamespace( + config=SimpleNamespace( + pretrained_model_name_or_path="org/config-model", + name_or_path="org/config-model", + ) + ) + + assert _model_pretrained_path(model_cfg) == "org/config-model" + + +def test_set_model_pretrained_path_retargets_config_based_recipe(): + nested_config = SimpleNamespace( + pretrained_model_name_or_path="org/source-model", + name_or_path="org/source-model", + ) + model_cfg = SimpleNamespace(config=nested_config) + + _set_model_pretrained_path(model_cfg, "/tmp/exported-model") + + assert nested_config.pretrained_model_name_or_path == "/tmp/exported-model" + assert nested_config.name_or_path == "/tmp/exported-model" + assert not hasattr(model_cfg, "pretrained_model_name_or_path") + + +def test_model_pretrained_path_rejects_recipe_without_source_path(): + with pytest.raises(ValueError, match="model.config.pretrained_model_name_or_path"): + _model_pretrained_path(SimpleNamespace()) + + def test_resolve_hf_model_class_uses_advertised_causal_lm_for_vlm_checkpoint(): from transformers import AutoModelForCausalLM, AutoModelForImageTextToText @@ -70,6 +123,18 @@ def test_resolve_hf_model_class_uses_advertised_causal_lm_for_vlm_checkpoint(): assert resolved_cls is AutoModelForCausalLM +def test_resolve_hf_model_class_uses_native_image_text_mapping_for_mistral3(): + from transformers import AutoModelForImageTextToText + + with patch( + "transformers.PretrainedConfig.get_config_dict", + return_value=({"model_type": "mistral3", "architectures": ["Mistral3ForConditionalGeneration"]}, {}), + ): + resolved_cls = _resolve_hf_model_class("model-path", AutoModelForCausalLM) + + assert resolved_cls is AutoModelForImageTextToText + + def test_hf_device_map_max_memory_caps_each_visible_gpu(): with patch("torch.cuda.device_count", return_value=8): max_memory = _hf_device_map_max_memory("55") @@ -277,8 +342,7 @@ def create_mask(config, inputs_embeds, attention_mask, past_key_values, cache_po assert masking_utils.create_sliding_window_causal_mask is create_mask -@pytest.mark.parametrize("metadata_api", ["legacy", "user"]) -def test_get_logits_pp_pads_prompt_to_static_stage_sequence_length(metadata_api): +def test_get_logits_pp_updates_pipeline_sequence_length(): class _Schedule: def __init__(self): self._loss_fn = None @@ -286,7 +350,7 @@ def __init__(self): self.attention_mask = None def eval(self, ids, *, target, losses, attention_mask): - """Capture a padded pipeline batch and invoke the active loss callback. + """Capture a pipeline batch and invoke the active loss callback. Args: ids: Tensor of shape [batch, sequence]. @@ -309,19 +373,13 @@ def get_group(): def size(): return 1 - if metadata_api == "legacy": - stage = SimpleNamespace(inputs_meta=(torch.empty(1, 16),)) - else: - tensor_meta = SimpleNamespace(shape=torch.Size([1, 16])) - stage = SimpleNamespace(_user_meta=SimpleNamespace(inputs=(tensor_meta,), outputs=())) - schedule = _Schedule() + update_seq_len = Mock() trainer = SimpleNamespace( pp=SimpleNamespace( - pp_seq_len=None, + update_seq_len=update_seq_len, info=SimpleNamespace( schedule=schedule, - stages=[stage], has_first_stage=True, has_last_stage=True, ), @@ -329,7 +387,6 @@ def size(): pipeline_config=SimpleNamespace(pp_batch_size=1), model_parts=[SimpleNamespace(eval=lambda: None, config=SimpleNamespace(vocab_size=7))], device_mesh={"pp": _PipelineMesh()}, - cfg=SimpleNamespace(get=lambda *_args: None), ) with ( @@ -341,48 +398,60 @@ def size(): ): logits = _get_logits_pp(trainer, [11, 12, 13], torch.device("cpu")) - assert schedule.ids.tolist() == [[11, 12, 13] + [0] * 13] - assert schedule.attention_mask.shape == (1, 16) - assert schedule.attention_mask.tolist() == [[1, 1, 1] + [0] * 13] + update_seq_len.assert_called_once_with(3) + assert schedule.ids.tolist() == [[11, 12, 13]] + assert schedule.attention_mask.shape == (1, 3) + assert schedule.attention_mask.tolist() == [[1, 1, 1]] assert logits.shape == (1, 3, 7) @pytest.mark.parametrize( ("model_type", "expected_attn_implementation"), - [("nemotron_h", "eager"), ("step3p7", "eager"), ("nemotron_flash", "flash_attention_2")], + [ + ("deepseek_v4", "eager"), + ("nemotron-nas", "eager"), + ("nemotron_h", "eager"), + ("step3p7", "eager"), + ("nemotron_flash", "eager"), + ], ) def test_remote_code_attention_implementation(model_type, expected_attn_implementation): with patch( - "transformers.AutoConfig.from_pretrained", - return_value=SimpleNamespace(model_type=model_type), - ) as from_pretrained: + "transformers.PretrainedConfig.get_config_dict", + return_value=({"model_type": model_type}, {}), + ) as get_config_dict: hf_kwargs = _hf_source_load_kwargs( {"revision": "model-revision", "token": "model-token"}, pretrained_model_name_or_path="model-path", source_dtype=torch.bfloat16, trust_remote_code=True, experts_implementation=None, + hf_model_cls=AutoModelForCausalLM, device=torch.device("cpu"), hf_device_map_auto=False, ) assert hf_kwargs["attn_implementation"] == expected_attn_implementation - from_pretrained.assert_called_once_with( + get_config_dict.assert_called_once_with( "model-path", - trust_remote_code=True, + local_files_only=False, revision="model-revision", token="model-token", ) def test_explicit_attention_implementation_is_preserved(): - with patch("transformers.AutoConfig.from_pretrained", side_effect=AssertionError("must not probe config")): + with patch( + "transformers.PretrainedConfig.get_config_dict", + return_value=({"model_type": "unknown_remote_model"}, {}), + ): hf_kwargs = _hf_source_load_kwargs( {"attn_implementation": "eager"}, pretrained_model_name_or_path="model-path", source_dtype=torch.bfloat16, trust_remote_code=True, experts_implementation=None, + hf_model_cls=AutoModelForCausalLM, device=torch.device("cpu"), hf_device_map_auto=False, ) @@ -390,6 +459,39 @@ def test_explicit_attention_implementation_is_preserved(): assert hf_kwargs["attn_implementation"] == "eager" +@pytest.mark.parametrize(("supported", "expected"), [(True, "sdpa"), (False, "eager")]) +def test_builtin_attention_implementation_uses_supported_recipe_backend_or_eager(supported, expected): + class FakeConfig: + pass + + concrete_model_cls = SimpleNamespace(_supports_sdpa=supported) + auto_model_cls = SimpleNamespace(_model_mapping={FakeConfig: concrete_model_cls}) + with patch("transformers.AutoConfig.from_pretrained", return_value=FakeConfig()): + implementation = _resolve_hf_attn_implementation( + "model-path", + "sdpa", + hf_model_cls=auto_model_cls, + trust_remote_code=False, + ) + + assert implementation == expected + + +def test_hf_source_load_kwargs_explicit_false_disables_recipe_remote_code(): + hf_kwargs = _hf_source_load_kwargs( + {"trust_remote_code": True}, + pretrained_model_name_or_path="model-path", + source_dtype=torch.bfloat16, + trust_remote_code=False, + experts_implementation=None, + hf_model_cls=AutoModelForCausalLM, + device=torch.device("cpu"), + hf_device_map_auto=False, + ) + + assert hf_kwargs["trust_remote_code"] is False + + def test_hf_source_load_kwargs_passes_grouped_experts_implementation(): hf_kwargs = _hf_source_load_kwargs( {}, @@ -397,6 +499,7 @@ def test_hf_source_load_kwargs_passes_grouped_experts_implementation(): source_dtype=torch.bfloat16, trust_remote_code=False, experts_implementation="grouped_mm", + hf_model_cls=AutoModelForCausalLM, device=torch.device("cpu"), hf_device_map_auto=False, ) @@ -475,6 +578,7 @@ def test_hf_source_load_kwargs_respects_hf_offline(monkeypatch, offline, expecte source_dtype=torch.bfloat16, trust_remote_code=False, experts_implementation=None, + hf_model_cls=AutoModelForCausalLM, device=torch.device("cpu"), hf_device_map_auto=False, ) @@ -488,7 +592,8 @@ def test_get_input_ids_respects_hf_offline(monkeypatch, offline, expected_local_ monkeypatch.delenv("HF_HUB_OFFLINE", raising=False) else: monkeypatch.setenv("HF_HUB_OFFLINE", offline) - tokenizer = SimpleNamespace(encode=lambda *args, **kwargs: [11, 12, 13]) + tokenizer = Mock() + tokenizer.encode.return_value = [11, 12, 13] with patch("nemo_automodel.NeMoAutoTokenizer.from_pretrained", return_value=tokenizer) as from_pretrained: input_ids = _get_input_ids("mistralai/Ministral-3-3B-Instruct-2512") @@ -499,6 +604,7 @@ def test_get_input_ids_respects_hf_offline(monkeypatch, offline, expected_local_ trust_remote_code=True, local_files_only=expected_local_files_only, ) + tokenizer.encode.assert_called_once_with(_get_parity_document(), add_special_tokens=False) @pytest.mark.parametrize(("offline", "expected_local_files_only"), [(None, False), ("1", True)]) @@ -507,7 +613,8 @@ def test_get_vlm_input_ids_uses_processor_tokenizer(monkeypatch, offline, expect monkeypatch.delenv("HF_HUB_OFFLINE", raising=False) else: monkeypatch.setenv("HF_HUB_OFFLINE", offline) - tokenizer = SimpleNamespace(encode=lambda *args, **kwargs: [21, 22, 23]) + tokenizer = Mock() + tokenizer.encode.return_value = [21, 22, 23] processor = SimpleNamespace(tokenizer=tokenizer) with patch("transformers.AutoProcessor.from_pretrained", return_value=processor) as from_pretrained: @@ -519,6 +626,22 @@ def test_get_vlm_input_ids_uses_processor_tokenizer(monkeypatch, offline, expect trust_remote_code=True, local_files_only=expected_local_files_only, ) + tokenizer.encode.assert_called_once_with(_get_parity_document(), add_special_tokens=False) + + +def test_parity_document_is_the_frozen_long_form_finetuning_guide_snapshot(): + document = _get_parity_document() + + assert "Supervised Fine-Tuning (SFT) and Parameter-Efficient Fine-Tuning (PEFT)" in document[:200] + assert "## Configure Your Training Recipe" in document + assert "## Next Steps" in document + assert len(document.split()) > 6_000 + + +@pytest.mark.parametrize("input_ids_loader", [_get_input_ids, _get_vlm_input_ids]) +def test_parity_input_requires_a_model_tokenizer(input_ids_loader): + with pytest.raises(ValueError, match="tokenizer_name is required"): + input_ids_loader(None) def test_load_input_ids_once_shares_rank0_result(tmp_path, monkeypatch): @@ -528,22 +651,37 @@ def test_load_input_ids_once_shares_rank0_result(tmp_path, monkeypatch): monkeypatch.setenv("SLURM_JOB_ID", "input-id-test") monkeypatch.setenv("RANK", "0") - assert _load_input_ids_once(cfg, rank0_loader, "model/tokenizer") == [31, 32, 33] + assert _load_input_ids_once(cfg, rank0_loader, "model/tokenizer", sequence_length=3) == [31, 32, 33] rank0_loader.assert_called_once_with("model/tokenizer") rank1_loader = Mock(side_effect=AssertionError("nonzero rank must not load the tokenizer")) monkeypatch.setenv("RANK", "1") - assert _load_input_ids_once(cfg, rank1_loader, "model/tokenizer") == [31, 32, 33] + assert _load_input_ids_once(cfg, rank1_loader, "model/tokenizer", sequence_length=3) == [31, 32, 33] rank1_loader.assert_not_called() rank0_reuse_loader = Mock(side_effect=AssertionError("rank 0 must reuse the published input IDs")) monkeypatch.setenv("RANK", "0") - assert _load_input_ids_once(cfg, rank0_reuse_loader, "model/tokenizer") == [31, 32, 33] + assert _load_input_ids_once(cfg, rank0_reuse_loader, "model/tokenizer", sequence_length=3) == [31, 32, 33] rank0_reuse_loader.assert_not_called() +def test_load_input_ids_once_rejects_short_document_for_parity_length(tmp_path): + cfg = SimpleNamespace(checkpoint=SimpleNamespace(checkpoint_dir=tmp_path / "checkpoints")) + + with pytest.raises(ValueError, match="contains 3 tokens, but parity_sequence_length requires 8"): + _load_input_ids_once(cfg, Mock(return_value=[7, 8, 9]), None, sequence_length=8) + + +def test_load_input_ids_once_truncates_long_document_to_parity_length(tmp_path): + cfg = SimpleNamespace(checkpoint=SimpleNamespace(checkpoint_dir=tmp_path / "checkpoints")) + + input_ids = _load_input_ids_once(cfg, Mock(return_value=[7, 8, 9, 10, 11]), None, sequence_length=3) + + assert input_ids == [7, 8, 9] + + def test_load_input_ids_once_waits_for_payload_visibility(tmp_path, monkeypatch): cfg = SimpleNamespace(checkpoint=SimpleNamespace(checkpoint_dir=tmp_path / "checkpoints")) monkeypatch.setenv("WORLD_SIZE", "2") @@ -563,7 +701,7 @@ def publish_payload(_seconds): "tests.functional_tests.checkpoint_robustness.test_checkpoint_robustness_llm.time.sleep", side_effect=publish_payload, ): - assert _load_input_ids_once(cfg, loader, "model/tokenizer") == [41, 42, 43] + assert _load_input_ids_once(cfg, loader, "model/tokenizer", sequence_length=3) == [41, 42, 43] loader.assert_not_called() @@ -575,11 +713,16 @@ def test_load_input_ids_once_propagates_rank0_failure(tmp_path, monkeypatch): monkeypatch.setenv("RANK", "0") with pytest.raises(ValueError, match="tokenizer failed"): - _load_input_ids_once(cfg, Mock(side_effect=ValueError("tokenizer failed")), "model/tokenizer") + _load_input_ids_once( + cfg, + Mock(side_effect=ValueError("tokenizer failed")), + "model/tokenizer", + sequence_length=3, + ) monkeypatch.setenv("RANK", "1") with pytest.raises(RuntimeError, match="Rank 0 input-ID loading failed"): - _load_input_ids_once(cfg, Mock(), "model/tokenizer") + _load_input_ids_once(cfg, Mock(), "model/tokenizer", sequence_length=3) def test_vllm_deploy_tokenization_omits_token_type_ids(): @@ -617,6 +760,70 @@ def test_extract_custom_args_accepts_isolated_phase(): assert remaining == ["--other-arg"] +def test_extract_custom_args_enables_core_source_and_resume_checks_by_default(): + custom, remaining = _extract_custom_args(["--other-arg"]) + + assert custom["source_load_parity_enabled"] is True + assert custom["resume_enabled"] is True + assert remaining == ["--other-arg"] + + +def test_extract_custom_args_reads_semantic_skips_and_parity_settings(tmp_path): + config_path = tmp_path / "recipe.yaml" + config_path.write_text( + "ci:\n" + " checkpoint_robustness:\n" + " skip_source_load_parity: true\n" + " skip_source_load_logit_parity: true\n" + " skip_resume: true\n" + " skip_automodel_reload_logit_parity: true\n" + " skip_hf_reload_logit_parity: true\n" + " trust_remote_code: false\n" + " parity_sequence_length: 1024\n" + " parity_tolerance_profile: relaxed\n" + " parity_tolerance_profile_overrides:\n" + " source_load: strict\n" + " hf_reload: standard\n" + " parity_threshold_overrides:\n" + " source_load: {mean_kl: 0.01}\n" + " automodel_reload: {mean_kl: 0.04, cosine_similarity: 0.99}\n" + " hf_reload: {p95_kl: 0.08}\n" + " cross_tp: {cosine_similarity: 0.997}\n" + " hf_reload_timeout_seconds: 3600\n" + ) + + custom, remaining = _extract_custom_args(["--config", str(config_path)]) + + assert custom["source_load_parity_enabled"] is False + assert custom["resume_enabled"] is False + assert custom["skip_source_load_logit_parity"] is True + assert custom["skip_automodel_reload_logit_parity"] is True + assert custom["skip_hf_reload_logit_parity"] is True + assert custom["trust_remote_code"] is False + assert custom["parity_sequence_length"] == "1024" + assert custom["parity_tolerance_profile"] == "relaxed" + assert custom["parity_tolerance_profile_overrides"] == { + "source_load": "strict", + "hf_reload": "standard", + } + assert custom["parity_threshold_overrides"] == { + "source_load": {"mean_kl": 0.01}, + "automodel_reload": {"mean_kl": 0.04, "cosine_similarity": 0.99}, + "hf_reload": {"p95_kl": 0.08}, + "cross_tp": {"cosine_similarity": 0.997}, + } + assert custom["hf_reload_timeout_seconds"] == "3600" + assert remaining == ["--config", str(config_path)] + + +def test_extract_custom_args_rejects_removed_config_fields(tmp_path): + config_path = tmp_path / "recipe.yaml" + config_path.write_text("ci:\n checkpoint_robustness:\n hf_kl_threshold: 0.01\n") + + with pytest.raises(ValueError, match="Removed checkpoint-robustness fields.*hf_kl_threshold"): + _extract_custom_args(["--config", str(config_path)]) + + def test_extract_custom_args_accepts_resume_tolerance_profile_and_numeric_override(): custom, remaining = _extract_custom_args( ["--resume_tolerance_profile", "relaxed", "--resume_loss_threshold", "0.02", "--other-arg"] @@ -627,17 +834,12 @@ def test_extract_custom_args_accepts_resume_tolerance_profile_and_numeric_overri assert remaining == ["--other-arg"] -def test_extract_custom_args_accepts_skip_hf_logit_parity(): - custom, remaining = _extract_custom_args(["--skip_hf_logit_parity", "--other-arg"]) - - assert custom["skip_hf_logit_parity"] is True - assert remaining == ["--other-arg"] - - -def test_extract_custom_args_accepts_skip_automodel_logit_parity(): - custom, remaining = _extract_custom_args(["--skip_automodel_logit_parity", "--other-arg"]) +def test_extract_custom_args_accepts_cli_profile_overrides(): + custom, remaining = _extract_custom_args( + ["--parity_tolerance_profile_overrides", "{hf_reload: relaxed}", "--other-arg"] + ) - assert custom["skip_automodel_logit_parity"] is True + assert custom["parity_tolerance_profile_overrides"] == {"hf_reload": "relaxed"} assert remaining == ["--other-arg"] @@ -677,7 +879,7 @@ def test_process_isolated_hf_reload_runs_rank0_hf_loader(tmp_path): reference_logits = torch.randn(1, 2, 3) recipe_cls = Mock() hf_model_cls = Mock() - custom_args = {"hf_device_map_auto": True, "no_check_resume": True, "trust_remote_code": True} + custom_args = {"hf_device_map_auto": True, "skip_resume": True, "trust_remote_code": True} with ( patch( @@ -732,22 +934,109 @@ def test_process_isolated_hf_reload_runs_rank0_hf_loader(tmp_path): recipe_cls.assert_not_called() -def test_process_isolated_resume_rejects_no_check_resume(): - with pytest.raises(ValueError, match="conflicts with no_check_resume=true"): +def test_hf_reload_applies_remote_code_compatibility_patches(): + """Phase 3 installs the same remote-code compatibility setup as Phase 0.""" + with patch( + "nemo_automodel._transformers.utils.apply_cache_compatibility_patches", + side_effect=RuntimeError("compatibility sentinel"), + ) as apply_compatibility: + error = _run_vanilla_hf_reload( + SimpleNamespace(), + [], + torch.empty(1, 0, 0), + hf_model_cls=Mock(), + custom_args={}, + ) + + apply_compatibility.assert_called_once_with() + assert error is not None + assert "RuntimeError: compatibility sentinel" in error + + +def test_process_isolated_cross_tp_reload_uses_exported_weights_and_reports_parity(tmp_path): + checkpoint_dir = tmp_path / "checkpoint" + consolidated_dir = checkpoint_dir / "epoch_0_step_5/model/consolidated" + consolidated_dir.mkdir(parents=True) + artifact_dir = checkpoint_dir / ".checkpoint_robustness" + artifact_dir.mkdir() + reference_logits = torch.randn(1, 2, 3) + candidate_logits = reference_logits.clone() + torch.save(reference_logits, artifact_dir / "reference_logits.pt") + cfg = SimpleNamespace( + checkpoint=SimpleNamespace(checkpoint_dir=checkpoint_dir, enabled=True), + model=SimpleNamespace(pretrained_model_name_or_path="source-model"), + distributed=SimpleNamespace(tp_size=1, dp_size=8), + ) + model_part = torch.nn.Linear(2, 2, bias=False) + cross_tp_trainer = SimpleNamespace(model_parts=[model_part], setup=Mock()) + recipe_cls = Mock(return_value=cross_tp_trainer) + custom_args = {"cross_tp_size": "2", "parity_sequence_length": "2"} + + with ( + patch( + "tests.functional_tests.checkpoint_robustness.test_checkpoint_robustness_llm.parse_args_and_load_config", + return_value=cfg, + ), + patch( + "tests.functional_tests.checkpoint_robustness.test_checkpoint_robustness_llm." + "_disable_distributed_atexit_teardown" + ), + patch( + "tests.functional_tests.checkpoint_robustness.test_checkpoint_robustness_llm._load_input_ids_once", + return_value=[11, 12], + ), + patch( + "tests.functional_tests.checkpoint_robustness.test_checkpoint_robustness_llm." + "_prepare_consolidated_hf_cache_once" + ) as prepare_cache, + patch( + "tests.functional_tests.checkpoint_robustness.test_checkpoint_robustness_llm._get_logits", + return_value=candidate_logits, + ), + patch( + "tests.functional_tests.checkpoint_robustness.test_checkpoint_robustness_llm._compare_logits", + return_value=None, + ) as compare_logits, + patch( + "tests.functional_tests.checkpoint_robustness.test_checkpoint_robustness_llm._raise_distributed_failure" + ) as raise_distributed_failure, + patch("tests.functional_tests.checkpoint_robustness.test_checkpoint_robustness_llm._release_recipe_memory"), + ): + _run_process_isolated_checkpoint_phase( + "cross_tp_reload", + custom_args=custom_args, + recipe_cls=recipe_cls, + hf_model_cls=Mock(), + input_ids_loader=Mock(), + ) + + prepare_cache.assert_called_once_with(cfg, consolidated_dir) + assert cfg.model.pretrained_model_name_or_path == str(consolidated_dir) + assert cfg.checkpoint.enabled is False + assert cfg.distributed.tp_size == 2 + assert cfg.distributed.dp_size is None + recipe_cls.assert_called_once_with(cfg) + cross_tp_trainer.setup.assert_called_once_with() + compare_args = compare_logits.call_args.args + assert compare_args[0] == artifact_dir + torch.testing.assert_close(compare_args[1], reference_logits) + torch.testing.assert_close(compare_args[2], candidate_logits) + assert compare_args[3].phase == "phase_5" + assert compare_args[3].comparison == "cross_tp_reload" + raise_distributed_failure.assert_called_once_with(None) + + +def test_process_isolated_resume_rejects_skip_resume(): + with pytest.raises(ValueError, match="conflicts with skip_resume=true"): _run_process_isolated_checkpoint_phase( "resume", - custom_args={"no_check_resume": True}, + custom_args={"skip_resume": True}, recipe_cls=Mock(), hf_model_cls=Mock(), input_ids_loader=Mock(), ) -@pytest.mark.parametrize("non_finite_kl", [float("nan"), float("inf"), float("-inf")]) -def test_hf_reload_rejects_non_finite_kl(non_finite_kl): - assert "non-finite KL divergence" in _hf_reload_kl_error(non_finite_kl, 7e-2) - - def test_process_isolated_source_load_reference_persists_hf_artifacts(tmp_path): cfg = SimpleNamespace(checkpoint=SimpleNamespace(checkpoint_dir=tmp_path)) reference_logits = torch.randn(1, 2, 3) @@ -755,7 +1044,7 @@ def test_process_isolated_source_load_reference_persists_hf_artifacts(tmp_path): recipe_cls = Mock() hf_model_cls = Mock() custom_args = { - "check_source_load_parity": True, + "source_load_parity_enabled": True, "hf_device_map_auto": True, "trust_remote_code": True, } @@ -795,6 +1084,7 @@ def test_process_isolated_source_load_reference_persists_hf_artifacts(tmp_path): experts_implementation=None, hf_device_map_auto=True, hf_source_post_load_dequantize=False, + parity_tolerance_profile="standard", ) persisted_logits = torch.load( tmp_path / ".checkpoint_robustness" / "source_load_reference_logits.pt", @@ -831,6 +1121,38 @@ def publish_artifacts(_seconds): assert sleep_calls == 2 +def test_prepare_consolidated_hf_cache_once_serializes_preinit_workers(tmp_path, monkeypatch): + consolidated_dir = tmp_path / "checkpoint/model/consolidated" + consolidated_dir.mkdir(parents=True) + cfg = SimpleNamespace(checkpoint=SimpleNamespace(checkpoint_dir=tmp_path / "checkpoint")) + monkeypatch.delenv("SLURM_NTASKS", raising=False) + monkeypatch.delenv("SLURM_PROCID", raising=False) + monkeypatch.setenv("WORLD_SIZE", "2") + monkeypatch.setenv("RANK", "0") + + with ( + patch( + "tests.functional_tests.checkpoint_robustness.test_checkpoint_robustness_llm." + "_prepopulate_hf_dynamic_modules_cache" + ) as prepopulate, + patch("transformers.AutoConfig.from_pretrained") as auto_config, + ): + _prepare_consolidated_hf_cache_once(cfg, consolidated_dir) + + prepopulate.assert_called_once_with(consolidated_dir) + auto_config.assert_called_once_with(str(consolidated_dir), trust_remote_code=True) + + monkeypatch.setenv("RANK", "1") + with patch( + "tests.functional_tests.checkpoint_robustness.test_checkpoint_robustness_llm." + "_prepopulate_hf_dynamic_modules_cache", + side_effect=AssertionError("nonzero rank must reuse the completed cache marker"), + ) as nonzero_prepopulate: + _prepare_consolidated_hf_cache_once(cfg, consolidated_dir) + + nonzero_prepopulate.assert_not_called() + + def test_process_isolated_source_load_parity_compares_persisted_reference(tmp_path): artifact_dir = tmp_path / ".checkpoint_robustness" artifact_dir.mkdir() @@ -845,10 +1167,8 @@ def test_process_isolated_source_load_parity_compares_persisted_reference(tmp_pa source_trainer = SimpleNamespace(model_parts=[model_part], setup=Mock()) recipe_cls = Mock(return_value=source_trainer) custom_args = { - "check_source_load_parity": True, - "source_load_kl_threshold": "4e-2", - "source_load_mean_kl_threshold": "1e-2", - "source_load_cosine_threshold": "0.9985", + "source_load_parity_enabled": True, + "parity_tolerance_profile": "standard", } with ( @@ -899,9 +1219,8 @@ def test_process_isolated_source_load_parity_compares_persisted_reference(tmp_pa assert compare_args.args[0][1:] == (False, False) assert compare_args.args[1:] == (candidate_logits, False) assert compare_args.kwargs == { - "source_load_kl_threshold": 4e-2, - "source_load_mean_kl_threshold": 1e-2, - "source_load_cosine_threshold": 0.9985, + "artifact_dir": artifact_dir, + "policy": _source_load_parity_policy(custom_args), } cleanup_source_load.assert_called_once_with(cfg) raise_distributed_failure.assert_called_once_with(None) @@ -973,7 +1292,13 @@ class TinyAutoModel: "nemo_automodel._transformers.model_init._resolve_custom_model_cls_for_config", return_value=TinyAutoModel, ): - assert _hf_fp32_module_names(hf_config) == ("rotary_emb", "router.e_score_correction_bias") + # Dotted strict names also register their distinctive leaf so vanilla + # layouts with a different parent path stay covered. + assert _hf_fp32_module_names(hf_config) == ( + "rotary_emb", + "router.e_score_correction_bias", + "e_score_correction_bias", + ) def test_hf_fp32_module_names_combines_gdn_and_generic_contracts_without_duplicates(): @@ -996,7 +1321,7 @@ def test_hf_fp32_module_names_is_empty_without_model_contract(): assert _hf_fp32_module_names(SimpleNamespace(architectures=["LlamaForCausalLM"])) == () -def test_source_load_parity_failure_is_returned_for_later_reporting(): +def test_source_load_parity_failure_is_returned_for_later_reporting(tmp_path): reference_logits = torch.tensor([[[2.0, -2.0], [1.0, -1.0]]]) candidate_logits = -reference_logits @@ -1004,30 +1329,153 @@ def test_source_load_parity_failure_is_returned_for_later_reporting(): (reference_logits, None, None), candidate_logits, None, - source_load_kl_threshold=0.0, - source_load_mean_kl_threshold=0.0, - source_load_cosine_threshold=1.0, + artifact_dir=tmp_path, + policy=_source_load_parity_policy({"parity_tolerance_profile": "strict"}), ) assert failure is not None - assert "KL divergence between original HF source load and constructed trainer model too large" in failure + assert "source_load parity failed" in failure -def test_source_load_parity_success_returns_no_deferred_failure(): +def test_source_load_parity_success_returns_no_deferred_failure(tmp_path): logits = torch.tensor([[[2.0, -2.0], [1.0, -1.0]]]) failure = _compare_source_load_parity( (logits, None, None), logits.clone(), None, - source_load_kl_threshold=0.0, - source_load_mean_kl_threshold=0.0, - source_load_cosine_threshold=1.0, + artifact_dir=tmp_path, + policy=_source_load_parity_policy({"parity_tolerance_profile": "strict"}), ) assert failure is None +def test_source_load_logit_skip_keeps_metrics_informational(): + policy = _source_load_parity_policy({"skip_source_load_logit_parity": True}) + + assert policy.enforce is False + + +def test_repeatability_policy_is_same_implementation_and_informational(): + policy = _repeatability_policy( + phase="phase_2", + comparison="automodel_reload_self_repeat", + profile="relaxed", + ) + + assert policy.comparison_kind == "same_implementation" + assert policy.profile == "relaxed" + assert policy.enforce is False + + +def test_parity_policies_use_structured_per_comparison_profile_and_threshold_overrides(): + custom_args = { + "parity_tolerance_profile": "standard", + "parity_tolerance_profile_overrides": { + "source_load": "strict", + "hf_reload": "relaxed", + "cross_tp": "relaxed", + }, + "parity_threshold_overrides": { + "source_load": {"mean_kl": 0.01}, + "automodel_reload": {"mean_kl": 0.04, "cosine_similarity": 0.99}, + "hf_reload": {"p95_kl": 0.08}, + "cross_tp": {"cosine_similarity": 0.997}, + }, + } + + source = _source_load_parity_policy(custom_args) + automodel = _automodel_reload_parity_policy(custom_args) + hf = _hf_reload_parity_policy(custom_args) + cross_tp = _cross_tp_parity_policy(custom_args) + + assert source.profile == "strict" + assert source.mean_kl_threshold_override == 0.01 + assert automodel.profile == "standard" + assert automodel.mean_kl_threshold_override == 0.04 + assert automodel.p95_kl_threshold_override is None + assert automodel.cosine_threshold_override == 0.99 + assert hf.profile == "relaxed" + assert hf.p95_kl_threshold_override == 0.08 + assert cross_tp.profile == "relaxed" + assert cross_tp.cosine_threshold_override == 0.997 + + +def test_compare_logits_persists_machine_readable_metrics(tmp_path): + logits = torch.tensor([[[2.0, -2.0], [1.0, -1.0]]]) + policy = _LogitParityPolicy( + phase="phase_2", + comparison="automodel_model_reload", + comparison_kind="same_implementation", + profile="standard", + ) + + failure = _compare_logits(tmp_path, logits, logits.clone(), policy) + + assert failure is None + payload = json.loads((tmp_path / "parity_metrics/phase_2_automodel_model_reload.json").read_text()) + assert payload["schema_version"] == 1 + assert payload["parity_document_sha256"] == _PARITY_DOCUMENT_SHA256 + assert payload["threshold_mode"] == "profile" + assert payload["passed"] is True + assert payload["within_active_thresholds"] is True + assert payload["would_pass_profile"] is True + assert payload["reference_logits"] == {"dtype": "torch.float32", "shape": [1, 2, 2]} + assert payload["candidate_logits"] == {"dtype": "torch.float32", "shape": [1, 2, 2]} + assert payload["metrics"]["token_count"] == 2 + assert payload["metrics"]["mean_kl"] == pytest.approx(0.0, abs=1e-8) + + +def test_compare_logits_marks_skipped_gate_as_informational(tmp_path): + reference_logits = torch.tensor([[[20.0, -20.0]]]) + candidate_logits = -reference_logits + policy = _LogitParityPolicy( + phase="phase_3", + comparison="hf_export_reload", + comparison_kind="cross_framework", + profile="strict", + enforce=False, + ) + + failure = _compare_logits(tmp_path, reference_logits, candidate_logits, policy) + + assert failure is None + payload = json.loads((tmp_path / "parity_metrics/phase_3_hf_export_reload.json").read_text()) + assert payload["enforced"] is False + assert payload["passed"] is True + assert payload["within_active_thresholds"] is False + assert payload["failures"] == [] + assert payload["threshold_failures"] + + +def test_compare_logits_reports_and_applies_targeted_profile_threshold_overrides(tmp_path): + reference_logits = torch.tensor([[[20.0, -20.0]]]) + candidate_logits = -reference_logits + policy = _LogitParityPolicy( + phase="phase_2", + comparison="automodel_model_reload", + comparison_kind="same_implementation", + profile="strict", + mean_kl_threshold_override=100.0, + p95_kl_threshold_override=100.0, + cosine_threshold_override=-1.0, + ) + + failure = _compare_logits(tmp_path, reference_logits, candidate_logits, policy) + + assert failure is None + payload = json.loads((tmp_path / "parity_metrics/phase_2_automodel_model_reload.json").read_text()) + assert payload["threshold_mode"] == "profile_with_numeric_overrides" + assert payload["profile_failures"] + assert payload["threshold_failures"] == [] + assert payload["threshold_overrides"] == { + "mean_kl": 100.0, + "p95_kl": 100.0, + "cosine_similarity": -1.0, + } + + def test_dequantize_hf_fp8_weights_in_place_handles_linear_and_expert_parameters(): class FakeFP8Module(torch.nn.Module): def __init__(self): @@ -1181,19 +1629,28 @@ def test_hf_reload_wait_has_separate_timeout(tmp_path, monkeypatch): _wait_for_hf_reload_rank0(tmp_path / "done") +def test_hf_reload_wait_accepts_explicit_timeout(tmp_path): + with pytest.raises(TimeoutError, match="Timed out waiting 0s"): + _wait_for_hf_reload_rank0(tmp_path / "done", timeout_s=0) + + def test_hf_reload_finish_returns_error_without_distributed_sync(): assert _finish_hf_reload_sync(None, "HF parity failed") == "HF parity failed" -def test_biencoder_robustness_reads_hf_reload_settings_from_config(tmp_path): +def test_biencoder_robustness_reads_current_settings_from_config(tmp_path): config_path = tmp_path / "recipe.yaml" config_path.write_text( "ci:\n" " checkpoint_robustness:\n" - " check_hf_reload: true\n" - " check_resume: true\n" - " cosine_threshold: 0.998\n" - " hf_cosine_threshold: 0.997\n" + " skip_hf_reload: true\n" + " skip_resume: true\n" + " parity_tolerance_profile: relaxed\n" + " parity_tolerance_profile_overrides:\n" + " hf_reload: standard\n" + " parity_threshold_overrides:\n" + " automodel_reload: {cosine_similarity: 0.997}\n" + " hf_reload: {cosine_similarity: 0.996}\n" " resume_tolerance_profile: relaxed\n" " dataloader.num_workers: 0\n" ) @@ -1201,15 +1658,56 @@ def test_biencoder_robustness_reads_hf_reload_settings_from_config(tmp_path): custom, remaining = _extract_biencoder_custom_args(["--config", str(config_path)]) assert custom == { - "check_hf_reload": True, - "check_resume": True, - "cosine_threshold": "0.998", - "hf_cosine_threshold": "0.997", + "skip_hf_reload": True, + "skip_resume": True, + "parity_tolerance_profile": "relaxed", + "parity_tolerance_profile_overrides": {"hf_reload": "standard"}, + "parity_threshold_overrides": { + "automodel_reload": {"cosine_similarity": 0.997}, + "hf_reload": {"cosine_similarity": 0.996}, + }, "resume_tolerance_profile": "relaxed", } assert remaining == ["--config", str(config_path)] +def test_biencoder_robustness_defaults_to_standard_profile_and_default_on_phases(): + custom, remaining = _extract_biencoder_custom_args(["--other-arg"]) + + assert custom.get("parity_tolerance_profile", "standard") == "standard" + assert custom.get("skip_hf_reload", False) is False + assert custom.get("skip_resume", False) is False + assert remaining == ["--other-arg"] + + +def test_biencoder_robustness_rejects_removed_config_fields(tmp_path): + config_path = tmp_path / "recipe.yaml" + config_path.write_text("ci:\n checkpoint_robustness:\n cosine_threshold: 0.999\n") + + with pytest.raises(ValueError, match="Removed retrieval checkpoint-robustness fields.*cosine_threshold"): + _extract_biencoder_custom_args(["--config", str(config_path)]) + + +def test_biencoder_robustness_rejects_non_cosine_threshold_overrides(tmp_path): + config_path = tmp_path / "recipe.yaml" + config_path.write_text( + "ci:\n checkpoint_robustness:\n parity_threshold_overrides:\n hf_reload: {mean_kl: 0.01}\n" + ) + + with pytest.raises(ValueError, match="hf_reload supports only cosine_similarity"): + _extract_biencoder_custom_args(["--config", str(config_path)]) + + +def test_biencoder_robustness_rejects_unsupported_profile_comparison(tmp_path): + config_path = tmp_path / "recipe.yaml" + config_path.write_text( + "ci:\n checkpoint_robustness:\n parity_tolerance_profile_overrides:\n source_load: relaxed\n" + ) + + with pytest.raises(ValueError, match="supports only automodel_reload and hf_reload"): + _extract_biencoder_custom_args(["--config", str(config_path)]) + + def test_record_deferred_failure_preserves_all_comparison_failures(): failures = [] @@ -1217,3 +1715,94 @@ def test_record_deferred_failure_preserves_all_comparison_failures(): _record_deferred_failure(failures, "Phase 4 HF reload parity", "HF parity failed") assert failures == ["Phase 4 HF reload parity:\nHF parity failed"] + + +def _legacy_partial_rotary_minimax_config(**overrides): + """Tiny in-tree MiniMax-M2 config built from checkpoint-style legacy fields.""" + from transformers import AutoConfig + + kwargs = dict( + vocab_size=128, + hidden_size=64, + intermediate_size=32, + num_hidden_layers=1, + num_attention_heads=4, + num_key_value_heads=2, + head_dim=32, + rotary_dim=16, + num_local_experts=4, + num_experts_per_tok=2, + max_position_embeddings=128, + ) + kwargs.update(overrides) + return AutoConfig.for_model("minimax_m2", **kwargs) + + +def test_repair_legacy_partial_rotary_derives_factor_from_rotary_dim(): + config = _legacy_partial_rotary_minimax_config() + + factor_missing_before = not config.rope_parameters.get("partial_rotary_factor") + assert _repair_legacy_partial_rotary_config(config) is factor_missing_before + assert config.rope_parameters["partial_rotary_factor"] == pytest.approx(0.5) + # A second pass finds the factor present and must not report a repair. + assert _repair_legacy_partial_rotary_config(config) is False + + +def test_repaired_minimax_m2_config_rotates_only_rotary_dim(): + from transformers import AutoModelForCausalLM as HFAutoModelForCausalLM + + config = _legacy_partial_rotary_minimax_config() + _repair_legacy_partial_rotary_config(config) + + model = HFAutoModelForCausalLM.from_config(config) + # inv_freq carries one frequency per rotated dim pair: rotary_dim // 2, + # not head_dim // 2 (the full-rotation failure mode from AMINT-286). + assert model.model.rotary_emb.inv_freq.shape[0] == config.rotary_dim // 2 + + +def test_repair_legacy_partial_rotary_is_noop_without_legacy_spec(): + from transformers import AutoConfig + + llama = AutoConfig.for_model( + "llama", + vocab_size=128, + hidden_size=64, + intermediate_size=32, + num_hidden_layers=1, + num_attention_heads=4, + num_key_value_heads=2, + ) + rope_before = deepcopy(getattr(llama, "rope_parameters", None)) + assert _repair_legacy_partial_rotary_config(llama) is False + assert getattr(llama, "rope_parameters", None) == rope_before + + full_rotary = _legacy_partial_rotary_minimax_config(rotary_dim=32) + rope_before = deepcopy(full_rotary.rope_parameters) + assert _repair_legacy_partial_rotary_config(full_rotary) is False + assert full_rotary.rope_parameters == rope_before + + + +def test_hf_fp32_module_names_cover_vanilla_layout_differences(): + """The fp32 contract must reach tensors whose vanilla-HF parent path differs. + + AutoModel's strict name is ``mlp.gate.e_score_correction_bias``, but in-tree + MiniMax-M2 keeps the buffer at ``mlp.e_score_correction_bias``; without the + leaf entry the HF reference silently casts the router bias to bf16. + """ + + class TinyAutoModel: + _keep_in_fp32_modules_strict = ["mlp.gate.e_score_correction_bias", "router.weight", "norm.bias"] + + hf_config = SimpleNamespace(architectures=["TinyForCausalLM"]) + with patch( + "nemo_automodel._transformers.model_init._resolve_custom_model_cls_for_config", + return_value=TinyAutoModel, + ): + names = _hf_fp32_module_names(hf_config) + + assert "mlp.gate.e_score_correction_bias" in names + assert "e_score_correction_bias" in names + # Generic leaves would pin every weight/bias fp32 and must not be added. + assert "weight" not in names + assert "bias" not in names diff --git a/tests/unit_tests/ci_tests/test_config_resolver.py b/tests/unit_tests/ci_tests/test_config_resolver.py index fe92fd4d3e..3e9a56c239 100644 --- a/tests/unit_tests/ci_tests/test_config_resolver.py +++ b/tests/unit_tests/ci_tests/test_config_resolver.py @@ -19,6 +19,7 @@ from pathlib import Path import pytest +import yaml as pyyaml from ruamel.yaml import YAML SCRIPTS_DIR = Path(__file__).resolve().parents[3] / "tests" / "ci_tests" / "scripts" @@ -316,20 +317,134 @@ def test_nemotron_flash_peft_robustness_keeps_supported_tp_topology(tmp_path): assert "resume_first_loss_threshold" not in resolved["ci"]["checkpoint_robustness"] -def test_qwen3_moe_lora_robustness_keeps_source_and_checkpoint_gates(tmp_path): - """Qwen MoE LoRA retains both source-load and checkpoint reload coverage.""" - recipe_path = REPO_ROOT / "examples/llm_finetune/qwen/qwen3_moe_30b_lora.yaml" +@pytest.mark.parametrize( + "recipe_path", + [ + "examples/llm_finetune/deepseek_v4/deepseek_v4_flash_hellaswag_lora.yaml", + "examples/llm_finetune/ernie4_5/ernie4_5_21b_a3b_hellaswag.yaml", + "examples/llm_finetune/glm/glm_4.7_flash_te_deepep.yaml", + "examples/llm_finetune/gpt_oss/gpt_oss_20b.yaml", + "examples/llm_finetune/kimi/kimi_linear_48b_a3b_hellaswag.yaml", + "examples/llm_finetune/llama3_1/customizer_llama_3_1_8b_full_sft_tp.yaml", + "examples/llm_finetune/llama3_2/llama3_2_1b_hellaswag.yaml", + "examples/llm_finetune/minimax_m2/minimax_m2.7_hellaswag_lora.yaml", + "examples/llm_finetune/mistral/mistral_7b_hellaswag_fp8.yaml", + "examples/llm_finetune/nemotron/nemotron_nano_v3_hellaswag.yaml", + "examples/llm_finetune/nemotron/nemotron_super_v3_hellaswag.yaml", + "examples/llm_finetune/nemotron_flash/nemotron_flash_1b_squad.yaml", + "examples/llm_finetune/qwen/qwen3_moe_30b_lora.yaml", + "examples/vlm_finetune/gemma4/gemma4_26b_a4b_moe.yaml", + "examples/vlm_finetune/mistral4/mistral4_medpix.yaml", + "examples/vlm_finetune/qwen3/qwen3_vl_moe_30b_te_deepep.yaml", + "examples/vlm_finetune/qwen3_5_moe/qwen3_5_35b.yaml", + "examples/vlm_finetune/qwen3_8/qwen3_8_27b.yaml", + "examples/vlm_finetune/stepfun/step3p7_medpix_200b_lora_pp8ep8_8node.yaml", + ], +) +def test_calibration_recipes_enable_all_checkpoint_gates(tmp_path, recipe_path): + """The calibration cohort runs every phase and documents any informational logit gate.""" + recipe_path = REPO_ROOT / recipe_path out = tmp_path / "resolved.yaml" - env = {"PIPELINE_DIR": str(tmp_path), "TEST_NAME": recipe_path.stem} + env = {"PIPELINE_DIR": str(tmp_path), "TEST_NAME": recipe_path.stem, "NEMO_CI_PATH": "/mnt/nci"} + _run_resolver( + ["--base", str(recipe_path), "--phase", "checkpoint_robustness", "--output", str(out)], + env=env, + ) + + # The checkpoint harness consumes the resolver output with PyYAML, whose YAML 1.1 scalar rules differ from ruamel. + robustness = pyyaml.safe_load(out.read_text())["ci"]["checkpoint_robustness"] + assert "check_source_load_parity" not in robustness + assert "skip_source_load_parity" not in robustness + if recipe_path.stem == "step3p7_medpix_200b_lora_pp8ep8_8node": + assert robustness["parity_tolerance_profile"] == "relaxed" + assert robustness["parity_threshold_overrides"] == { + "automodel_reload": {"mean_kl": 4e-2, "cosine_similarity": 0.99} + } + assert robustness["resume_tolerance_profile"] == "relaxed" + assert "skip_automodel_reload_logit_parity" not in robustness + informational_source_and_hf = { + "gemma4_26b_a4b_moe", + "step3p7_medpix_200b_lora_pp8ep8_8node", + } + if recipe_path.stem in informational_source_and_hf: + assert robustness["skip_source_load_logit_parity"] is True + assert robustness["skip_hf_reload_logit_parity"] is True + else: + assert "skip_source_load_logit_parity" not in robustness + assert "skip_hf_reload_logit_parity" not in robustness + assert "skip_hf_reload" not in robustness + assert "skip_resume" not in robustness + + +@pytest.mark.parametrize( + "recipe_name", + [ + "customizer_gpt_oss_full_sft.yaml", + "customizer_gpt_oss_full_sft_chat.yaml", + "customizer_gpt_oss_peft.yaml", + "customizer_gpt_oss_peft_packing.yaml", + ], +) +def test_gpt_oss_customizers_use_routed_moe_parity_profile(tmp_path, recipe_name): + """GPT-OSS Customizer variants retain parity gates with the routed-MoE profile.""" + recipe_path = REPO_ROOT / "examples/llm_finetune/gpt_oss" / recipe_name + out = tmp_path / "resolved.yaml" + env = {"PIPELINE_DIR": str(tmp_path), "TEST_NAME": recipe_path.stem, "NEMO_CI_PATH": "/mnt/nci"} + _run_resolver( + ["--base", str(recipe_path), "--phase", "checkpoint_robustness", "--output", str(out)], + env=env, + ) + + robustness = yaml.load(out.open())["ci"]["checkpoint_robustness"] + assert robustness["parity_tolerance_profile"] == "relaxed" + assert "skip_source_load_logit_parity" not in robustness + assert "skip_hf_reload_logit_parity" not in robustness + assert "skip_resume" not in robustness + + +@pytest.mark.parametrize( + ("recipe_path", "expected_resume_profile"), + [ + ("examples/llm_finetune/nemotron/customizer_nemotron_nano_full_sft.yaml", None), + ("examples/llm_finetune/nemotron/customizer_nemotron_nano_full_sft_chat.yaml", "relaxed"), + ("examples/llm_finetune/qwen/qwen3_moe_30b_hellaswag.yaml", None), + ], +) +def test_additional_routed_moe_configs_enable_resume(tmp_path, recipe_path, expected_resume_profile): + """Expanded routed-MoE coverage keeps resume active with narrowly calibrated loss drift.""" + recipe_path = REPO_ROOT / recipe_path + out = tmp_path / "resolved.yaml" + env = {"PIPELINE_DIR": str(tmp_path), "TEST_NAME": recipe_path.stem, "NEMO_CI_PATH": "/mnt/nci"} _run_resolver( ["--base", str(recipe_path), "--phase", "checkpoint_robustness", "--output", str(out)], env=env, ) robustness = yaml.load(out.open())["ci"]["checkpoint_robustness"] - assert robustness["check_source_load_parity"] is True - assert "skip_automodel_logit_parity" not in robustness - assert robustness["skip_hf_logit_parity"] is True + assert "skip_resume" not in robustness + if expected_resume_profile is None: + assert "resume_tolerance_profile" not in robustness + else: + assert robustness["resume_tolerance_profile"] == expected_resume_profile + assert "parity_tolerance_profile" not in robustness + assert "resume_first_loss_threshold" not in robustness + assert "resume_loss_threshold" not in robustness + + +def test_nemotron_nano_4b_peft_uses_cached_family_tokenizer(tmp_path): + """The offline recipe uses one complete family tokenizer for training and parity.""" + recipe_path = REPO_ROOT / "examples/llm_finetune/nemotron/nemotron_nano_4b_squad_peft.yaml" + out = tmp_path / "resolved.yaml" + env = {"PIPELINE_DIR": str(tmp_path), "TEST_NAME": recipe_path.stem, "NEMO_CI_PATH": "/mnt/nci"} + _run_resolver( + ["--base", str(recipe_path), "--phase", "checkpoint_robustness", "--output", str(out)], + env=env, + ) + + resolved = yaml.load(out.open()) + tokenizer_name = "nvidia/NVIDIA-Nemotron-3-Nano-30B-A3B-BF16" + assert resolved["dataset"]["tokenizer"]["pretrained_model_name_or_path"] == tokenizer_name + assert resolved["ci"]["checkpoint_robustness"]["tokenizer_name"] == tokenizer_name def test_end_to_end_fixture_keys_not_applied_as_overrides(tmp_path): @@ -339,17 +454,23 @@ def test_end_to_end_fixture_keys_not_applied_as_overrides(tmp_path): "step_scheduler: {global_batch_size: 8}\n" "ci:\n" " checkpoint_robustness:\n" - " check_source_load_parity: true # fixture arg, must NOT become top-level\n" - " skip_automodel_logit_parity: true # fixture arg, must NOT become top-level\n" - " skip_hf_logit_parity: true # fixture arg, must NOT become top-level\n" + " skip_source_load_parity: true # fixture arg, must NOT become top-level\n" + " skip_source_load_logit_parity: true # fixture arg, must NOT become top-level\n" + " skip_automodel_reload_logit_parity: true # fixture arg, must NOT become top-level\n" + " skip_hf_reload_logit_parity: true # fixture arg, must NOT become top-level\n" " hf_adapter_ignored_key_prefix: base_model.model.mtp. # fixture arg, must NOT become top-level\n" - " hf_kl_threshold: 5e-3 # fixture arg, must NOT become top-level\n" + " parity_threshold_overrides: # fixture arg, must NOT become top-level\n" + " source_load: {mean_kl: 1e-2}\n" + " automodel_reload: {p95_kl: 2e-2, cosine_similarity: 0.998}\n" + " hf_reload: {mean_kl: 3e-2}\n" + " cross_tp: {cosine_similarity: 0.997}\n" " training_reproducibility_loss_threshold: 1e-2 # fixture arg, must NOT become top-level\n" " resume_tolerance_profile: relaxed # fixture arg, must NOT become top-level\n" " resume_first_loss_threshold: 1e-6 # fixture arg, must NOT become top-level\n" - " source_load_kl_threshold: 1e-2 # fixture arg, must NOT become top-level\n" - " source_load_mean_kl_threshold: 1e-3 # fixture arg, must NOT become top-level\n" - " source_load_cosine_threshold: 0.999 # fixture arg, must NOT become top-level\n" + " parity_sequence_length: 1024 # fixture arg, must NOT become top-level\n" + " parity_tolerance_profile: strict # fixture arg, must NOT become top-level\n" + " parity_tolerance_profile_overrides: # fixture arg, must NOT become top-level\n" + " hf_reload: relaxed\n" " tokenizer_name: nvidia/Test # fixture arg, must NOT become top-level\n" " dataset.limit_dataset_samples: 500 # dotted -> applied as override\n" ) @@ -362,29 +483,36 @@ def test_end_to_end_fixture_keys_not_applied_as_overrides(tmp_path): assert resolved["dataset"]["limit_dataset_samples"] == 500 # Fixture args stay under ci.checkpoint_robustness for the consumer (pytest) to read, # and do NOT pollute the top level. - assert "hf_kl_threshold" not in resolved - assert "check_source_load_parity" not in resolved - assert "skip_automodel_logit_parity" not in resolved - assert "skip_hf_logit_parity" not in resolved + assert "parity_threshold_overrides" not in resolved + assert "skip_source_load_parity" not in resolved + assert "skip_source_load_logit_parity" not in resolved + assert "skip_automodel_reload_logit_parity" not in resolved + assert "skip_hf_reload_logit_parity" not in resolved assert "hf_adapter_ignored_key_prefix" not in resolved assert "training_reproducibility_loss_threshold" not in resolved assert "resume_tolerance_profile" not in resolved assert "resume_first_loss_threshold" not in resolved - assert "source_load_kl_threshold" not in resolved - assert "source_load_mean_kl_threshold" not in resolved - assert "source_load_cosine_threshold" not in resolved + assert "parity_sequence_length" not in resolved + assert "parity_tolerance_profile" not in resolved + assert "parity_tolerance_profile_overrides" not in resolved assert "tokenizer_name" not in resolved - assert resolved["ci"]["checkpoint_robustness"]["hf_kl_threshold"] == 5e-3 - assert resolved["ci"]["checkpoint_robustness"]["check_source_load_parity"] is True - assert resolved["ci"]["checkpoint_robustness"]["skip_automodel_logit_parity"] is True - assert resolved["ci"]["checkpoint_robustness"]["skip_hf_logit_parity"] is True + assert resolved["ci"]["checkpoint_robustness"]["parity_threshold_overrides"] == { + "source_load": {"mean_kl": 1e-2}, + "automodel_reload": {"p95_kl": 2e-2, "cosine_similarity": 0.998}, + "hf_reload": {"mean_kl": 3e-2}, + "cross_tp": {"cosine_similarity": 0.997}, + } + assert resolved["ci"]["checkpoint_robustness"]["skip_source_load_parity"] is True + assert resolved["ci"]["checkpoint_robustness"]["skip_source_load_logit_parity"] is True + assert resolved["ci"]["checkpoint_robustness"]["skip_automodel_reload_logit_parity"] is True + assert resolved["ci"]["checkpoint_robustness"]["skip_hf_reload_logit_parity"] is True assert resolved["ci"]["checkpoint_robustness"]["hf_adapter_ignored_key_prefix"] == "base_model.model.mtp." assert resolved["ci"]["checkpoint_robustness"]["training_reproducibility_loss_threshold"] == 1e-2 assert resolved["ci"]["checkpoint_robustness"]["resume_tolerance_profile"] == "relaxed" assert resolved["ci"]["checkpoint_robustness"]["resume_first_loss_threshold"] == 1e-6 - assert resolved["ci"]["checkpoint_robustness"]["source_load_kl_threshold"] == 1e-2 - assert resolved["ci"]["checkpoint_robustness"]["source_load_mean_kl_threshold"] == 1e-3 - assert resolved["ci"]["checkpoint_robustness"]["source_load_cosine_threshold"] == 0.999 + assert resolved["ci"]["checkpoint_robustness"]["parity_sequence_length"] == 1024 + assert resolved["ci"]["checkpoint_robustness"]["parity_tolerance_profile"] == "strict" + assert resolved["ci"]["checkpoint_robustness"]["parity_tolerance_profile_overrides"] == {"hf_reload": "relaxed"} @pytest.mark.parametrize( @@ -412,12 +540,8 @@ def test_vlm_checkpoint_robustness_recipes_resolve(tmp_path, recipe_path): assert resolved["checkpoint"]["enabled"] is True assert resolved["checkpoint"]["model_save_format"] == "safetensors" assert resolved["checkpoint"]["save_consolidated"] is True - if Path(recipe_path).stem == "gemma4_26b_a4b_moe": - # Opted out: source-load logit KL tracks the host's reduction order for - # this recipe's DeepEP MoE routing, not checkpoint integrity. - assert robustness["check_source_load_parity"] is False - else: - assert robustness["check_source_load_parity"] is True + assert "skip_source_load_parity" not in robustness + assert "check_source_load_parity" not in robustness assert robustness["tokenizer_name"] == resolved["model"]["pretrained_model_name_or_path"] if Path(recipe_path).stem == "gemma4_26b_a4b_moe": assert resolved["distributed"]["multimodal"]["frozen_sharding"] == "replicate" @@ -428,28 +552,37 @@ def test_vlm_checkpoint_robustness_recipes_resolve(tmp_path, recipe_path): assert robustness["hf_device_map_auto"] is True if "/mistral4/" in recipe_path: assert robustness["hf_source_post_load_dequantize"] is True - assert robustness["kl_threshold"] == 5e-2 - assert robustness["source_load_kl_threshold"] == 1e-2 - assert robustness["source_load_mean_kl_threshold"] == 2e-3 - assert robustness["source_load_cosine_threshold"] == 0.999 - assert robustness["hf_kl_threshold"] == 5e-2 + assert "parity_tolerance_profile" not in robustness + assert robustness["parity_tolerance_profile_overrides"] == {"hf_reload": "relaxed"} + for key in ( + "kl_threshold", + "source_load_kl_threshold", + "source_load_mean_kl_threshold", + "source_load_cosine_threshold", + "hf_kl_threshold", + ): + assert key not in robustness if Path(recipe_path).stem == "qwen3_vl_moe_30b_te_deepep": - assert robustness["hf_kl_threshold"] == 2.5e-2 assert "resume_loss_threshold" not in robustness assert robustness["training_reproducibility_loss_threshold"] == 2e-2 - assert robustness["source_load_kl_threshold"] == 4e-2 - assert robustness["source_load_mean_kl_threshold"] == 7e-3 + for key in ( + "hf_kl_threshold", + "source_load_kl_threshold", + "source_load_mean_kl_threshold", + "source_load_cosine_threshold", + ): + assert key not in robustness if Path(recipe_path).stem == "qwen3_5_35b": assert robustness["experts_implementation"] == "grouped_mm" for key in ( "hf_keep_in_fp32_modules", "resume_loss_threshold", + "hf_kl_threshold", + "source_load_cosine_threshold", + "source_load_kl_threshold", + "source_load_mean_kl_threshold", ): assert key not in robustness - assert robustness["hf_kl_threshold"] == 1e-1 - assert robustness["source_load_cosine_threshold"] == 0.9985 - assert robustness["source_load_kl_threshold"] == 1e-1 - assert robustness["source_load_mean_kl_threshold"] == 1e-2 assert resolved["loss_fn"]["_target_"] == ("nemo_automodel.components.loss.chunked_ce.ChunkedCrossEntropy") assert resolved["model"]["backend"]["experts"] == "torch_mm" assert resolved["step_scheduler"]["global_batch_size"] == 16 @@ -473,6 +606,9 @@ def test_retrieval_checkpoint_robustness_retains_calibrated_resume_threshold(tmp ) robustness = yaml.load(out.open())["ci"]["checkpoint_robustness"] + assert robustness["parity_tolerance_profile"] == "standard" + for removed_key in ("check_hf_reload", "check_resume", "cosine_threshold", "hf_cosine_threshold"): + assert removed_key not in robustness assert robustness["resume_loss_threshold"] == 5e-2 assert robustness["training_reproducibility_loss_threshold"] == 5e-2 diff --git a/tests/unit_tests/ci_tests/test_generate_ci_tests.py b/tests/unit_tests/ci_tests/test_generate_ci_tests.py index 29ca83cb09..2045c270cb 100644 --- a/tests/unit_tests/ci_tests/test_generate_ci_tests.py +++ b/tests/unit_tests/ci_tests/test_generate_ci_tests.py @@ -19,6 +19,38 @@ from tests.ci_tests.utils.generate_ci_tests import generate_job, generate_pipeline +def test_example_checkpoint_robustness_configs_do_not_use_removed_fields(): + removed_keys = { + "automodel_reload_cosine_threshold", + "automodel_reload_mean_kl_threshold", + "automodel_reload_p95_kl_threshold", + "check_hf_reload", + "check_resume", + "check_source_load_parity", + "cosine_threshold", + "hf_cosine_threshold", + "kl_threshold", + "hf_kl_threshold", + "source_load_kl_threshold", + "source_load_mean_kl_threshold", + "source_load_cosine_threshold", + "cross_tp_kl_threshold", + "no_check_resume", + "skip_automodel_logit_parity", + "skip_hf_logit_parity", + } + violations = [] + + for config in Path("examples").rglob("*.yaml"): + recipe = YAML(typ="safe").load(config) or {} + robustness = (recipe.get("ci") or {}).get("checkpoint_robustness") or {} + found = sorted(removed_keys & robustness.keys()) + if found: + violations.append(f"{config}: {', '.join(found)}") + + assert not violations, "Removed checkpoint-robustness fields remain:\n" + "\n".join(violations) + + def test_generate_deepseek_v4_pretrain_nightly_job(): pipeline = generate_pipeline(".", "nightly", "llm_pretrain") @@ -91,8 +123,7 @@ def test_generate_checkpoint_robustness_process_isolation_derives_phases(tmp_pat ci: checkpoint_robustness: process_isolation: true - check_source_load_parity: true - no_check_resume: true + skip_resume: true trust_remote_code: true hf_device_map_auto: true hf_device_map_max_memory_gib: 55 @@ -114,15 +145,16 @@ def test_generate_checkpoint_robustness_process_isolation_preserves_full_default (tmp_path / config).write_text( """ ci: - checkpoint_robustness: - process_isolation: true + checkpoint_robustness: {} """, encoding="utf-8", ) jobs = dict(generate_job(config, {}, "release", "llm_finetune", str(tmp_path))) - assert jobs[""]["variables"]["CHECKPOINT_ROBUSTNESS_PHASES"] == ("train_and_save automodel_reload hf_reload resume") + assert jobs[""]["variables"]["CHECKPOINT_ROBUSTNESS_PHASES"] == ( + "source_load_reference source_load_parity train_and_save automodel_reload hf_reload resume" + ) def test_generate_checkpoint_robustness_process_isolation_honors_skips(tmp_path): @@ -132,8 +164,9 @@ def test_generate_checkpoint_robustness_process_isolation_honors_skips(tmp_path) ci: checkpoint_robustness: process_isolation: true + skip_source_load_parity: true skip_hf_reload: true - no_check_resume: true + skip_resume: true """, encoding="utf-8", ) @@ -143,6 +176,42 @@ def test_generate_checkpoint_robustness_process_isolation_honors_skips(tmp_path) assert jobs[""]["variables"]["CHECKPOINT_ROBUSTNESS_PHASES"] == "train_and_save automodel_reload" +def test_generate_checkpoint_robustness_process_isolation_is_default_and_cross_tp_is_last(tmp_path): + config = Path("dense_model.yaml") + (tmp_path / config).write_text( + """ +ci: + checkpoint_robustness: + cross_tp_size: 2 +""", + encoding="utf-8", + ) + + jobs = dict(generate_job(config, {}, "release", "llm_finetune", str(tmp_path))) + + assert jobs[""]["variables"]["CHECKPOINT_ROBUSTNESS_PROCESS_ISOLATION"] == "true" + assert jobs[""]["variables"]["CHECKPOINT_ROBUSTNESS_PHASES"] == ( + "source_load_reference source_load_parity train_and_save automodel_reload hf_reload resume cross_tp_reload" + ) + + +def test_generate_checkpoint_robustness_allows_single_process_fallback(tmp_path): + config = Path("fallback_model.yaml") + (tmp_path / config).write_text( + """ +ci: + checkpoint_robustness: + process_isolation: false +""", + encoding="utf-8", + ) + + jobs = dict(generate_job(config, {}, "release", "llm_finetune", str(tmp_path))) + + assert "CHECKPOINT_ROBUSTNESS_PROCESS_ISOLATION" not in jobs[""]["variables"] + assert "CHECKPOINT_ROBUSTNESS_PHASES" not in jobs[""]["variables"] + + def test_generate_checkpoint_robustness_process_isolation_allows_phase_override(tmp_path): config = Path("custom_phases.yaml") (tmp_path / config).write_text( @@ -161,7 +230,7 @@ def test_generate_checkpoint_robustness_process_isolation_allows_phase_override( assert jobs[""]["variables"]["CHECKPOINT_ROBUSTNESS_PHASES"] == "train_and_save automodel_reload" -def test_generate_qwen3_moe_lora_uses_isolated_reload_phases(): +def test_generate_qwen3_moe_lora_uses_all_isolated_checkpoint_phases(): config = Path("examples/llm_finetune/qwen/qwen3_moe_30b_lora.yaml") jobs = dict(generate_job(config, {}, "release", "llm_finetune", ".")) @@ -171,15 +240,51 @@ def test_generate_qwen3_moe_lora_uses_isolated_reload_phases(): variables = jobs[""]["variables"] assert "CHECKPOINT_ROBUSTNESS_PHASES" not in ci_config.get("env_vars", {}) assert variables["PYTORCH_CUDA_ALLOC_CONF"] == "expandable_segments:True" + assert variables["TIME"] == "00:30:00" assert variables["CHECKPOINT_ROBUSTNESS_PROCESS_ISOLATION"] == "true" assert variables["CHECKPOINT_ROBUSTNESS_PHASES"] == ( - "source_load_reference source_load_parity train_and_save automodel_reload hf_reload" + "source_load_reference source_load_parity train_and_save automodel_reload hf_reload resume" ) - assert robustness["check_source_load_parity"] is True - assert robustness["skip_hf_logit_parity"] is True - assert robustness["source_load_kl_threshold"] == 3e-2 - assert robustness["source_load_mean_kl_threshold"] == 6e-3 - assert robustness["source_load_cosine_threshold"] == 0.997 + assert "check_source_load_parity" not in robustness + assert "skip_source_load_parity" not in robustness + assert "skip_automodel_reload_logit_parity" not in robustness + assert "skip_hf_reload_logit_parity" not in robustness + assert "skip_resume" not in robustness + for key in ( + "source_load_kl_threshold", + "source_load_mean_kl_threshold", + "source_load_cosine_threshold", + ): + assert key not in robustness + + +def test_generate_nemotron_resume_cohort_preserves_known_issue_gating(): + expected_times = { + "customizer_nemotron_nano_peft": "00:30:00", + "customizer_nemotron_nano_peft_packing": "00:30:00", + "nemotron_nano_4b_squad_peft": "00:30:00", + "nemotron_nano_8b_v1_squad": "00:25:00", + "nemotron_nano_8b_v1_squad_peft": "00:25:00", + "nemotron_nano_9b_squad": "00:30:00", + "nemotron_nano_9b_squad_peft": "00:30:00", + "nemotron_nano_v3_hellaswag_peft": "00:30:00", + } + + for recipe_name, expected_time in expected_times.items(): + config = Path(f"examples/llm_finetune/nemotron/{recipe_name}.yaml") + jobs = dict(generate_job(config, {}, "release", "llm_finetune", ".")) + + job = jobs[""] + assert job.get("allow_failure") is None + assert job["variables"]["TIME"] == expected_time + assert job["variables"]["CHECKPOINT_ROBUSTNESS_PROCESS_ISOLATION"] == "true" + expected_phases = "source_load_reference source_load_parity train_and_save automodel_reload hf_reload resume" + if recipe_name == "nemotron_nano_8b_v1_squad": + expected_phases += " cross_tp_reload" + assert job["variables"]["CHECKPOINT_ROBUSTNESS_PHASES"] == expected_phases + + known_issue_config = Path("examples/llm_finetune/nemotron/nemotron_nano_4b_squad.yaml") + assert generate_job(known_issue_config, {}, "release", "llm_finetune", ".") == [] def test_generate_qwen3_moe_te_deepep_uses_isolated_source_and_reload_phases(): @@ -191,14 +296,16 @@ def test_generate_qwen3_moe_te_deepep_uses_isolated_source_and_reload_phases(): variables = jobs[""]["variables"] assert "CHECKPOINT_ROBUSTNESS_PHASES" not in ci_config.get("env_vars", {}) - assert variables["TIME"] == "00:20:00" + assert variables["TIME"] == "00:30:00" assert variables["CHECKPOINT_ROBUSTNESS_PROCESS_ISOLATION"] == "true" assert variables["CHECKPOINT_ROBUSTNESS_PHASES"] == ( - "source_load_reference source_load_parity train_and_save automodel_reload hf_reload" + "source_load_reference source_load_parity train_and_save automodel_reload hf_reload resume" ) - assert robustness["check_source_load_parity"] is True - assert robustness["source_load_kl_threshold"] == 3e-2 - assert robustness["source_load_mean_kl_threshold"] == 6e-3 - assert robustness["source_load_cosine_threshold"] == 0.997 + assert "check_source_load_parity" not in robustness + assert "skip_source_load_parity" not in robustness + assert "source_load_kl_threshold" not in robustness + assert "source_load_mean_kl_threshold" not in robustness + assert "source_load_cosine_threshold" not in robustness + assert "skip_resume" not in robustness assert robustness["trust_remote_code"] is True assert robustness["hf_device_map_auto"] is True diff --git a/tests/unit_tests/ci_tests/test_resume_trajectory.py b/tests/unit_tests/ci_tests/test_resume_trajectory.py index 7917e2f698..9300d45e8c 100644 --- a/tests/unit_tests/ci_tests/test_resume_trajectory.py +++ b/tests/unit_tests/ci_tests/test_resume_trajectory.py @@ -29,6 +29,7 @@ _compare_training_reproducibility, _configure_resumed_run, _configure_uninterrupted_run, + _disable_checkpoint_saves_after_restore, _report_resume_comparison, _report_training_reproducibility, _resolve_resume_loss_tolerance, @@ -44,7 +45,12 @@ def _config(max_steps: int = 5) -> SimpleNamespace: return SimpleNamespace( - step_scheduler=SimpleNamespace(max_steps=max_steps, ckpt_every_steps=max_steps), + step_scheduler=SimpleNamespace( + max_steps=max_steps, + num_epochs=1, + ckpt_every_steps=max_steps, + save_checkpoint_every_epoch=True, + ), lr_scheduler=SimpleNamespace(lr_decay_steps=None), checkpoint=SimpleNamespace( checkpoint_dir="/tmp/checkpoint-robustness", @@ -100,7 +106,9 @@ def test_shared_resume_plan_extends_phase_one_from_the_checkpoint_boundary(tmp_p assert plan.boundary_step == 5 assert plan.comparison_steps == (5, 6, 7) assert cfg.step_scheduler.max_steps == 8 + assert cfg.step_scheduler.num_epochs == 8 assert cfg.step_scheduler.ckpt_every_steps == 5 + assert cfg.step_scheduler.save_checkpoint_every_epoch is False assert cfg.lr_scheduler.lr_decay_steps == 5 assert cfg.checkpoint.save_consolidated == "final" @@ -109,6 +117,16 @@ def test_shared_resume_plan_extends_phase_one_from_the_checkpoint_boundary(tmp_p assert cfg.checkpoint.restore_from == str(checkpoint_path) assert cfg.checkpoint.checkpoint_dir == str(plan.resume_checkpoint_dir) assert cfg.checkpoint.save_consolidated is False + assert cfg.step_scheduler.num_epochs == 8 + assert cfg.step_scheduler.save_checkpoint_every_epoch is False + + +def test_resume_continuation_disables_checkpoint_writes_after_restore(): + trainer = SimpleNamespace(checkpointer=SimpleNamespace(config=SimpleNamespace(enabled=True))) + + _disable_checkpoint_saves_after_restore(trainer) + + assert trainer.checkpointer.config.enabled is False def test_resume_state_check_detects_omitted_rng_state(): @@ -202,11 +220,13 @@ def make_dataloader() -> StatefulDataLoader: def test_shared_trajectory_harness_runs_checkpoint_and_resume_locally(tmp_path): - plan = _ResumePlan(checkpoint_dir=tmp_path, boundary_step=2, continuation_steps=2) + # Three optimizer steps fit in each epoch, so the step-five checkpoint and + # its continuation exercise dataloader restoration across epoch boundaries. + plan = _ResumePlan(checkpoint_dir=tmp_path, boundary_step=5, continuation_steps=3) checkpoints: dict[int, dict[str, object]] = {} def make_trainer() -> SimpleNamespace: - dataset = TensorDataset(torch.arange(16, dtype=torch.float32)) + dataset = TensorDataset(torch.arange(6, dtype=torch.float32)) sampler = StatefulDistributedSampler( dataset, seed=42, diff --git a/tests/unit_tests/models/minimax_m2/__init__.py b/tests/unit_tests/models/minimax_m2/__init__.py index 26496bfed7..c03552b02d 100644 --- a/tests/unit_tests/models/minimax_m2/__init__.py +++ b/tests/unit_tests/models/minimax_m2/__init__.py @@ -1 +1,13 @@ # Copyright (c) 2026, NVIDIA CORPORATION. All rights reserved. +# +# 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. diff --git a/tests/unit_tests/models/minimax_m2/test_minimax_m2_router_precision.py b/tests/unit_tests/models/minimax_m2/test_minimax_m2_router_precision.py new file mode 100644 index 0000000000..b9159b7d9f --- /dev/null +++ b/tests/unit_tests/models/minimax_m2/test_minimax_m2_router_precision.py @@ -0,0 +1,119 @@ +# Copyright (c) 2026, NVIDIA CORPORATION. All rights reserved. +# +# 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. + +"""MiniMax-M2 router precision contract. + +Released MiniMax-M2 checkpoints store the router gate weight in fp32 and the +HF reference projects with ``hidden_states.to(weight.dtype)``, so the +checkpoint-faithful router is fp32 end to end: fp32 parameter, fp32 +projection, fp32 scoring, fp32 selected weights (AMINT-286). +""" + +import torch +from transformers import AutoConfig + +from nemo_automodel.components.models.common import BackendConfig +from nemo_automodel.components.models.minimax_m2.model import MiniMaxM2ForCausalLM + +TINY = dict( + vocab_size=128, + hidden_size=64, + intermediate_size=32, + num_hidden_layers=2, + num_attention_heads=4, + num_key_value_heads=2, + head_dim=32, + rotary_dim=16, + num_local_experts=8, + num_experts_per_tok=2, + max_position_embeddings=128, +) + + +def _cpu_backend() -> BackendConfig: + return BackendConfig( + linear="torch", + attn="sdpa", + rms_norm="torch", + rope_fusion=False, + dispatcher="torch", + experts="torch", + fake_balanced_gate=False, + enable_hf_state_dict_adapter=False, + ) + + +def test_router_fp32_contract_is_model_owned(): + config = AutoConfig.for_model("minimax_m2", torch_dtype="bfloat16", **TINY) + model = MiniMaxM2ForCausalLM(config, backend=_cpu_backend()).eval() + + assert model.model.backend.gate_precision == torch.float32 + assert model.model.moe_config.router_weights_fp32 is True + assert "mlp.gate.weight" in MiniMaxM2ForCausalLM._keep_in_fp32_modules_strict + assert "mlp.gate.e_score_correction_bias" in MiniMaxM2ForCausalLM._keep_in_fp32_modules_strict + + # After the model-wide bf16 cast, the fp32 contract keeps the router gate + # parameter and bias in fp32 while the rest of the model is bf16. + model.initialize_weights(buffer_device=torch.device("cpu"), dtype=torch.bfloat16) + gate = model.model.layers["0"].mlp.gate + assert gate.weight.dtype == torch.float32 + assert gate.e_score_correction_bias.dtype == torch.float32 + assert model.model.layers["0"].self_attn.q_proj.weight.dtype == torch.bfloat16 + + # Selected routing weights stay fp32 through the gate output, matching the + # HF reference's top_k_weights.to(router_logits.dtype) with an fp32 weight. + # x: Tensor of shape [tokens, hidden] in bf16, like real routed inputs. + x = torch.randn(8, TINY["hidden_size"], dtype=torch.bfloat16) + weights, indices, _aux = gate(x, torch.ones(8, dtype=torch.bool), None) + assert weights.dtype == torch.float32 + assert indices.shape == (8, TINY["num_experts_per_tok"]) + + +def test_gate_is_fp32_at_construction_for_fsdp_dtype_grouping(): + """The gate must be fp32 from allocation, before any init or checkpoint cast. + + FSDP shards the freshly constructed (meta/from_pretrained) module: a + bf16-allocated gate weight with an fp32-pinned compute dtype shares its + module with the fp32 correction-bias buffer, which FSDP cannot isolate + (pipeline 64344786: "FSDP could not isolate parameters with a distinct + dtype from siblings in the same module: mlp.gate.weight"). + """ + import torch.distributed.fsdp as fsdp + + from nemo_automodel.components.distributed.parallelizer_utils import fully_shard_by_dtype + + config = AutoConfig.for_model("minimax_m2", torch_dtype="bfloat16", **TINY) + model = MiniMaxM2ForCausalLM(config, backend=_cpu_backend()) + block = model.model.layers["0"] + + # No initialize_weights on purpose: this is the state FSDP shards. + assert block.mlp.gate.weight.dtype == torch.float32 + assert block.mlp.gate.e_score_correction_bias.dtype == torch.float32 + + fully_shard_by_dtype( + block, + mesh=None, + mp_policy=fsdp.MixedPrecisionPolicy(param_dtype=torch.bfloat16, reduce_dtype=torch.float32), + offload_policy=None, + fp32_compute_module_names=tuple(MiniMaxM2ForCausalLM._keep_in_fp32_modules_strict), + fully_shard_fn=lambda *args, **kwargs: None, + ) + + +def test_explicit_gate_precision_override_is_preserved(): + config = AutoConfig.for_model("minimax_m2", torch_dtype="bfloat16", **TINY) + backend = _cpu_backend() + backend.gate_precision = torch.bfloat16 + model = MiniMaxM2ForCausalLM(config, backend=backend).eval() + assert model.model.backend.gate_precision == torch.bfloat16 diff --git a/tests/unit_tests/models/nemotron_v3/test_nemotron_v3_state_dict_adapter.py b/tests/unit_tests/models/nemotron_v3/test_nemotron_v3_state_dict_adapter.py index ddb731fea9..e5b133a821 100644 --- a/tests/unit_tests/models/nemotron_v3/test_nemotron_v3_state_dict_adapter.py +++ b/tests/unit_tests/models/nemotron_v3/test_nemotron_v3_state_dict_adapter.py @@ -90,13 +90,24 @@ def test_adapter_init(self, config, moe_config, backend): assert adapter.moe_config == moe_config assert adapter.backend == backend assert adapter.dtype == torch.bfloat16 - assert adapter._uses_model_prefix is True + assert adapter._uses_model_prefix is False def test_hf_prefix_property(self, config, moe_config, backend): - """Test _hf_prefix property returns 'backbone.'""" + """The output prefix tracks the namespace detected from the source checkpoint.""" adapter = NemotronV3StateDictAdapter(config, moe_config, backend) assert adapter._hf_prefix == "backbone." + adapter._uses_model_prefix = True + assert adapter._hf_prefix == "model." + + def test_peft_target_module_mapping_tracks_hf_prefix(self, config, moe_config, backend): + """PEFT target-module metadata uses the same public namespace as exported tensors.""" + adapter = NemotronV3StateDictAdapter(config, moe_config, backend) + native_name = "model.layers.0.mixer.in_proj" + + assert adapter.map_peft_target_module_to_hf(native_name) == "backbone.layers.0.mixer.in_proj" + adapter._uses_model_prefix = True + assert adapter.map_peft_target_module_to_hf(native_name) == native_name def test_expert_path_segment_property(self, config, moe_config, backend): """Test _expert_path_segment property returns 'mixer.experts'.""" @@ -161,6 +172,7 @@ def test_from_hf_renames_without_experts(self, adapter): assert "lm_head.weight" in native assert not any(k.startswith("backbone.") for k in native) assert not any(k.endswith("norm_f.weight") for k in native) + assert adapter._uses_model_prefix is False def test_round_trip_dense(self, adapter): hf_sd = { @@ -173,6 +185,18 @@ def test_round_trip_dense(self, adapter): assert set(back.keys()) == set(hf_sd.keys()) + def test_peft_outer_prefix_round_trip(self, adapter): + hf_key = "base_model.model.backbone.layers.0.mixer.in_proj.lora_A.weight" + native_key = "base_model.model.model.layers.0.mixer.in_proj.lora_A.weight" + tensor = torch.randn(8, 256) + + exported = adapter.to_hf({native_key: tensor}) + restored = adapter.from_hf(dict(exported)) + + assert list(exported) == [hf_key] + assert list(restored) == [native_key] + torch.testing.assert_close(restored[native_key], tensor) + class TestNemotronV3AdapterMTP: """MTP checkpoint namespace regressions.""" diff --git a/tests/unit_tests/models/test_moe_peft_v5_state_dict_adapters.py b/tests/unit_tests/models/test_moe_peft_v5_state_dict_adapters.py index a421ef30de..cdbb57d3d3 100644 --- a/tests/unit_tests/models/test_moe_peft_v5_state_dict_adapters.py +++ b/tests/unit_tests/models/test_moe_peft_v5_state_dict_adapters.py @@ -108,6 +108,9 @@ def _make_adapter_and_state(family: str, rank: int): adapter = NemotronV3StateDictAdapter( SimpleNamespace(num_hidden_layers=1), moe_config, backend, dtype=torch.float32 ) + # This fixture uses Transformers v5's native ``model.*`` hierarchy; + # remote-code Nemotron-H checkpoints instead select ``backbone.*``. + adapter._uses_model_prefix = True expert_path = "mixer.experts" base = f"base_model.model.model.layers.0.{expert_path}" diff --git a/tests/unit_tests/moe/test_layers.py b/tests/unit_tests/moe/test_layers.py index 759ffc01b2..cd17b6e6f6 100644 --- a/tests/unit_tests/moe/test_layers.py +++ b/tests/unit_tests/moe/test_layers.py @@ -1876,3 +1876,67 @@ def test_apply_bias_is_not_compiled(self): # torch.compile wraps functions in OptimizedModule or similar assert not hasattr(_apply_bias, "_torchdynamo_orig_callable"), "_apply_bias should not be torch.compiled" + + +class TestSigmoidGateScoringPrecision: + """Sigmoid routing must score in fp32 by default, like the softmax path. + + HF sigmoid-router references compute ``sigmoid(logits.float())``; scoring in + bf16 quantizes scores at ~2e-3, which flips knife-edge selections against + fine-grained ``e_score_correction_bias`` lattices (AMINT-286). + """ + + def _sigmoid_config(self): + return MoEConfig( + n_routed_experts=16, + n_shared_experts=0, + n_activated_experts=4, + n_expert_groups=0, + n_limited_groups=0, + train_gate=True, + gate_bias_update_factor=0.0, + aux_loss_coeff=0.0, + score_func="sigmoid", + route_scale=1.0, + dim=64, + inter_dim=128, + moe_inter_dim=128, + norm_topk_prob=True, + router_bias=False, + expert_bias=False, + expert_activation="swiglu", + force_e_score_correction_bias=True, + dtype=torch.bfloat16, + ) + + def test_sigmoid_scoring_matches_fp32_reference_on_bf16_inputs(self): + torch.manual_seed(0) + config = self._sigmoid_config() + gate = Gate(config) + with torch.no_grad(): + gate.weight.copy_(torch.randn_like(gate.weight) * 0.05) + # Knife-edge lattice bias like the MiniMax-M2.7 / GLM-4.7 checkpoints: + # large magnitude, 1e-3 spacing. + gate.e_score_correction_bias.copy_(8.0 + torch.arange(config.n_routed_experts, dtype=torch.float32) * 1e-3) + gate.eval() + + # x: Tensor of shape [tokens, hidden] in bf16, like real routed inputs. + x = torch.randn(512, config.dim, dtype=torch.bfloat16) + token_mask = torch.ones(x.shape[0], dtype=torch.bool) + weights, indices, _aux = gate(x, token_mask, None) + + # fp32 reference: the bf16 gate matmul followed by fp32 sigmoid, + # bias-augmented selection, and fp32 top-k normalization. + scores_bf16 = F.linear(x, gate.weight.to(x.dtype)) + ref_scores = torch.sigmoid(scores_bf16.float()) + choice = ref_scores + gate.e_score_correction_bias + ref_indices = torch.topk(choice, config.n_activated_experts, dim=-1)[1] + + assert torch.equal(indices.sort(dim=-1).values, ref_indices.sort(dim=-1).values) + + ref_weights = ref_scores.gather(1, indices) + ref_weights = ref_weights / (ref_weights.sum(dim=-1, keepdim=True) + 1e-20) + # The gate casts final weights back to the input dtype, mirroring the + # HF reference's top_k_weights.to(router_logits.dtype). + assert weights.dtype == x.dtype + assert torch.equal(weights, ref_weights.to(x.dtype))