diff --git a/docs/fern/versions/nightly/pages/models/qwen/qwen.mdx b/docs/fern/versions/nightly/pages/models/qwen/qwen.mdx index 9ac9ac2c73..18e0eb9080 100644 --- a/docs/fern/versions/nightly/pages/models/qwen/qwen.mdx +++ b/docs/fern/versions/nightly/pages/models/qwen/qwen.mdx @@ -218,9 +218,9 @@ uv run python examples/conversion/hf_to_megatron_generate_text.py \ | Model | Mode | TP | PP | EP | Total GPUs | Use Case | |-------|------|----|----|----|-----------:|----------| -| **Qwen3-30B-A3B** | Pretrain | 1 | 1 | 8 | 8 | Pre-training (single node) | -| **Qwen3-30B-A3B** | Full SFT | 1 | 1 | 8 | 8 | Full supervised finetuning | -| **Qwen3-30B-A3B** | LoRA/DoRA | 1 | 1 | 8 | 8 | PEFT finetuning (single node) | +| **Qwen3-30B-A3B** | Pretrain | 1 | 1 | 16 | 16 | Pre-training (2 nodes) | +| **Qwen3-30B-A3B** | Full SFT | 1 | 1 | 16 | 16 | Full supervised finetuning (2 nodes) | +| **Qwen3-30B-A3B** | LoRA/DoRA | 4 | 1 | 4 | 4 | PEFT finetuning | | **Qwen3-235B-A22B** | Pretrain | 2 | 8 | 32 | 512 | Pre-training (64 nodes) | | **Qwen3-235B-A22B** | Full SFT | 2 | 8 | 32 | 512 | Full supervised finetuning (64 nodes) | | **Qwen3-235B-A22B** | LoRA/DoRA | 2 | 8 | 32 | 512 | PEFT finetuning (64 nodes) | @@ -239,7 +239,7 @@ config = qwen3_30b_a3b_pretrain_config( train_iters=500_000, global_batch_size=2048, seq_length=4096, - # Uses TP=1, PP=1, EP=8 (8 GPUs) automatically + # Uses TP=1, PP=1, EP=16 (16 GPUs) automatically ) ``` @@ -272,7 +272,7 @@ config = qwen3_30b_a3b_sft_config( train_iters=1000, global_batch_size=64, finetune_lr=5e-6, - # Uses TP=1, PP=1, EP=8 (8 GPUs) automatically + # Uses TP=1, PP=1, EP=16 (16 GPUs) automatically ) ``` @@ -288,7 +288,7 @@ config = qwen3_30b_a3b_peft_config( train_iters=1000, global_batch_size=128, finetune_lr=1e-4, - # Uses TP=1, PP=1, EP=8 (8 GPUs) automatically + # Uses TP=4, PP=1, EP=4 (4 GPUs) automatically ) ``` diff --git a/docs/models/qwen/qwen.md b/docs/models/qwen/qwen.md index 9ac9ac2c73..18e0eb9080 100644 --- a/docs/models/qwen/qwen.md +++ b/docs/models/qwen/qwen.md @@ -218,9 +218,9 @@ uv run python examples/conversion/hf_to_megatron_generate_text.py \ | Model | Mode | TP | PP | EP | Total GPUs | Use Case | |-------|------|----|----|----|-----------:|----------| -| **Qwen3-30B-A3B** | Pretrain | 1 | 1 | 8 | 8 | Pre-training (single node) | -| **Qwen3-30B-A3B** | Full SFT | 1 | 1 | 8 | 8 | Full supervised finetuning | -| **Qwen3-30B-A3B** | LoRA/DoRA | 1 | 1 | 8 | 8 | PEFT finetuning (single node) | +| **Qwen3-30B-A3B** | Pretrain | 1 | 1 | 16 | 16 | Pre-training (2 nodes) | +| **Qwen3-30B-A3B** | Full SFT | 1 | 1 | 16 | 16 | Full supervised finetuning (2 nodes) | +| **Qwen3-30B-A3B** | LoRA/DoRA | 4 | 1 | 4 | 4 | PEFT finetuning | | **Qwen3-235B-A22B** | Pretrain | 2 | 8 | 32 | 512 | Pre-training (64 nodes) | | **Qwen3-235B-A22B** | Full SFT | 2 | 8 | 32 | 512 | Full supervised finetuning (64 nodes) | | **Qwen3-235B-A22B** | LoRA/DoRA | 2 | 8 | 32 | 512 | PEFT finetuning (64 nodes) | @@ -239,7 +239,7 @@ config = qwen3_30b_a3b_pretrain_config( train_iters=500_000, global_batch_size=2048, seq_length=4096, - # Uses TP=1, PP=1, EP=8 (8 GPUs) automatically + # Uses TP=1, PP=1, EP=16 (16 GPUs) automatically ) ``` @@ -272,7 +272,7 @@ config = qwen3_30b_a3b_sft_config( train_iters=1000, global_batch_size=64, finetune_lr=5e-6, - # Uses TP=1, PP=1, EP=8 (8 GPUs) automatically + # Uses TP=1, PP=1, EP=16 (16 GPUs) automatically ) ``` @@ -288,7 +288,7 @@ config = qwen3_30b_a3b_peft_config( train_iters=1000, global_batch_size=128, finetune_lr=1e-4, - # Uses TP=1, PP=1, EP=8 (8 GPUs) automatically + # Uses TP=4, PP=1, EP=4 (4 GPUs) automatically ) ``` diff --git a/examples/conversion/compare_hf_and_megatron/compare.py b/examples/conversion/compare_hf_and_megatron/compare.py index a392ff310a..610e14b5b4 100644 --- a/examples/conversion/compare_hf_and_megatron/compare.py +++ b/examples/conversion/compare_hf_and_megatron/compare.py @@ -100,6 +100,8 @@ import torch import torch.distributed as dist from megatron.core import parallel_state +from megatron.core.inference.contexts import StaticInferenceContext +from megatron.core.inference.utils import InferenceMode from megatron.core.pipeline_parallel.schedules import get_forward_backward_func from transformers import AutoConfig, AutoModelForCausalLM, AutoProcessor, AutoTokenizer @@ -125,8 +127,8 @@ from megatron.bridge.utils.safe_url import is_safe_public_http_url, safe_url_open -# Cosine similarity threshold: require at least 98% similarity (2% tolerance) -SIMILARITY_THRESHOLD = 0.98 +# Cosine similarity threshold: require at least 99% similarity (1% cosine distance) +SIMILARITY_THRESHOLD = 0.99 sys.path.append(os.path.dirname(__file__)) @@ -210,7 +212,16 @@ def get_model_class(model_class_name: str = None, is_vl_model: bool = False): return AutoModelForCausalLM -def is_vision_language_model(model_path: str, trust_remote_code: bool | None = None) -> bool: +def _hf_revision_kwargs(revision: str | None) -> dict[str, str]: + """Return an optional immutable Hugging Face revision argument.""" + return {"revision": revision} if revision is not None else {} + + +def is_vision_language_model( + model_path: str, + trust_remote_code: bool | None = None, + revision: str | None = None, +) -> bool: """Check if the model is a vision-language model. Args: @@ -226,6 +237,7 @@ def is_vision_language_model(model_path: str, trust_remote_code: bool | None = N trust_remote_code=trust_remote_code, hf_path=model_path, ), + **_hf_revision_kwargs(revision), ) # Check for VL model indicators in config @@ -264,11 +276,20 @@ class SingleBatchIterator: then raises StopIteration. Used for single-step inference in the forward pass. """ - def __init__(self, input_ids, position_ids, attention_mask, pixel_values=None, image_grid_thw=None): + def __init__( + self, + input_ids, + position_ids, + attention_mask, + pixel_values=None, + image_grid_thw=None, + inference_context=None, + ): self.batch = dict( tokens=input_ids, position_ids=position_ids, attention_mask=attention_mask, + inference_context=inference_context, ) # Add vision inputs if provided @@ -329,6 +350,47 @@ def loss_func(x, **kwargs): return output_tensor, loss_func +def inference_forward_step(data_iterator, model, **kwargs) -> torch.Tensor: + """Run a text-model forward step with an explicit inference context.""" + batch = next(data_iterator) + + def loss_func(x, **kwargs): + return x + + model_output = model( + input_ids=batch["tokens"], + position_ids=batch["position_ids"], + attention_mask=batch.get("attention_mask"), + inference_context=batch["inference_context"], + runtime_gather_output=True, + ) + if isinstance(model_output, tuple): + model_output = model_output[0] + return model_output, loss_func + + +def _run_megatron_forward(fwd_bwd_function, **kwargs): + """Run a Megatron forward pass with the inference execution paths active.""" + with InferenceMode.active(): + return fwd_bwd_function(**kwargs) + + +def _maybe_gather_tensor_parallel_logits(megatron_output, hf_vocab_size: int, world_size: int, group): + """Gather sharded TP logits while preserving an already gathered full-vocabulary tensor.""" + if megatron_output.size(-1) >= hf_vocab_size: + return megatron_output + + gathered_tensors = [torch.zeros_like(megatron_output) for _ in range(world_size)] + dist.all_gather(gathered_tensors, megatron_output, group=group) + gathered_output = torch.cat(gathered_tensors, dim=2) + if gathered_output.size(-1) < hf_vocab_size: + raise ValueError( + f"Gathered Megatron vocabulary ({gathered_output.size(-1)}) is smaller than " + f"the Hugging Face vocabulary ({hf_vocab_size})." + ) + return gathered_output + + def load_image(image_path: str) -> Image.Image: """Load an image from URL or file path. @@ -453,6 +515,7 @@ def _load_hf_model(args, is_vl_model: bool): trust_remote_code=args.trust_remote_code, hf_path=args.hf_model_path, ), + **_hf_revision_kwargs(args.hf_revision), ) hf_model = hf_model.eval() print_rank_0(f"Loaded with {model_class.__name__}") @@ -581,7 +644,14 @@ def _load_megatron_model(args): if args.megatron_model_path: # Load from Megatron checkpoint - bridge = AutoBridge.from_hf_pretrained(args.hf_model_path) + bridge = AutoBridge.from_hf_pretrained( + args.hf_model_path, + trust_remote_code=is_safe_repo( + trust_remote_code=args.trust_remote_code, + hf_path=args.hf_model_path, + ), + **_hf_revision_kwargs(args.hf_revision), + ) model_provider = bridge.to_megatron_provider(load_weights=False) model_provider.tensor_model_parallel_size = tp model_provider.pipeline_model_parallel_size = pp @@ -608,6 +678,7 @@ def _load_megatron_model(args): trust_remote_code=args.trust_remote_code, hf_path=args.hf_model_path, ), + **_hf_revision_kwargs(args.hf_revision), ) model_provider = bridge.to_megatron_provider(load_weights=True) model_provider.tensor_model_parallel_size = tp @@ -650,6 +721,7 @@ def _setup_tokenizer_and_processor(args, is_vl_model: bool): trust_remote_code=args.trust_remote_code, hf_path=args.hf_model_path, ), + **_hf_revision_kwargs(args.hf_revision), ) if tokenizer.pad_token is None: tokenizer.pad_token = tokenizer.eos_token @@ -663,6 +735,7 @@ def _setup_tokenizer_and_processor(args, is_vl_model: bool): trust_remote_code=args.trust_remote_code, hf_path=args.hf_model_path, ), + **_hf_revision_kwargs(args.hf_revision), ) except Exception as e: print_rank_0(f"Warning: Could not load processor for VL model: {e}") @@ -671,6 +744,29 @@ def _setup_tokenizer_and_processor(args, is_vl_model: bool): return tokenizer, processor +def _broadcast_hf_results(hf_logits, hf_next_token, device): + """Broadcast rank-0 HF results using the model's actual output vocabulary size.""" + if hf_logits is not None: + hf_logits = hf_logits.float() + + hf_logits_size = torch.tensor( + [hf_logits.numel() if hf_logits is not None else 0], + device=device, + dtype=torch.long, + ) + torch.distributed.broadcast(hf_logits_size, 0) + + if hf_next_token is None: + hf_next_token = torch.zeros(1, device=device, dtype=torch.long) + if hf_logits is None: + hf_logits = torch.zeros(hf_logits_size.item(), device=device, dtype=torch.float32) + + torch.distributed.broadcast(hf_next_token, 0) + torch.distributed.broadcast(hf_logits, 0) + torch.distributed.barrier() + return hf_logits, hf_next_token + + def compare_models_one_step(args) -> None: """Compare 1-step generation between HF and Megatron models with debugging. @@ -685,7 +781,7 @@ def compare_models_one_step(args) -> None: print_rank_0(f"Set CUDA device to: {torch.cuda.current_device()}") # Detect model type - is_vl_model = is_vision_language_model(args.hf_model_path, args.trust_remote_code) + is_vl_model = is_vision_language_model(args.hf_model_path, args.trust_remote_code, args.hf_revision) print_rank_0(f"Detected model type: {'Vision-Language' if is_vl_model else 'Text-only LLM'}") # Validate vision requirements @@ -733,28 +829,10 @@ def compare_models_one_step(args) -> None: # Broadcast HF results to all ranks if torch.distributed.is_initialized(): - # Ensure consistent dtype across ranks: rank 0 has bfloat16 logits from the HF model, - # so all ranks must use the same dtype for NCCL broadcast to work correctly. - if hf_logits is not None: - hf_logits = hf_logits.float() - - # Create tensors for broadcasting if they don't exist on non-rank-0 - if hf_next_token is None: - hf_next_token = torch.zeros(1, device=input_ids.device, dtype=torch.long) - if hf_logits is None: - # Get vocab size from tokenizer for proper tensor size - vocab_size = getattr( - tokenizer, "vocab_size", len(tokenizer.vocab) if hasattr(tokenizer, "vocab") else 32000 - ) - hf_logits = torch.zeros(vocab_size, device=input_ids.device, dtype=torch.float32) - - # Ensure consistent dtype across ranks before broadcast - hf_logits = hf_logits.float() - - # Broadcast from rank 0 to all ranks - torch.distributed.broadcast(hf_next_token, 0) - torch.distributed.broadcast(hf_logits, 0) - torch.distributed.barrier() + # The model's output vocabulary can be larger than the tokenizer vocabulary. + # Broadcast the actual logits length before allocating receive buffers so every + # rank participates in the logits broadcast with the same tensor shape. + hf_logits, hf_next_token = _broadcast_hf_results(hf_logits, hf_next_token, input_ids.device) print_rank_0("HF results broadcast complete.") # Run Megatron model forward pass @@ -773,18 +851,44 @@ def compare_models_one_step(args) -> None: attention_mask = None fwd_bwd_function = get_forward_backward_func() - iterator = SingleBatchIterator(input_ids, position_ids, attention_mask, pixel_values, image_grid_thw) - - megatron_output = fwd_bwd_function( - forward_step_func=vlm_forward_step, - data_iterator=iterator, - model=megatron_model, - num_microbatches=1, - forward_only=True, - seq_length=input_ids.size(1), - micro_batch_size=1, - collect_non_loss_data=True, - ) + forward_kwargs = { + "model": megatron_model, + "num_microbatches": 1, + "forward_only": True, + "seq_length": input_ids.size(1), + "micro_batch_size": 1, + "collect_non_loss_data": True, + } + if is_vl_model: + iterator = SingleBatchIterator( + input_ids, + position_ids, + attention_mask, + pixel_values, + image_grid_thw, + ) + megatron_output = fwd_bwd_function( + forward_step_func=vlm_forward_step, + data_iterator=iterator, + **forward_kwargs, + ) + else: + inference_context = StaticInferenceContext( + max_batch_size=input_ids.size(0), + max_sequence_length=input_ids.size(1), + ) + iterator = SingleBatchIterator( + input_ids, + position_ids, + attention_mask, + inference_context=inference_context, + ) + megatron_output = _run_megatron_forward( + fwd_bwd_function, + forward_step_func=inference_forward_step, + data_iterator=iterator, + **forward_kwargs, + ) if isinstance(megatron_output, list) and len(megatron_output) > 0: megatron_output = megatron_output[0] @@ -796,11 +900,12 @@ def compare_models_one_step(args) -> None: # Gather tensor parallel results if using TP if torch.distributed.is_initialized() and parallel_state.get_tensor_model_parallel_world_size() > 1: world_size = parallel_state.get_tensor_model_parallel_world_size() - gathered_tensors = [torch.zeros_like(megatron_output) for _ in range(world_size)] - dist.all_gather( - gathered_tensors, megatron_output, group=parallel_state.get_tensor_model_parallel_group() + megatron_output = _maybe_gather_tensor_parallel_logits( + megatron_output, + hf_logits.size(0), + world_size, + parallel_state.get_tensor_model_parallel_group(), ) - megatron_output = torch.cat(gathered_tensors, dim=2) megatron_logits = megatron_output[0, -1, :] megatron_next_token = torch.argmax(megatron_logits, dim=-1) @@ -838,8 +943,12 @@ def compare_models_one_step(args) -> None: cos_val = cosine_sim.item() percent = cos_val * 100.0 status_emoji = "✅" if cos_val >= SIMILARITY_THRESHOLD else "❌" - tolerance_text = "within ±2%" if cos_val >= SIMILARITY_THRESHOLD else "outside ±2%" - print(f"Cosine similarity: {cos_val:.6f} ({percent:.2f}%) {status_emoji} ({tolerance_text} tolerance)") + limit_text = "within" if cos_val >= SIMILARITY_THRESHOLD else "outside" + distance_limit = 1.0 - SIMILARITY_THRESHOLD + print( + f"Cosine similarity: {cos_val:.6f} ({percent:.2f}%) {status_emoji} " + f"({limit_text} {distance_limit:.0%} cosine-distance limit)" + ) print("=== COMPARISON COMPLETE ===") else: @@ -860,6 +969,10 @@ def build_parser() -> argparse.ArgumentParser: required=True, help="Path to the HuggingFace model.", ) + parser.add_argument( + "--hf-revision", + help="Immutable Hugging Face Hub revision used for model, config, and tokenizer loading.", + ) parser.add_argument( "--prompt", type=str, diff --git a/examples/conversion/hf_to_megatron_generate_text.py b/examples/conversion/hf_to_megatron_generate_text.py index 7d173192dc..f8da1f4ad2 100644 --- a/examples/conversion/hf_to_megatron_generate_text.py +++ b/examples/conversion/hf_to_megatron_generate_text.py @@ -26,6 +26,8 @@ import torch import torch.distributed as dist from megatron.core import parallel_state +from megatron.core.inference.contexts import StaticInferenceContext +from megatron.core.inference.utils import InferenceMode from megatron.core.pipeline_parallel.schedules import get_forward_backward_func from transformers import AutoTokenizer @@ -43,10 +45,11 @@ class SingleBatchIterator: Used for single-step inference in the forward pass. """ - def __init__(self, input_ids, position_ids): + def __init__(self, input_ids, position_ids, inference_context=None): self.batch = dict( tokens=input_ids, position_ids=position_ids, + inference_context=inference_context, ) self._yielded = False @@ -60,6 +63,11 @@ def __next__(self): return self.batch +def _hf_revision_kwargs(revision: str | None) -> dict[str, str]: + """Return an optional immutable Hugging Face revision argument.""" + return {"revision": revision} if revision is not None else {} + + def text_forward_step(data_iterator, model, **kwargs) -> torch.Tensor: """Forward step function for text generation. Required by the forward_backward_func function. @@ -80,12 +88,39 @@ def text_forward_step(data_iterator, model, **kwargs) -> torch.Tensor: "input_ids": batch["tokens"], "position_ids": batch["position_ids"], "attention_mask": batch.get("attention_mask", None), + "inference_context": batch["inference_context"], + "runtime_gather_output": True, } def loss_func(x, **kwargs): return x - return model(**forward_args), loss_func + model_output = model(**forward_args) + if isinstance(model_output, tuple): + model_output = model_output[0] + return model_output, loss_func + + +def _run_megatron_forward(fwd_bwd_function, **kwargs): + """Run a Megatron forward pass with inference execution paths active.""" + with InferenceMode.active(): + return fwd_bwd_function(**kwargs) + + +def _maybe_gather_tensor_parallel_logits(output, full_vocab_size: int, world_size: int, group): + """Gather sharded TP logits without gathering an already complete vocabulary.""" + if output.size(-1) >= full_vocab_size: + return output + + gathered_tensors = [torch.zeros_like(output) for _ in range(world_size)] + dist.all_gather(gathered_tensors, output, group=group) + gathered_output = torch.cat(gathered_tensors, dim=2) + if gathered_output.size(-1) < full_vocab_size: + raise ValueError( + f"Gathered Megatron vocabulary ({gathered_output.size(-1)}) is smaller than " + f"the configured vocabulary ({full_vocab_size})." + ) + return gathered_output def _tokenize_prompt(tokenizer, prompt: str, *, apply_chat_template: bool, thinking_mode: str) -> torch.Tensor: @@ -142,6 +177,7 @@ def main(args) -> None: trust_remote_code=args.trust_remote_code, hf_path=args.hf_model_path, ), + **_hf_revision_kwargs(args.hf_revision), ) # Initialize model parallel before loading @@ -195,6 +231,7 @@ def main(args) -> None: trust_remote_code=args.trust_remote_code, hf_path=args.hf_model_path, ), + **_hf_revision_kwargs(args.hf_revision), ) model_provider = bridge.to_megatron_provider(load_weights=True) model_provider.tensor_model_parallel_size = tp @@ -220,6 +257,7 @@ def main(args) -> None: trust_remote_code=args.trust_remote_code, hf_path=args.hf_model_path, ), + **_hf_revision_kwargs(args.hf_revision), ) if tokenizer.pad_token is None: tokenizer.pad_token = tokenizer.eos_token @@ -246,9 +284,14 @@ def main(args) -> None: print_rank_0(f"Generation step {step}") fwd_bwd_function = get_forward_backward_func() - iterator = SingleBatchIterator(input_ids, position_ids) + inference_context = StaticInferenceContext( + max_batch_size=input_ids.size(0), + max_sequence_length=input_ids.size(1), + ) + iterator = SingleBatchIterator(input_ids, position_ids, inference_context) - output = fwd_bwd_function( + output = _run_megatron_forward( + fwd_bwd_function, forward_step_func=text_forward_step, data_iterator=iterator, model=model, @@ -263,11 +306,12 @@ def main(args) -> None: if parallel_state.is_pipeline_last_stage(): world_size = parallel_state.get_tensor_model_parallel_world_size() - gathered_tensors = [torch.zeros_like(output) for _ in range(world_size)] - # All-gather operation - dist.all_gather(gathered_tensors, output, group=parallel_state.get_tensor_model_parallel_group()) - # Concatenate along last dimension (dim=2) - output = torch.cat(gathered_tensors, dim=2) + output = _maybe_gather_tensor_parallel_logits( + output, + model_provider.vocab_size, + world_size, + parallel_state.get_tensor_model_parallel_group(), + ) next_token_ids = torch.argmax(output[:, -1], dim=-1, keepdim=True) # Debug: print token information @@ -306,7 +350,8 @@ def main(args) -> None: print_rank_0("=======================================") -if __name__ == "__main__": +def build_parser() -> argparse.ArgumentParser: + """Build the text-generation CLI parser.""" parser = argparse.ArgumentParser(description="Text Generation from HuggingFace Models") parser.add_argument( "--hf_model_path", @@ -314,6 +359,10 @@ def main(args) -> None: required=True, help="Path to the HuggingFace model.", ) + parser.add_argument( + "--hf-revision", + help="Immutable Hugging Face Hub revision used for config and tokenizer loading.", + ) parser.add_argument( "--prompt", type=str, @@ -343,7 +392,11 @@ def main(args) -> None: parser.add_argument("--etp", type=int, default=1, help="Expert tensor parallelism size") parser.add_argument("--megatron_model_path", type=str, default=None, help="Path to the Megatron model checkpoint") parser.add_argument("--trust-remote-code", action="store_true", help="if trust_remote_code") - args = parser.parse_args() + return parser + + +if __name__ == "__main__": + args = build_parser().parse_args() main(args) diff --git a/model_cards/moonlight-16b-a3b/card.yaml b/model_cards/moonlight-16b-a3b/card.yaml new file mode 100644 index 0000000000..fffba10da9 --- /dev/null +++ b/model_cards/moonlight-16b-a3b/card.yaml @@ -0,0 +1,344 @@ +# Agent-readable model support card. +# status: unverified | verified | unsupported | not_applicable + +title: moonlight_16b_a3b +model: + hf_id: moonshotai/Moonlight-16B-A3B + hf_revision: 476b36a473d4467f94469414bef6cee75c9c8172 # pragma: allowlist secret + architecture: DeepseekV3ForCausalLM + min_transformers_version: "5.8.1" +verification_environment: + base_container: nvcr.io/nvidia/pytorch:26.04-py3 + bridge_commit: 93c6930597fd3766acd177288b7a8607d7b96ad7 # pragma: allowlist secret +summary: Moonlight-16B-A3B support verification for conversion, inference, and training. + +items: + hf_to_megatron_cpu: + status: verified + precision: bf16 + command: > + ./scripts/conversion/convert.sh import --executor slurm --device cpu --nodes 1 + --hf-model moonshotai/Moonlight-16B-A3B + --megatron-path work/model-verification/moonlight-16b-a3b/cpu-megatron + --torch-dtype bfloat16 --trust-remote-code + last_verified: 2026-07-19 + expected_result: > + The command exits successfully, creates iter_0000000, and the checkpoint + round-trips through CPU export with exact model configuration. After + native reload, all 377 registered parameters and persistent buffers match + bitwise by key, shape, dtype, and value. In the serialized files, 26 router + correction-bias buffers widen losslessly from BF16 to FP32. + + hf_to_megatron_gpu: + status: verified + precision: bf16 + command: > + ./scripts/conversion/convert.sh import --executor slurm --device gpu + --nodes 1 --gpus-per-node 8 + --hf-model moonshotai/Moonlight-16B-A3B + --megatron-path work/model-verification/moonlight-16b-a3b/imported-megatron + --torch-dtype bfloat16 --tp 1 --pp 1 --ep 8 --etp 1 + --trust-remote-code + last_verified: 2026-07-19 + expected_result: > + The command exits successfully, creates a reloadable iter_0000000, and + an exact audit finds all 5,344 serialized source tensors and the model + configuration unchanged after the paired GPU export. + + megatron_to_hf_cpu: + status: verified + precision: bf16 + command: > + ./scripts/conversion/convert.sh export --executor slurm --device cpu --nodes 1 + --hf-model moonshotai/Moonlight-16B-A3B + --megatron-path work/model-verification/moonlight-16b-a3b/cpu-megatron/iter_0000000 + --hf-path work/model-verification/moonlight-16b-a3b/cpu-hf-export + --torch-dtype bfloat16 --trust-remote-code + last_verified: 2026-07-19 + expected_result: > + The command exits successfully, the model configuration is exact, and all + 377 registered parameters and persistent buffers match bitwise by key, + shape, dtype, and value after native reload. Transformers strictly reloads + the output as DeepseekV3ForCausalLM with trust_remote_code disabled. In the + serialized files, 26 router correction-bias buffers widen losslessly from + BF16 to FP32 and 27 stale, nonpersistent source inv_freq tensors are + omitted. + + megatron_to_hf_gpu: + status: verified + precision: bf16 + command: > + ./scripts/conversion/convert.sh export --executor slurm --device gpu + --nodes 1 --gpus-per-node 8 + --hf-model moonshotai/Moonlight-16B-A3B + --megatron-path work/model-verification/moonlight-16b-a3b/imported-megatron/iter_0000000 + --hf-path work/model-verification/moonlight-16b-a3b/hf-export + --torch-dtype bfloat16 --export-weight-dtype bfloat16 + --distributed-save --tp 1 --pp 1 --ep 8 --etp 1 + --trust-remote-code + last_verified: 2026-07-19 + expected_result: > + Strict distributed export exits successfully, all 5,344 serialized tensor + keys, shapes, dtypes, and values match the recorded HF revision, and native + Transformers strictly reloads the output as DeepseekV3ForCausalLM with + trust_remote_code disabled and no missing or unexpected keys. + + manual_forward_pass: + status: verified + precision: bf16 + bridge_commit: 056616080f2d8202bfe3f636069584406b5e465e # pragma: allowlist secret + command: > + uv run python -m torch.distributed.run --standalone --nproc_per_node=8 + examples/conversion/compare_hf_and_megatron/compare.py + --hf_model_path moonshotai/Moonlight-16B-A3B + --megatron_model_path work/model-verification/moonlight-16b-a3b/imported-megatron/iter_0000000 + --tp 1 --pp 1 --ep 8 --etp 1 + --prompt "The capital of France is" + --model_class DeepseekV3ForCausalLM --trust-remote-code + last_verified: 2026-07-19 + expected_result: > + Hugging Face and Megatron have next-token matches at token ID 17374 + (" Paris"). Cosine similarity is 0.999893, above the 0.99 correlation + gate. The maximum absolute logit difference is 0.472656 and the mean + absolute logit difference is 0.074414; both are report-only diagnostic + observations. This historical result predates explicit helper revision + pinning and is retained against the recorded immutable HF revision. + + inference: + status: verified + precision: bf16 + command: > + uv run python -m torch.distributed.run --standalone --nproc_per_node=8 + examples/conversion/hf_to_megatron_generate_text.py + --hf_model_path moonshotai/Moonlight-16B-A3B + --megatron_model_path work/model-verification/moonlight-16b-a3b/imported-megatron/iter_0000000 + --tp 1 --pp 1 --ep 8 --etp 1 + --prompt "The capital of France is" --max_new_tokens 32 + --trust-remote-code + last_verified: 2026-07-19 + expected_result: > + Two independent executions using greedy decoding each produce exactly 32 + new tokens with byte-identical token IDs and this literal completion, + including its leading space, " Paris. It is the largest city in France + and is known for its iconic landmarks such as the Eiffel Tower, the + Louvre Museum, and Notre-Dame". + + pretrain: + status: unverified + precision: bf16 + gpu_type: H100 + enabled_features: {} + command: > + ./scripts/training/train.sh --nodes 2 --gpus-per-node 8 + --recipe moonlight_16b_pretrain_16gpu_h100_bf16_config --mode pretrain + --dataset megatron-indexed --seq_length 4096 --max_steps 100 + --lr 3e-4 --min_lr 3e-5 --warmup_iters 40 + 'dataset.blend=[["work/data/rp2/head_01"],null]' + dataset.path_to_cache=work/cache/moonlight-16b-a3b/rp2-convergence-v1 + tokenizer.tokenizer_type=SentencePieceTokenizer + tokenizer.tokenizer_model=work/data/rp2/tokenizer.model + dataset.random_seed=1234 rng.seed=1234 scheduler.lr_decay_iters=100 + validation.eval_iters=0 validation.eval_interval=0 + ddp.check_for_nan_in_grad=true ddp.check_for_large_grads=true + rerun_state_machine.check_for_nan_in_loss=true checkpoint.load=null + dist.distributed_timeout_minutes=30 + --save_dir work/model-verification/moonlight-16b-a3b/pretrain-convergence-v1-reference-checkpoints + --save_interval 50 logger.log_interval=1 logger.log_throughput=true + logger.tensorboard_dir=null + last_verified: null + metrics: + initial_loss: null + final_loss: null + last_10_steps_step_time_ms_avg: null + last_10_steps_model_tflops_per_gpu_avg: null + expected_result: > + The pending uninterrupted 16-GPU bounded RP2 reference must complete + exactly 100 steps at recipe-owned GBS/MBS 1024/1 with finite losses and + no skipped or NaN iterations. It must reach peak learning rate at step + 40, complete cosine decay at step 100, and save complete iter_0000050 and + iter_0000100 checkpoints before this cohort item can be marked verified. + + sft: + status: unverified + precision: bf16 + gpu_type: H100 + enabled_features: + sequence_packing: offline + command: > + ./scripts/training/train.sh --nodes 1 --gpus-per-node 8 + --recipe moonlight_16b_sft_8gpu_h100_bf16_tp1_config --mode sft + --dataset tulu3 + --pretrained_checkpoint work/model-verification/moonlight-16b-a3b/imported-megatron/iter_0000000 + --max_steps 100 --seq_length 2048 + --lr 5e-6 --min_lr 0 --warmup_iters 10 + 'dataset.hf_dataset.split="train[:10000]"' + 'dataset.hf_dataset.load_kwargs={revision:"b14afda60f1bbebe55d5d2fa1e4df5042f97f8be"}' + '++tokenizer.hf_tokenizer_kwargs.revision="476b36a473d4467f94469414bef6cee75c9c8172"' + dataset.hf_output_root=work/data/tulu3/moonlight-16b-a3b-sft-b14afda60f1b + dataset.hf_rewrite=true dataset.seed=1234 rng.seed=5678 + dataset.do_validation=false dataset.hf_validation_proportion=null + dataset.enable_offline_packing=true + +dataset.offline_packing_specs.pad_seq_to_mult=1 scheduler.lr_decay_iters=100 + validation.eval_iters=0 validation.eval_interval=0 checkpoint.load=null + --save_dir work/model-verification/moonlight-16b-a3b/sft-convergence-v1-checkpoints + --save_interval 100 logger.log_interval=1 logger.log_throughput=true + last_verified: null + metrics: + initial_loss: null + final_loss: null + last_10_steps_step_time_ms_avg: null + last_10_steps_model_tflops_per_gpu_avg: null + expected_result: > + The pending TP1/PP1/EP8 run must preserve recipe-owned GBS/MBS 32/1 and + pad-1 offline packing, complete exactly 100 full-SFT steps with finite + losses and no skipped or NaN iterations, and save a complete full-model + iter_0000100 checkpoint. Packing efficiency and the actual supervised- + token count after label masking must be recorded before verification. + + sft_export_inference: + status: unverified + precision: bf16 + depends_on: sft + commands: + - > + ./scripts/conversion/convert.sh export --executor slurm --device gpu + --nodes 1 --gpus-per-node 8 + --hf-model moonshotai/Moonlight-16B-A3B + --megatron-path work/model-verification/moonlight-16b-a3b/sft-convergence-v1-checkpoints/iter_0000100 + --hf-path work/model-verification/moonlight-16b-a3b/sft-convergence-v1-hf-export + --torch-dtype bfloat16 --export-weight-dtype bfloat16 + --distributed-save --tp 1 --pp 1 --ep 8 --etp 1 --trust-remote-code + - > + uv run python + skills/create-model-verification-card/scripts/verify_hf_inference.py + --hf-model work/model-verification/moonlight-16b-a3b/sft-convergence-v1-hf-export + --prompt "In one short sentence, explain why Paris is important to France." + --max-new-tokens 32 --chat-template + last_verified: null + expected_result: > + After the pending SFT run passes, its step-100 checkpoint must export with + a 163842-row embedding and output head, all index-referenced BF16 shards + must reload natively with no missing or unexpected keys, and the helper + must produce exactly 32 new tokens twice with identical token IDs. The + byte-identical literal completion must be recorded before verification. + + sft_long_context: + status: verified + precision: bf16 + bridge_commit: f3ae2767b5e18aeb67b726cd8d5f1db58216dcc9 # pragma: allowlist secret + gpu_type: H100 + enabled_features: + sequence_packing: offline + context_parallel_size: 2 + command: > + ./scripts/training/train.sh --nodes 1 --gpus-per-node 8 + --recipe moonlight_16b_sft_8k_config --mode sft + --dataset tulu3 + --pretrained_checkpoint work/model-verification/moonlight-16b-a3b/imported-megatron/iter_0000000 + --max_steps 20 --seq_length 8192 --context_parallel_size 2 + --lr 1e-6 --min_lr 0 --warmup_iters 2 + 'dataset.hf_dataset.split="train[:10000]"' + 'dataset.hf_dataset.load_kwargs={revision:"b14afda60f1bbebe55d5d2fa1e4df5042f97f8be"}' + '++tokenizer.hf_tokenizer_kwargs.revision="476b36a473d4467f94469414bef6cee75c9c8172"' + dataset.hf_output_root=work/data/tulu3/moonlight-16b-a3b-long-context-b14afda60f1b + dataset.hf_rewrite=true dataset.seed=1234 rng.seed=5678 + dataset.do_validation=false dataset.hf_validation_proportion=null + dataset.enable_offline_packing=true + +dataset.offline_packing_specs.pad_seq_to_mult=4 scheduler.lr_decay_iters=20 + validation.eval_iters=0 validation.eval_interval=0 + checkpoint.load=null checkpoint.save=null + logger.log_interval=1 logger.log_throughput=true + last_verified: 2026-07-20 + metrics: + initial_loss: 1.306440 + final_loss: 1.228019 + last_10_steps_step_time_ms_avg: 60599.310 + last_10_steps_model_tflops_per_gpu_avg: 40.840 + expected_result: > + The immutable-revision 8-GPU run completes exactly 20 Tulu3 SFT steps at + the model's 8192-token context limit with recipe-owned + TP2/PP1/CP2/EP8/SP-on, GBS/MBS 128/1, and explicit pad-4 offline packing. + LM loss is 1.306440 to 1.228019; skipped/NaN totals are 0/0. The persisted + post-setup runtime config matches the command, packing is 99.69%, and the + sampled training window contains 13,765,732 actual supervised tokens. + + peft: + status: unverified + precision: bf16 + gpu_type: H100 + enabled_features: + sequence_packing: offline + command: > + ./scripts/training/train.sh --nodes 1 --gpus-per-node 4 + --recipe moonlight_16b_peft_4gpu_h100_bf16_config --mode lora + --dataset tulu3 + --pretrained_checkpoint work/model-verification/moonlight-16b-a3b/imported-megatron/iter_0000000 + --max_steps 100 --seq_length 2048 + --lr 1e-4 --min_lr 0 --warmup_iters 10 + 'dataset.hf_dataset.split="train[:10000]"' + 'dataset.hf_dataset.load_kwargs={revision:"b14afda60f1bbebe55d5d2fa1e4df5042f97f8be"}' + '++tokenizer.hf_tokenizer_kwargs.revision="476b36a473d4467f94469414bef6cee75c9c8172"' + dataset.hf_output_root=work/data/tulu3/moonlight-16b-a3b-peft-b14afda60f1b + dataset.hf_rewrite=true dataset.seed=1234 rng.seed=5678 + dataset.do_validation=false dataset.hf_validation_proportion=null + dataset.enable_offline_packing=true + +dataset.offline_packing_specs.pad_seq_to_mult=4 scheduler.lr_decay_iters=100 + validation.eval_iters=0 validation.eval_interval=0 checkpoint.load=null + --save_dir work/model-verification/moonlight-16b-a3b/peft-convergence-v1-checkpoints + --save_interval 100 logger.log_interval=1 logger.log_throughput=true + last_verified: null + metrics: + initial_loss: null + final_loss: null + last_10_steps_step_time_ms_avg: null + last_10_steps_model_tflops_per_gpu_avg: null + expected_result: > + The pending TP4/PP1/EP4 run must preserve recipe-owned GBS/MBS 32/1 and + pad-4 offline packing, train only rank-8, alpha-16, zero-dropout LoRA on + Moonlight's model-native attention projections, and complete exactly 100 + steps with finite losses and no skipped or NaN iterations. Packing + efficiency, actual supervised-token count, and a complete iter_0000100 + adapter checkpoint must be recorded before verification. + + checkpoint_resume: + status: unverified + precision: bf16 + gpu_type: H100 + depends_on: pretrain + command: > + ./scripts/training/train.sh --nodes 2 --gpus-per-node 8 + --recipe moonlight_16b_pretrain_16gpu_h100_bf16_config --mode pretrain + --dataset megatron-indexed --seq_length 4096 --max_steps 100 + --lr 3e-4 --min_lr 3e-5 --warmup_iters 40 + 'dataset.blend=[["work/data/rp2/head_01"],null]' + dataset.path_to_cache=work/cache/moonlight-16b-a3b/rp2-convergence-v1 + tokenizer.tokenizer_type=SentencePieceTokenizer + tokenizer.tokenizer_model=work/data/rp2/tokenizer.model + dataset.random_seed=1234 rng.seed=1234 scheduler.lr_decay_iters=100 + validation.eval_iters=0 validation.eval_interval=0 + ddp.check_for_nan_in_grad=true ddp.check_for_large_grads=true + rerun_state_machine.check_for_nan_in_loss=true + dist.distributed_timeout_minutes=30 + --load_dir work/model-verification/moonlight-16b-a3b/pretrain-convergence-v1-reference-checkpoints + --save_dir work/model-verification/moonlight-16b-a3b/pretrain-convergence-v1-resumed-checkpoints + --save_interval 50 checkpoint.ckpt_step=50 + logger.log_interval=1 logger.log_throughput=true logger.tensorboard_dir=null + last_verified: null + metrics: + initial_loss: null + final_loss: null + last_10_steps_step_time_ms_avg: null + last_10_steps_model_tflops_per_gpu_avg: null + resume_comparison: + reference_item: pretrain + sentinel_steps: [51, 100] + loss_relative_tolerance: 1.0e-2 + loss_absolute_tolerance: 1.0e-6 + sentinels_match: null + expected_result: > + After the pending uninterrupted reference completes, this command must + restore optimizer, scheduler, data-order, and RNG state directly from + iter_0000050, begin at step 51, and finish at step 100 in the distinct + resumed output root. Step-51 and step-100 loss must satisfy the declared + absolute-plus-relative tolerance against the reference, with finite + losses and no skipped or NaN iterations. diff --git a/model_cards/nemotron-3-nano-4b/card.yaml b/model_cards/nemotron-3-nano-4b/card.yaml new file mode 100644 index 0000000000..0f459d7410 --- /dev/null +++ b/model_cards/nemotron-3-nano-4b/card.yaml @@ -0,0 +1,348 @@ +# Agent-readable model support card. +# status: unverified | verified | unsupported | not_applicable + +title: nemotron_3_nano_4b +model: + hf_id: nvidia/NVIDIA-Nemotron-3-Nano-4B-BF16 + hf_revision: dfaf35de3e30f1867dd8dbc38a7fc9fb52d3914f # pragma: allowlist secret + architecture: NemotronHForCausalLM + min_transformers_version: "5.8.1" +verification_environment: + base_container: nvcr.io/nvidia/pytorch:26.04-py3 + bridge_commit: 0e6c1837edeb442b07c9175b855b3d288c5cad33 # pragma: allowlist secret +summary: Nemotron 3 Nano 4B support verification for conversion, inference, and training. + +items: + hf_to_megatron_cpu: + status: verified + precision: bf16 + command: > + ./scripts/conversion/convert.sh import --executor slurm --device cpu --nodes 1 + --hf-model nvidia/NVIDIA-Nemotron-3-Nano-4B-BF16 + --hf-revision dfaf35de3e30f1867dd8dbc38a7fc9fb52d3914f + --megatron-path work/model-verification/nemotron-3-nano-4b/cpu-megatron + --torch-dtype bfloat16 + last_verified: 2026-07-20 + expected_result: > + The command exits successfully and creates iter_0000000. After the paired + CPU export, all 263 tensors and 3,973,556,832 parameters match the pinned + source exactly in name, shape, dtype, and value, with maximum difference + zero. + + hf_to_megatron_gpu: + status: verified + precision: bf16 + command: > + ./scripts/conversion/convert.sh import --executor slurm --device gpu + --nodes 1 --gpus-per-node 1 --tp 1 + --hf-model nvidia/NVIDIA-Nemotron-3-Nano-4B-BF16 + --hf-revision dfaf35de3e30f1867dd8dbc38a7fc9fb52d3914f + --megatron-path work/model-verification/nemotron-3-nano-4b/imported-megatron + --torch-dtype bfloat16 + last_verified: 2026-07-20 + expected_result: > + The command exits successfully and creates a reloadable iter_0000000. + After the paired GPU export, all 263 tensors and 3,973,556,832 parameters + match the pinned source exactly in name, shape, dtype, and value. + + megatron_to_hf_cpu: + status: verified + precision: bf16 + command: > + ./scripts/conversion/convert.sh export --executor slurm --device cpu --nodes 1 + --hf-model nvidia/NVIDIA-Nemotron-3-Nano-4B-BF16 + --hf-revision dfaf35de3e30f1867dd8dbc38a7fc9fb52d3914f + --megatron-path work/model-verification/nemotron-3-nano-4b/cpu-megatron + --hf-path work/model-verification/nemotron-3-nano-4b/cpu-hf-export + --torch-dtype bfloat16 + last_verified: 2026-07-20 + expected_result: > + Strict export exits successfully; all 263 BF16 tensors match the pinned + source bitwise, the 171-byte generation configuration is preserved + byte-for-byte, and Transformers reloads the output natively as + NemotronHForCausalLM. + + megatron_to_hf_gpu: + status: verified + precision: bf16 + command: > + ./scripts/conversion/convert.sh export --executor slurm --device gpu + --nodes 1 --gpus-per-node 1 --tp 1 + --hf-model nvidia/NVIDIA-Nemotron-3-Nano-4B-BF16 + --hf-revision dfaf35de3e30f1867dd8dbc38a7fc9fb52d3914f + --megatron-path work/model-verification/nemotron-3-nano-4b/imported-megatron + --hf-path work/model-verification/nemotron-3-nano-4b/hf-export + --torch-dtype bfloat16 + last_verified: 2026-07-20 + expected_result: > + Strict export exits successfully; all 263 BF16 tensors match the pinned + source bitwise, the generation configuration is exact, and Transformers + reloads the output natively as NemotronHForCausalLM. A_log, D, dt_bias, + and out_proj are exact for all 21 Mamba layers. + + manual_forward_pass: + status: verified + precision: bf16 + command: > + uv run python -m torch.distributed.run --standalone --nproc_per_node=1 + examples/conversion/compare_hf_and_megatron/compare.py + --hf_model_path nvidia/NVIDIA-Nemotron-3-Nano-4B-BF16 + --hf-revision dfaf35de3e30f1867dd8dbc38a7fc9fb52d3914f + --megatron_model_path work/model-verification/nemotron-3-nano-4b/imported-megatron/iter_0000000 + --tp 1 --prompt "The capital of France is the city of" + last_verified: 2026-07-20 + expected_result: > + The eight-token one-step comparison exits successfully using native + Transformers Nemotron-H and a real Megatron inference context. Both paths + have next-token predictions that match at token ID 6993 (" Paris"); cosine + similarity is 0.999982 and the maximum and mean absolute logit + differences are 0.093750 and 0.018522. + + inference: + status: verified + precision: bf16 + bridge_commit: 78ecfad7644ee67efa4e8af4b0e7cf09f6bc9ecc # pragma: allowlist secret + command: > + uv run python -m torch.distributed.run --standalone --nproc_per_node=1 + examples/conversion/hf_to_megatron_generate_text.py + --hf_model_path nvidia/NVIDIA-Nemotron-3-Nano-4B-BF16 + --hf-revision dfaf35de3e30f1867dd8dbc38a7fc9fb52d3914f + --megatron_model_path work/model-verification/nemotron-3-nano-4b/imported-megatron/iter_0000000 + --tp 1 --prompt "The capital of France is" --max_new_tokens 2 + last_verified: 2026-07-20 + expected_result: > + Two independent runs use greedy decoding, produce byte-identical token + IDs, and return exactly 2 generated tokens with this literal completion, + including its leading space: ' Paris."'. + + pretrain: + status: verified + precision: bf16 + gpu_type: H100 + enabled_features: {} + command: > + ./scripts/training/train.sh --wait --nodes 1 --gpus-per-node 8 + --recipe nemotron_3_nano_4b_pretrain_8gpu_h100_bf16_config --mode pretrain + --dataset megatron-indexed --seq_length 4096 --max_steps 100 + --lr 3e-4 --min_lr 3e-5 --warmup_iters 40 + 'dataset.blend=[["work/data/rp2/head_01"],null]' + dataset.path_to_cache=work/cache/nemotron-3-nano-4b/rp2 + tokenizer.tokenizer_type=SentencePieceTokenizer + tokenizer.tokenizer_model=work/data/rp2/tokenizer.model + rng.seed=1234 dataset.random_seed=1234 + scheduler.lr_decay_iters=100 validation.eval_iters=0 validation.eval_interval=0 + dist.distributed_timeout_minutes=30 + ddp.check_for_nan_in_grad=true ddp.check_for_large_grads=true + rerun_state_machine.check_for_nan_in_loss=true checkpoint.load=null + --save_dir work/model-verification/nemotron-3-nano-4b/pretrain-reference-checkpoints + --save_interval 50 logger.log_interval=1 logger.log_throughput=true + logger.tensorboard_dir=null + last_verified: 2026-07-20 + metrics: + initial_loss: 12.432500 + final_loss: 5.511951 + last_10_steps_step_time_ms_avg: 26775.11 + last_10_steps_model_tflops_per_gpu_avg: 402.51 + expected_result: > + The uninterrupted bounded random-initialization RP2 reference finishes + all 100 steps at recipe-owned GBS/MBS 1024/1 with the four recorded + metrics, finite loss, and no skipped or NaN iterations. It reaches peak + learning rate 3e-4 at step 40 and completes cosine decay to 3e-5 at step + 100. Both iter_0000050 and iter_0000100 contain all eight distributed + shards plus metadata, optimizer/RNG train state, run config, and + tokenizer; iter_0000100 reloads at step 100 without an extra optimizer + step. + + sft: + status: verified + precision: bf16 + gpu_type: H100 + enabled_features: + sequence_packing: offline + command: > + ./scripts/training/train.sh --wait --nodes 1 --gpus-per-node 8 + --recipe nemotron_3_nano_4b_sft_8gpu_h100_bf16_config --mode sft + --dataset tulu3 + --pretrained_checkpoint work/model-verification/nemotron-3-nano-4b/imported-megatron/iter_0000000 + --max_steps 100 --seq_length 2048 + --lr 5e-6 --min_lr 0 --warmup_iters 10 + 'dataset.hf_dataset.split="train[:10000]"' + 'dataset.hf_dataset.load_kwargs={revision:"b14afda60f1bbebe55d5d2fa1e4df5042f97f8be"}' + dataset.hf_output_root=work/data/tulu3/nemotron-3-nano-4b-sft-b14afda60f1b + dataset.hf_rewrite=true dataset.seed=1234 rng.seed=5678 + dataset.do_validation=false dataset.hf_validation_proportion=null + dataset.enable_offline_packing=true + 'dataset.offline_packing_specs={packed_sequence_size:2048,pad_seq_to_mult:1}' + scheduler.lr_decay_iters=100 validation.eval_iters=0 validation.eval_interval=0 + checkpoint.load=null + --save_dir work/model-verification/nemotron-3-nano-4b/sft-checkpoints + --save_interval 100 logger.log_interval=1 logger.log_throughput=true + last_verified: 2026-07-20 + metrics: + initial_loss: 1.839282 + final_loss: 1.234906 + last_10_steps_step_time_ms_avg: 627.0 + last_10_steps_model_tflops_per_gpu_avg: 262.89 + expected_result: > + Pad-1 offline packing is 99.27% efficient. The TP1/DP8 run uses four + microbatches per optimizer step and completes all 100 steps with the four + recorded metrics, finite loss, and no skipped or NaN iterations. Across + 6,553,600 token slots, assistant-only loss masks contain 4,187,630 + supervised tokens. The complete eight-shard iter_0000100 full-model + checkpoint is saved and reloads for export. + + sft_export_inference: + status: verified + precision: bf16 + depends_on: sft + commands: + - > + ./scripts/conversion/convert.sh export --executor slurm --device gpu + --nodes 1 --gpus-per-node 1 --tp 1 + --hf-model nvidia/NVIDIA-Nemotron-3-Nano-4B-BF16 + --hf-revision dfaf35de3e30f1867dd8dbc38a7fc9fb52d3914f + --megatron-path work/model-verification/nemotron-3-nano-4b/sft-checkpoints/iter_0000100 + --hf-path work/model-verification/nemotron-3-nano-4b/sft-hf-export + --torch-dtype bfloat16 --export-weight-dtype bfloat16 + - > + uv run python + skills/create-model-verification-card/scripts/verify_hf_inference.py + --hf-model work/model-verification/nemotron-3-nano-4b/sft-hf-export + --prompt "Name the capital of France and explain its role in one sentence." + --max-new-tokens 45 --chat-template --disable-thinking + last_verified: 2026-07-20 + expected_result: > + The final full-SFT checkpoint exports as one BF16 safetensors file with + 263 weights and reloads natively as NemotronHForCausalLM. The helper runs + greedy generation twice, obtains identical token IDs and exactly 45 new + tokens, and produces this literal completion with twelve trailing newline + characters: "The capital of France is Paris, which serves as the country's + political, cultural, and economic center.\n\n\n\n\n\n\n\n\n\n\n\n". + + sft_long_context: + status: verified + precision: bf16 + gpu_type: H100 + enabled_features: + sequence_packing: offline + context_parallel_size: 2 + command: > + ./scripts/training/train.sh --wait --nodes 1 --gpus-per-node 8 + --recipe nemotron_3_nano_4b_sft_8gpu_h100_bf16_32k_config --mode sft + --dataset tulu3 + --pretrained_checkpoint work/model-verification/nemotron-3-nano-4b/imported-megatron/iter_0000000 + --max_steps 100 --seq_length 32768 + --tensor_model_parallel_size 2 --context_parallel_size 2 + --lr 5e-6 --min_lr 0 --warmup_iters 10 + 'dataset.hf_dataset.split="train[:10000]"' + 'dataset.hf_dataset.load_kwargs={revision:"b14afda60f1bbebe55d5d2fa1e4df5042f97f8be"}' + dataset.hf_output_root=work/data/tulu3/nemotron-3-nano-4b-long-b14afda60f1b + dataset.hf_rewrite=true dataset.seed=1234 rng.seed=5678 + dataset.do_validation=false dataset.hf_validation_proportion=null + dataset.enable_offline_packing=true + 'dataset.offline_packing_specs={packed_sequence_size:32768,pad_seq_to_mult:4}' + model.sequence_parallel=true model.cp_comm_type=a2a + model.cross_entropy_loss_fusion=false model.calculate_per_token_loss=true + ddp.average_in_collective=false scheduler.lr_decay_iters=100 + validation.eval_iters=0 validation.eval_interval=0 checkpoint.load=null + --save_dir work/model-verification/nemotron-3-nano-4b/long-sft-checkpoints + --save_interval 100 logger.log_interval=1 logger.log_throughput=true + last_verified: 2026-07-20 + metrics: + initial_loss: 1.658660 + final_loss: 1.100197 + last_10_steps_step_time_ms_avg: 5575.59 + last_10_steps_model_tflops_per_gpu_avg: 142.06 + expected_result: > + Pad-4 offline packing is 99.35% efficient. The TP2/CP2/SP-on/DP2 run uses + four microbatches per optimizer step and completes all 100 32K steps with + the four recorded metrics, finite loss, and no skipped or NaN iterations. + Across 26,214,400 token slots, assistant-only loss masks contain + 16,734,347 supervised tokens. The complete eight-shard iter_0000100 + checkpoint is saved and reloads at the same topology. + + peft: + status: verified + precision: bf16 + gpu_type: H100 + enabled_features: + sequence_packing: offline + command: > + ./scripts/training/train.sh --wait --nodes 1 --gpus-per-node 8 + --recipe nemotron_3_nano_4b_peft_8gpu_h100_bf16_config --mode lora + --dataset tulu3 + --pretrained_checkpoint work/model-verification/nemotron-3-nano-4b/imported-megatron/iter_0000000 + --max_steps 100 --seq_length 2048 + --lr 1e-4 --min_lr 0 --warmup_iters 10 + 'dataset.hf_dataset.split="train[:10000]"' + 'dataset.hf_dataset.load_kwargs={revision:"b14afda60f1bbebe55d5d2fa1e4df5042f97f8be"}' + dataset.hf_output_root=work/data/tulu3/nemotron-3-nano-4b-peft-b14afda60f1b + dataset.hf_rewrite=true dataset.seed=1234 rng.seed=5678 + dataset.do_validation=false dataset.hf_validation_proportion=null + dataset.enable_offline_packing=true + 'dataset.offline_packing_specs={packed_sequence_size:2048,pad_seq_to_mult:4}' + scheduler.lr_decay_iters=100 validation.eval_iters=0 validation.eval_interval=0 + checkpoint.load=null + --save_dir work/model-verification/nemotron-3-nano-4b/peft-checkpoints + --save_interval 100 logger.log_interval=1 logger.log_throughput=true + last_verified: 2026-07-20 + metrics: + initial_loss: 1.769892 + final_loss: 1.256916 + last_10_steps_step_time_ms_avg: 381.69 + last_10_steps_model_tflops_per_gpu_avg: 431.90 + expected_result: > + Pad-4 offline packing is 99.36% efficient. The frozen-base LoRA run uses + rank 8, alpha 16, zero dropout, and linear_qkv/linear_proj targets. Its + TP1/DP8 layout uses four microbatches per optimizer step and completes all + 100 steps with finite loss and no skipped or NaN iterations. Across + 6,553,600 token slots, assistant-only loss masks contain 4,178,503 + supervised tokens. The complete eight-shard iter_0000100 adapter + checkpoint reloads over the pinned base model at step 100. + + checkpoint_resume: + status: verified + precision: bf16 + gpu_type: H100 + depends_on: pretrain + command: > + ./scripts/training/train.sh --wait --nodes 1 --gpus-per-node 8 + --recipe nemotron_3_nano_4b_pretrain_8gpu_h100_bf16_config --mode pretrain + --dataset megatron-indexed --seq_length 4096 --max_steps 100 + --lr 3e-4 --min_lr 3e-5 --warmup_iters 40 + 'dataset.blend=[["work/data/rp2/head_01"],null]' + dataset.path_to_cache=work/cache/nemotron-3-nano-4b/rp2 + tokenizer.tokenizer_type=SentencePieceTokenizer + tokenizer.tokenizer_model=work/data/rp2/tokenizer.model + rng.seed=1234 dataset.random_seed=1234 + scheduler.lr_decay_iters=100 validation.eval_iters=0 validation.eval_interval=0 + dist.distributed_timeout_minutes=30 + ddp.check_for_nan_in_grad=true ddp.check_for_large_grads=true + rerun_state_machine.check_for_nan_in_loss=true + --load_dir work/model-verification/nemotron-3-nano-4b/pretrain-reference-checkpoints + checkpoint.ckpt_step=50 + --save_dir work/model-verification/nemotron-3-nano-4b/pretrain-resumed-checkpoints + --save_interval 50 logger.log_interval=1 logger.log_throughput=true + logger.tensorboard_dir=null + last_verified: 2026-07-20 + metrics: + initial_loss: 6.173212 + final_loss: 5.512038 + last_10_steps_step_time_ms_avg: 26943.07 + last_10_steps_model_tflops_per_gpu_avg: 400.11 + resume_comparison: + reference_item: pretrain + sentinel_steps: [51, 100] + loss_relative_tolerance: 1.0e-2 + loss_absolute_tolerance: 1.0e-6 + sentinels_match: true + expected_result: > + The continuation restores optimizer, scheduler, data-order, and RNG state + directly from iter_0000050 and executes exactly steps 51 through 100 into + a distinct output root with the four recorded metrics, finite loss, and + no skipped or NaN iterations. Step 51 matches the uninterrupted reference + exactly at 6.173212. At step 100, resumed loss 5.512038 differs from + reference loss 5.511951 by 0.000087 (0.00158%), inside the declared 1% + tolerance. Its complete eight-shard iter_0000100 checkpoint contains + metadata, optimizer/RNG train state, run config, and tokenizer, and + reloads at step 100 without an extra optimizer step. diff --git a/model_cards/qwen3-30b-a3b/card.yaml b/model_cards/qwen3-30b-a3b/card.yaml new file mode 100644 index 0000000000..3b9118d788 --- /dev/null +++ b/model_cards/qwen3-30b-a3b/card.yaml @@ -0,0 +1,387 @@ +# Agent-readable model support card. +# status: unverified | verified | unsupported | not_applicable + +title: qwen3_30b_a3b +model: + hf_id: Qwen/Qwen3-30B-A3B + hf_revision: ad44e777bcd18fa416d9da3bd8f70d33ebb85d39 # pragma: allowlist secret + architecture: Qwen3MoeForCausalLM + min_transformers_version: "5.8.1" +verification_environment: + base_container: nvcr.io/nvidia/pytorch:26.04-py3 + bridge_commit: aabb29d0ed15fc7ed881f0538352c6bd2d2fd52d # pragma: allowlist secret +summary: Qwen3-30B-A3B support verification for conversion, inference, and training. + +items: + hf_to_megatron_cpu: + status: verified + precision: bf16 + command: > + ./scripts/conversion/convert.sh import --executor slurm --device cpu --nodes 1 + --hf-model Qwen/Qwen3-30B-A3B + --megatron-path work/model-verification/qwen3-30b-a3b/cpu-megatron + --torch-dtype bfloat16 + last_verified: 2026-07-17 + expected_result: > + The command exits successfully, creates iter_0000000, and the checkpoint + round-trips through CPU export with all 18,867 BF16 tensors matching the + recorded HF revision bitwise. + + hf_to_megatron_gpu: + status: verified + precision: bf16 + command: > + ./scripts/conversion/convert.sh import --executor slurm --device gpu + --nodes 1 --gpus-per-node 8 + --hf-model Qwen/Qwen3-30B-A3B + --megatron-path work/model-verification/qwen3-30b-a3b/imported-megatron + --torch-dtype bfloat16 --tp 4 --pp 2 --ep 4 + last_verified: 2026-07-17 + expected_result: > + The command exits successfully, creates iter_0000000, and all 18,867 BF16 + tensors reload with keys, shapes, dtypes, and values exactly matching the + recorded HF revision. + + megatron_to_hf_cpu: + status: verified + precision: bf16 + command: > + ./scripts/conversion/convert.sh export --executor slurm --device cpu --nodes 1 + --hf-model Qwen/Qwen3-30B-A3B + --megatron-path work/model-verification/qwen3-30b-a3b/cpu-megatron/iter_0000000 + --hf-path work/model-verification/qwen3-30b-a3b/cpu-hf-export + --torch-dtype bfloat16 + last_verified: 2026-07-17 + expected_result: > + The command exits successfully; all 18,867 exported BF16 tensors match + the recorded HF revision bitwise, norm_topk_prob remains true, and the + export reloads as Qwen3MoeForCausalLM without missing or unexpected keys. + + megatron_to_hf_gpu: + status: verified + precision: bf16 + command: > + ./scripts/conversion/convert.sh export --executor slurm --device gpu + --nodes 1 --gpus-per-node 8 + --hf-model Qwen/Qwen3-30B-A3B + --megatron-path work/model-verification/qwen3-30b-a3b/imported-megatron/iter_0000000 + --hf-path work/model-verification/qwen3-30b-a3b/hf-export + --torch-dtype bfloat16 --export-weight-dtype bfloat16 + --distributed-save --tp 4 --pp 2 --ep 4 + last_verified: 2026-07-17 + expected_result: > + Strict export exits successfully, all 18,867 written BF16 tensors match + the recorded HF revision bitwise, norm_topk_prob remains true, and + Transformers reloads the output as Qwen3MoeForCausalLM without missing or + unexpected keys. + + manual_forward_pass: + status: verified + precision: bf16 + bridge_commit: 24f3dcde1fe27a1c3f1097b94d69b38e37a57415 # pragma: allowlist secret + command: > + uv run python -m torch.distributed.run --standalone --nproc_per_node=8 + examples/conversion/compare_hf_and_megatron/compare.py + --hf_model_path Qwen/Qwen3-30B-A3B + --hf-revision ad44e777bcd18fa416d9da3bd8f70d33ebb85d39 + --megatron_model_path work/model-verification/qwen3-30b-a3b/imported-megatron/iter_0000000 + --tp 1 --pp 8 --prompt "The capital of France is the city of" + last_verified: 2026-07-20 + expected_result: > + The pinned-revision one-step comparison exits successfully; the Hugging + Face and Megatron next-token predictions match at token ID 12095 + (" Paris"), and cosine similarity is 0.999317, above the 0.99 correlation + gate. The maximum and mean absolute logit differences are 0.593750 and + 0.130064, respectively; both are report-only diagnostic observations. + + inference: + status: verified + precision: bf16 + command: > + uv run python -m torch.distributed.run --standalone --nproc_per_node=8 + examples/conversion/hf_to_megatron_generate_text.py + --hf_model_path Qwen/Qwen3-30B-A3B + --megatron_model_path work/model-verification/qwen3-30b-a3b/imported-megatron/iter_0000000 + --tp 4 --pp 2 --ep 4 + --prompt "The capital of France is" --max_new_tokens 32 + last_verified: 2026-07-17 + expected_result: > + Two independent executions exit successfully after exactly 32 new tokens + and print this byte-identical completion, including its leading + space, " Paris. The capital of the United Kingdom is London. The capital + of the United States is Washington, D.C. The capital of Brazil is + Brasília. The". + + pretrain: + status: verified + precision: bf16 + bridge_commit: 5b9d9cf501193277e1ca47a99b97c640b8f39f90 # pragma: allowlist secret + gpu_type: H100 + enabled_features: + cuda_graph: + implementation: transformer_engine + scopes: [moe_router, moe_preprocess] + moe_dispatcher: hybridep + command: > + ./scripts/training/train.sh --nodes 2 --gpus-per-node 8 + --recipe qwen3_30b_a3b_pretrain_config --mode pretrain + --dataset megatron-indexed --seq_length 4096 --max_steps 100 + --lr 3e-4 --min_lr 3e-5 --warmup_iters 40 + 'dataset.blend=[["work/data/rp2/head_01"],null]' + dataset.path_to_cache=work/cache/qwen3-30b-a3b/rp2 + tokenizer.tokenizer_type=SentencePieceTokenizer + tokenizer.tokenizer_model=work/data/rp2/tokenizer.model + scheduler.lr_decay_iters=100 + model.moe_router_force_load_balancing=false + ddp.check_for_nan_in_grad=true ddp.check_for_large_grads=true + rerun_state_machine.check_for_nan_in_loss=true checkpoint.load=null + validation.eval_iters=0 validation.eval_interval=0 + dataset.random_seed=1234 dataset.num_workers=8 rng.seed=1234 + dist.distributed_timeout_minutes=30 + --save_dir work/model-verification/qwen3-30b-a3b/pretrain-convergence-v1-reference-checkpoints + --save_interval 50 logger.log_interval=1 logger.log_throughput=true + logger.tensorboard_dir=null + last_verified: 2026-07-19 + metrics: + initial_loss: 12.41145 + final_loss: 6.139116 + last_10_steps_step_time_ms_avg: 30289.550 + last_10_steps_model_tflops_per_gpu_avg: 199.120 + expected_result: > + On 16x H100, the public alias resolves to the 16-GPU recipe and completes + exactly 100 bounded RP2 optimizer steps with TP1/PP1/CP1/EP16/ETP1, + DP16, SP off, GBS/MBS 1024/1, and 64-way gradient accumulation. Natural + routing, HybridEP, and Transformer Engine CUDA graphs for moe_router and + moe_preprocess remain active. Loss is finite from 12.41145 to 6.139116 + with no skipped or NaN iterations, all four metrics are recorded, and + complete iter_0000050 and iter_0000100 checkpoints are saved. + + sft: + status: unverified + precision: bf16 + gpu_type: H100 + enabled_features: + sequence_packing: offline + command: > + ./scripts/training/train.sh --nodes 2 --gpus-per-node 8 + --recipe qwen3_30b_a3b_sft_config --mode sft + --dataset tulu3 + --pretrained_checkpoint work/model-verification/qwen3-30b-a3b/imported-megatron/iter_0000000 + --max_steps 100 --seq_length 2048 + --lr 5e-6 --min_lr 0 --warmup_iters 10 + 'dataset.hf_dataset.split="train[:10000]"' + 'dataset.hf_dataset.load_kwargs={revision:"b14afda60f1bbebe55d5d2fa1e4df5042f97f8be"}' + '++tokenizer.hf_tokenizer_kwargs.revision="ad44e777bcd18fa416d9da3bd8f70d33ebb85d39"' + dataset.hf_output_root=work/data/tulu3/qwen3-30b-a3b-sft-b14afda60f1b + dataset.hf_rewrite=true dataset.seed=1234 rng.seed=5678 + dataset.do_validation=false dataset.hf_validation_proportion=null + dataset.enable_offline_packing=true + +dataset.offline_packing_specs.pad_seq_to_mult=1 scheduler.lr_decay_iters=100 + validation.eval_iters=0 validation.eval_interval=0 + model.moe_router_force_load_balancing=false checkpoint.load=null + ddp.check_for_nan_in_grad=true ddp.check_for_large_grads=true + rerun_state_machine.check_for_nan_in_loss=true + --save_dir work/model-verification/qwen3-30b-a3b/sft-convergence-v1-checkpoints + --save_interval 100 + logger.log_interval=1 logger.log_throughput=true logger.tensorboard_dir=null + last_verified: null + metrics: + initial_loss: null + final_loss: null + last_10_steps_step_time_ms_avg: null + last_10_steps_model_tflops_per_gpu_avg: null + expected_result: > + The pending 16x H100 run must use the immutable Tulu3 and tokenizer + revisions, fresh pad-1 offline packing, dataset seed 1234, model RNG seed + 5678, and recipe-owned TP1/PP1/CP1/EP16/SP-off topology. Exactly 100 + full-SFT steps must complete at GBS/MBS 32/1 with natural routing, finite + losses, no skipped or NaN iterations, all four metrics, packing + efficiency, the actual supervised-token count, and a complete reloadable + iter_0000100 checkpoint before this item can be verified. + + sft_export_inference: + status: unverified + precision: bf16 + depends_on: sft + commands: + - > + ./scripts/conversion/convert.sh export --executor slurm --device gpu + --nodes 2 --gpus-per-node 8 --hf-model Qwen/Qwen3-30B-A3B + --megatron-path work/model-verification/qwen3-30b-a3b/sft-convergence-v1-checkpoints/iter_0000100 + --hf-path work/model-verification/qwen3-30b-a3b/sft-convergence-v1-hf-export + --torch-dtype bfloat16 --export-weight-dtype bfloat16 + --distributed-save --tp 1 --pp 1 --ep 16 --etp 1 + - > + uv run python + skills/create-model-verification-card/scripts/verify_hf_inference.py + --hf-model work/model-verification/qwen3-30b-a3b/sft-convergence-v1-hf-export + --prompt "In one short sentence, explain why Paris is important to France." + --max-new-tokens 32 --chat-template --disable-thinking + last_verified: null + expected_result: > + After the pending immutable-revision SFT run passes, its step-100 + checkpoint must export into indexed BF16 shards whose index covers every + weight, reload with Transformers without missing or unexpected keys, and + produce exactly 32 new tokens in two independent greedy runs with + byte-identical token IDs. The exact literal completion must be recorded + before verification. + + sft_long_context: + status: verified + precision: bf16 + bridge_commit: f3ae2767b5e18aeb67b726cd8d5f1db58216dcc9 # pragma: allowlist secret + gpu_type: H100 + enabled_features: + sequence_packing: offline + context_parallel_size: 2 + moe_dispatcher: deepep + command: > + ./scripts/training/train.sh --nodes 2 --gpus-per-node 8 + --recipe qwen3_30b_a3b_sft_8gpu_h100_bf16_config --mode sft + --dataset tulu3 + --pretrained_checkpoint work/model-verification/qwen3-30b-a3b/imported-megatron/iter_0000000 + --max_steps 20 --seq_length 32768 --context_parallel_size 2 + -tp 8 -pp 1 -ep 8 + --lr 1e-6 --min_lr 0 --warmup_iters 2 + 'dataset.hf_dataset.split="train[:10000]"' + 'dataset.hf_dataset.load_kwargs={revision:"b14afda60f1bbebe55d5d2fa1e4df5042f97f8be"}' + '++tokenizer.hf_tokenizer_kwargs.revision="ad44e777bcd18fa416d9da3bd8f70d33ebb85d39"' + dataset.hf_output_root=work/data/tulu3/qwen3-30b-a3b-long-context-b14afda60f1b + dataset.hf_rewrite=true dataset.seed=1234 rng.seed=5678 + dataset.do_validation=false dataset.hf_validation_proportion=null + dataset.enable_offline_packing=true + +dataset.offline_packing_specs.pad_seq_to_mult=16 model.cp_comm_type=p2p + model.cross_entropy_loss_fusion=false model.recompute_granularity=full + model.recompute_method=uniform model.recompute_num_layers=1 + scheduler.lr_decay_iters=20 validation.eval_iters=0 validation.eval_interval=0 + checkpoint.load=null checkpoint.save=null + logger.log_interval=1 logger.log_throughput=true + last_verified: 2026-07-20 + metrics: + initial_loss: 1.645009 + final_loss: 1.468103 + last_10_steps_step_time_ms_avg: 142663.710 + last_10_steps_model_tflops_per_gpu_avg: 26.120 + expected_result: > + The immutable-revision 16-GPU run completes exactly 20 Tulu3 SFT steps at + sequence length 32768 with TP8/PP1/CP2/EP8/SP-on, DeepEP, and explicit + pad-16 offline packing. LM loss is 1.645009 to 1.468103; skipped/NaN + totals are 0/0. The persisted post-setup runtime config matches the + command, packing is 99.28%, and the sampled training window contains + 13,573,663 actual supervised tokens. PP=1 keeps tokens, labels, loss + masks, and packed-sequence boundaries on one pipeline stage. + + peft: + status: unverified + precision: bf16 + gpu_type: H100 + enabled_features: + sequence_packing: offline + moe_dispatcher: deepep + command: > + ./scripts/training/train.sh --nodes 1 --gpus-per-node 4 + --recipe qwen3_30b_a3b_peft_4gpu_h100_bf16_config --mode lora + --dataset tulu3 + --pretrained_checkpoint work/model-verification/qwen3-30b-a3b/imported-megatron/iter_0000000 + --max_steps 100 --seq_length 2048 + --lr 1e-4 --min_lr 0 --warmup_iters 10 + 'dataset.hf_dataset.split="train[:10000]"' + 'dataset.hf_dataset.load_kwargs={revision:"b14afda60f1bbebe55d5d2fa1e4df5042f97f8be"}' + '++tokenizer.hf_tokenizer_kwargs.revision="ad44e777bcd18fa416d9da3bd8f70d33ebb85d39"' + dataset.hf_output_root=work/data/tulu3/qwen3-30b-a3b-peft-b14afda60f1b + dataset.hf_rewrite=true dataset.seed=1234 rng.seed=5678 + dataset.do_validation=false dataset.hf_validation_proportion=null + dataset.enable_offline_packing=true + +dataset.offline_packing_specs.pad_seq_to_mult=4 scheduler.lr_decay_iters=100 + validation.eval_iters=0 validation.eval_interval=0 + model.moe_router_force_load_balancing=false checkpoint.load=null + ddp.check_for_nan_in_grad=true ddp.check_for_large_grads=true + rerun_state_machine.check_for_nan_in_loss=true + --save_dir work/model-verification/qwen3-30b-a3b/peft-tp4-ep4-pad4-checkpoints + --save_interval 100 + logger.log_interval=1 logger.log_throughput=true + last_verified: null + metrics: + initial_loss: null + final_loss: null + last_10_steps_step_time_ms_avg: null + last_10_steps_model_tflops_per_gpu_avg: null + expected_result: > + The pending 4x H100 run must use the immutable Tulu3 and tokenizer + revisions, fresh pad-4 offline packing, dataset seed 1234, model RNG seed + 5678, and recipe-owned TP4/PP1/EP4/SP-on topology. Exactly 100 LoRA steps + must complete with natural routing, finite losses, no skipped or NaN + iterations, and all four metrics recorded before a complete + iter_0000100 adapter checkpoint can be verified. + + checkpoint_resume: + status: verified + precision: bf16 + bridge_commit: 5b9d9cf501193277e1ca47a99b97c640b8f39f90 # pragma: allowlist secret + gpu_type: H100 + depends_on: pretrain + command: > + ./scripts/training/train.sh --nodes 2 --gpus-per-node 8 + --recipe qwen3_30b_a3b_pretrain_config --mode pretrain + --dataset megatron-indexed --seq_length 4096 --max_steps 100 + --lr 3e-4 --min_lr 3e-5 --warmup_iters 40 + 'dataset.blend=[["work/data/rp2/head_01"],null]' + dataset.path_to_cache=work/cache/qwen3-30b-a3b/rp2 + tokenizer.tokenizer_type=SentencePieceTokenizer + tokenizer.tokenizer_model=work/data/rp2/tokenizer.model + scheduler.lr_decay_iters=100 + model.moe_router_force_load_balancing=false + ddp.check_for_nan_in_grad=true ddp.check_for_large_grads=true + rerun_state_machine.check_for_nan_in_loss=true + validation.eval_iters=0 validation.eval_interval=0 + dataset.random_seed=1234 dataset.num_workers=8 rng.seed=1234 + dist.distributed_timeout_minutes=30 + --load_dir work/model-verification/qwen3-30b-a3b/pretrain-convergence-v1-reference-checkpoints + --save_dir work/model-verification/qwen3-30b-a3b/pretrain-convergence-v1-resumed-checkpoints + --save_interval 50 checkpoint.ckpt_step=50 + train.empty_unused_memory_level=2 + logger.log_interval=1 logger.log_throughput=true + logger.tensorboard_dir=null + last_verified: 2026-07-19 + metrics: + initial_loss: 6.989006 + final_loss: 6.145390 + last_10_steps_step_time_ms_avg: 31142.650 + last_10_steps_model_tflops_per_gpu_avg: 193.640 + resume_comparison: + reference_item: pretrain + sentinel_steps: [51, 100] + loss_relative_tolerance: 1.0e-2 + loss_absolute_tolerance: 1.0e-6 + sentinels_match: true + expected_result: > + The command restores optimizer, scheduler, data-order, and RNG state from + iter_0000050, begins at step 51, and finishes at step 100 in the distinct + resumed root with finite losses and no skipped or NaN iterations. Releasing + unused cached memory after optimizer steps is execution-only. Step-51 loss + 6.989006 matches the uninterrupted reference exactly; step-100 loss + 6.145390 differs from reference 6.139116 by 0.102197%, within the declared + one-percent gate, so both sentinels match and all four metrics are recorded. + + pretrain_performance: + status: verified + precision: bf16 + bridge_commit: 7160fd35c8f01f09d795e10f2057259d122c1e7e # pragma: allowlist secret + gpu_type: H100 + command: > + ./scripts/training/train.sh --nodes 2 --gpus-per-node 8 + --recipe qwen3_30b_a3b_pretrain_16gpu_h100_bf16_config --max_steps 50 + last_verified: 2026-07-18 + metrics: + initial_loss: 12.34643 + final_loss: 8.143753 + last_10_steps_step_time_ms_avg: 25118.160 + last_10_steps_model_tflops_per_gpu_avg: 240.110 + expected_result: > + On two nodes with 16x H100, the exact mock-data performance recipe + completes exactly 50 steps with finite losses, no skipped or NaN + iterations, and all four metrics recorded. HybridEP is active and + Transformer Engine CUDA graph capture completes for all 48 graphable + layers with moe_router and moe_preprocess scopes. Over steps 41-50, the + run averages at most 27 seconds per step and at least 225 model + TFLOP/s/GPU, while peak allocated memory remains below 65 GiB/GPU. diff --git a/model_cards/qwen3-8b/card.yaml b/model_cards/qwen3-8b/card.yaml new file mode 100644 index 0000000000..62373d8a46 --- /dev/null +++ b/model_cards/qwen3-8b/card.yaml @@ -0,0 +1,338 @@ +# Draft agent-readable model support card. +# status: unverified | verified | unsupported | not_applicable +# last_verified: YYYY-MM-DD or null +# verified requires an immutable hf_revision, successful command(s), a concrete +# expected_result, and non-null last_verified. +# unverified may retain an older last_verified date when validation is stale. +# Every training item records the GPU type used for verification. + +title: qwen3_8b +model: + hf_id: Qwen/Qwen3-8B + hf_revision: b968826d9c46dd6066d109eabc6255188de91218 # pragma: allowlist secret + architecture: Qwen3ForCausalLM + min_transformers_version: "5.8.1" +verification_environment: + base_container: nvcr.io/nvidia/pytorch:26.04-py3 + bridge_commit: 5c56eab34c540fad08544c38b7cc39d662fb7475 # pragma: allowlist secret +summary: Qwen3-8B support verification for conversion, inference, and training. + +items: + hf_to_megatron_cpu: + status: verified + precision: bf16 + command: > + ./scripts/conversion/convert.sh import --executor slurm --device cpu --nodes 1 + --hf-model Qwen/Qwen3-8B + --megatron-path work/model-verification/qwen3-8b/cpu-megatron + --torch-dtype bfloat16 + last_verified: 2026-07-16 + expected_result: > + The command exits successfully, creates iter_0000000, and the checkpoint + round-trips through CPU export with all 399 HF tensors matching the + recorded HF revision exactly in keys, shapes, dtypes, and values. + + hf_to_megatron_gpu: + status: verified + precision: bf16 + command: > + ./scripts/conversion/convert.sh import --executor slurm --device gpu + --nodes 1 --gpus-per-node 4 + --hf-model Qwen/Qwen3-8B + --megatron-path work/model-verification/qwen3-8b/imported-megatron + --torch-dtype bfloat16 --tp 4 + last_verified: 2026-07-16 + expected_result: > + The command exits successfully, creates iter_0000000, and the checkpoint + reloads at TP=4 with weights exactly matching the recorded HF revision. + + megatron_to_hf_cpu: + status: verified + precision: bf16 + command: > + ./scripts/conversion/convert.sh export --executor slurm --device cpu --nodes 1 + --hf-model Qwen/Qwen3-8B + --megatron-path work/model-verification/qwen3-8b/cpu-megatron/iter_0000000 + --hf-path work/model-verification/qwen3-8b/cpu-hf-export + last_verified: 2026-07-16 + expected_result: > + The command exits successfully; all 399 exported tensors match the + recorded HF revision exactly, and the export reloads on CPU as + Qwen3ForCausalLM. + + megatron_to_hf_gpu: + status: verified + precision: bf16 + command: > + ./scripts/conversion/convert.sh export --executor slurm --device gpu + --nodes 1 --gpus-per-node 4 + --hf-model Qwen/Qwen3-8B + --megatron-path work/model-verification/qwen3-8b/imported-megatron/iter_0000000 + --hf-path work/model-verification/qwen3-8b/hf-export --torch-dtype bfloat16 + --export-weight-dtype bfloat16 --distributed-save --tp 4 + last_verified: 2026-07-16 + expected_result: > + Strict export exits successfully and the Hugging Face output reloads with + AutoModelForCausalLM as Qwen3ForCausalLM. + + manual_forward_pass: + status: verified + precision: bf16 + command: > + uv run python -m torch.distributed.run --standalone --nproc_per_node=4 + examples/conversion/compare_hf_and_megatron/compare.py + --hf_model_path Qwen/Qwen3-8B + --megatron_model_path work/model-verification/qwen3-8b/imported-megatron/iter_0000000 + --tp 4 --prompt "The capital of France is the city of" + last_verified: 2026-07-17 + expected_result: > + The one-step comparison of the eight-token prompt exits successfully; the + Hugging Face and Megatron next-token predictions match at token ID 12095 + (" Paris"), and cosine similarity is 0.999969, above the 0.99 correlation + gate. This historical result predates explicit helper revision pinning + and is retained against the recorded immutable HF revision. The maximum + and mean absolute logit differences are 0.187500 and 0.030649, + respectively; both are report-only diagnostic observations. + + inference: + status: verified + precision: bf16 + command: > + uv run python -m torch.distributed.run --standalone --nproc_per_node=4 + examples/conversion/hf_to_megatron_generate_text.py + --hf_model_path Qwen/Qwen3-8B + --megatron_model_path work/model-verification/qwen3-8b/imported-megatron/iter_0000000 --tp 4 + --prompt "The capital of France is" --max_new_tokens 32 + last_verified: 2026-07-16 + expected_result: > + Two independent executions exit successfully after exactly 32 generation + steps and print this byte-identical completion, including its leading space, + " Paris. The capital of Italy is Rome. The capital of Spain is Madrid. The + capital of Germany is Berlin. The capital of the Netherlands is Amsterdam. The". + + pretrain: + status: unverified + precision: bf16 + gpu_type: H100 + enabled_features: {} + command: > + ./scripts/training/train.sh --nodes 2 --gpus-per-node 8 + --recipe qwen3_8b_pretrain_config --mode pretrain + --dataset megatron-indexed --seq_length 4096 + --max_steps 100 + --lr 3e-4 --min_lr 3e-5 --warmup_iters 40 + 'dataset.blend=[["work/data/rp2/head_01"],null]' + dataset.path_to_cache=work/cache/qwen3-8b/rp2 + tokenizer.tokenizer_type=SentencePieceTokenizer + tokenizer.tokenizer_model=work/data/rp2/tokenizer.model + rng.seed=1234 dataset.random_seed=1234 + scheduler.lr_decay_iters=100 validation.eval_iters=0 validation.eval_interval=0 + dist.distributed_timeout_minutes=30 + ddp.check_for_nan_in_grad=true ddp.check_for_large_grads=true + rerun_state_machine.check_for_nan_in_loss=true checkpoint.load=null + --save_dir work/model-verification/qwen3-8b/pretrain-reference-checkpoints --save_interval 50 + logger.log_interval=1 logger.log_throughput=true logger.tensorboard_dir=null + last_verified: null + metrics: + initial_loss: null + final_loss: null + last_10_steps_step_time_ms_avg: null + last_10_steps_model_tflops_per_gpu_avg: null + expected_result: > + The pending uninterrupted bounded RP2 reference must complete exactly 100 + steps at recipe-owned GBS/MBS 1024/1 with finite losses and no skipped or + NaN iterations. It must reach peak learning rate at step 40, complete + decay at step 100, and save complete iter_0000050 and iter_0000100 + checkpoints before this cohort item can be marked verified. + + sft: + status: verified + precision: bf16 + bridge_commit: f3ae2767b5e18aeb67b726cd8d5f1db58216dcc9 # pragma: allowlist secret + gpu_type: H100 + enabled_features: + sequence_packing: offline + command: > + ./scripts/training/train.sh --nodes 1 --gpus-per-node 4 + --recipe qwen3_8b_sft_4gpu_h100_bf16_config --mode sft + --dataset tulu3 + --pretrained_checkpoint work/model-verification/qwen3-8b/imported-megatron/iter_0000000 + --max_steps 100 --seq_length 2048 + --lr 5e-6 --min_lr 0 --warmup_iters 10 + 'dataset.hf_dataset.split="train[:10000]"' + 'dataset.hf_dataset.load_kwargs={revision:"b14afda60f1bbebe55d5d2fa1e4df5042f97f8be"}' + '++tokenizer.hf_tokenizer_kwargs.revision="b968826d9c46dd6066d109eabc6255188de91218"' + dataset.hf_output_root=work/data/tulu3/qwen3-8b-sft-b14afda60f1b + dataset.hf_rewrite=true dataset.seed=1234 rng.seed=5678 + dataset.do_validation=false dataset.hf_validation_proportion=null + dataset.enable_offline_packing=true + +dataset.offline_packing_specs.pad_seq_to_mult=1 scheduler.lr_decay_iters=100 + validation.eval_iters=0 validation.eval_interval=0 checkpoint.load=null + --save_dir work/model-verification/qwen3-8b/sft-checkpoints --save_interval 100 + logger.log_interval=1 logger.log_throughput=true + last_verified: 2026-07-19 + metrics: + initial_loss: 1.657044 + final_loss: 0.9733383 + last_10_steps_step_time_ms_avg: 7424.240 + last_10_steps_model_tflops_per_gpu_avg: 101.630 + expected_result: > + Immutable Tulu3 pad-1 offline packing is 99.30% efficient. Full SFT uses + DP=1 with 32 gradient-accumulation steps and reaches step 100 with the four + recorded metrics, finite loss, and no skipped or NaN iterations. Across + 6,553,600 token slots, the sampled assistant-only loss masks contain + 4,350,004 supervised tokens. The complete four-shard iter_0000100 + full-model checkpoint is saved. + + sft_export_inference: + status: verified + precision: bf16 + bridge_commit: f3ae2767b5e18aeb67b726cd8d5f1db58216dcc9 # pragma: allowlist secret + depends_on: sft + commands: + - > + ./scripts/conversion/convert.sh export --executor slurm --device gpu + --nodes 1 --gpus-per-node 4 --hf-model Qwen/Qwen3-8B + --megatron-path work/model-verification/qwen3-8b/sft-checkpoints/iter_0000100 + --hf-path work/model-verification/qwen3-8b/sft-hf-export + --torch-dtype bfloat16 --export-weight-dtype bfloat16 + --distributed-save --tp 4 + - > + uv run python + skills/create-model-verification-card/scripts/verify_hf_inference.py + --hf-model work/model-verification/qwen3-8b/sft-hf-export + --prompt "Name the capital of France and explain its role in one sentence." + --max-new-tokens 45 --chat-template --disable-thinking + last_verified: 2026-07-19 + expected_result: > + The current full-SFT checkpoint exports as five indexed BF16 shards with + 399 weights, config vocabulary size 151936, and embedding and output-head + shapes [151936, 4096]. Transformers natively reloads all 399 weights, and + two independent runs using greedy generation produce byte-identical token + IDs and exactly 45 new tokens with this literal completion, including the + escaped blank line: "The capital of France is Paris. It is the political, economic, and cultural center of France and serves as the seat of the French government. ✅\n\nParis is also a major tourist destination, known for its iconic landmarks". + + sft_long_context: + status: verified + precision: bf16 + bridge_commit: f3ae2767b5e18aeb67b726cd8d5f1db58216dcc9 # pragma: allowlist secret + gpu_type: H100 + enabled_features: + sequence_packing: offline + context_parallel_size: 2 + command: > + ./scripts/training/train.sh --nodes 1 --gpus-per-node 8 + --recipe qwen3_8b_sft_32k_config --mode sft + --dataset tulu3 + --pretrained_checkpoint work/model-verification/qwen3-8b/imported-megatron/iter_0000000 + --max_steps 20 --seq_length 32768 --context_parallel_size 2 + --lr 1e-6 --min_lr 0 --warmup_iters 2 + 'dataset.hf_dataset.split="train[:10000]"' + 'dataset.hf_dataset.load_kwargs={revision:"b14afda60f1bbebe55d5d2fa1e4df5042f97f8be"}' + '++tokenizer.hf_tokenizer_kwargs.revision="b968826d9c46dd6066d109eabc6255188de91218"' + dataset.hf_output_root=work/data/tulu3/qwen3-8b-long-context-b14afda60f1b + dataset.hf_rewrite=true dataset.seed=1234 rng.seed=5678 + dataset.do_validation=false dataset.hf_validation_proportion=null + dataset.enable_offline_packing=true + +dataset.offline_packing_specs.pad_seq_to_mult=8 + model.sequence_parallel=true + model.cp_comm_type=a2a model.cross_entropy_loss_fusion=false + scheduler.lr_decay_iters=20 validation.eval_iters=0 validation.eval_interval=0 + checkpoint.load=null checkpoint.save=null + logger.log_interval=1 logger.log_throughput=true + last_verified: 2026-07-20 + metrics: + initial_loss: 1.611622 + final_loss: 1.490618 + last_10_steps_step_time_ms_avg: 71668.780 + last_10_steps_model_tflops_per_gpu_avg: 34.080 + expected_result: > + The immutable-revision 8-GPU run completes exactly 20 Tulu3 SFT steps at + sequence length 32768 with recipe-owned TP4/PP1/CP2/SP-on, GBS/MBS 8/1, + and explicit pad-8 offline packing. LM loss is 1.611622 to 1.490618; + skipped/NaN totals are 0/0. The persisted post-setup runtime config + matches the command, packing is 99.97%, and the sampled training window + contains 3,444,917 actual supervised tokens. + + peft: + status: verified + precision: bf16 + bridge_commit: f3ae2767b5e18aeb67b726cd8d5f1db58216dcc9 # pragma: allowlist secret + gpu_type: H100 + enabled_features: + sequence_packing: offline + command: > + ./scripts/training/train.sh --nodes 1 --gpus-per-node 1 + --recipe qwen3_8b_peft_1gpu_h100_bf16_config --mode lora + --dataset tulu3 + --pretrained_checkpoint work/model-verification/qwen3-8b/imported-megatron/iter_0000000 + --max_steps 100 --seq_length 2048 + --lr 1e-4 --min_lr 0 --warmup_iters 10 + 'dataset.hf_dataset.split="train[:10000]"' + 'dataset.hf_dataset.load_kwargs={revision:"b14afda60f1bbebe55d5d2fa1e4df5042f97f8be"}' + '++tokenizer.hf_tokenizer_kwargs.revision="b968826d9c46dd6066d109eabc6255188de91218"' + dataset.hf_output_root=work/data/tulu3/qwen3-8b-peft-b14afda60f1b + dataset.hf_rewrite=true dataset.seed=1234 rng.seed=5678 + dataset.do_validation=false dataset.hf_validation_proportion=null + dataset.enable_offline_packing=true + +dataset.offline_packing_specs.pad_seq_to_mult=4 scheduler.lr_decay_iters=100 + validation.eval_iters=0 validation.eval_interval=0 checkpoint.load=null + --save_dir work/model-verification/qwen3-8b/peft-checkpoints --save_interval 100 + logger.log_interval=1 logger.log_throughput=true + last_verified: 2026-07-19 + metrics: + initial_loss: 1.527745 + final_loss: 1.170038 + last_10_steps_step_time_ms_avg: 7774.790 + last_10_steps_model_tflops_per_gpu_avg: 387.430 + expected_result: > + Immutable Tulu3 pad-4 offline packing is 99.39% efficient. The 100 LoRA + steps use DP=1 with 32 gradient-accumulation steps and complete with the + four recorded metrics, finite loss, and no skipped or NaN iterations. + Across 6,553,600 token slots, the sampled assistant-only loss masks + contain 4,332,480 supervised tokens. The complete single-shard + iter_0000100 adapter checkpoint is saved. + + checkpoint_resume: + status: unverified + precision: bf16 + gpu_type: H100 + depends_on: pretrain + command: > + ./scripts/training/train.sh --nodes 2 --gpus-per-node 8 + --recipe qwen3_8b_pretrain_config --mode pretrain + --dataset megatron-indexed --seq_length 4096 + --max_steps 100 + --lr 3e-4 --min_lr 3e-5 --warmup_iters 40 + 'dataset.blend=[["work/data/rp2/head_01"],null]' + dataset.path_to_cache=work/cache/qwen3-8b/rp2 + tokenizer.tokenizer_type=SentencePieceTokenizer + tokenizer.tokenizer_model=work/data/rp2/tokenizer.model + rng.seed=1234 dataset.random_seed=1234 + scheduler.lr_decay_iters=100 validation.eval_iters=0 validation.eval_interval=0 + dist.distributed_timeout_minutes=30 + ddp.check_for_nan_in_grad=true ddp.check_for_large_grads=true + rerun_state_machine.check_for_nan_in_loss=true + --load_dir work/model-verification/qwen3-8b/pretrain-reference-checkpoints + --save_dir work/model-verification/qwen3-8b/pretrain-resumed-from-reference-checkpoints + --save_interval 50 checkpoint.ckpt_step=50 + logger.log_interval=1 logger.log_throughput=true logger.tensorboard_dir=null + last_verified: null + metrics: + initial_loss: null + final_loss: null + last_10_steps_step_time_ms_avg: null + last_10_steps_model_tflops_per_gpu_avg: null + resume_comparison: + reference_item: pretrain + sentinel_steps: [51, 100] + loss_relative_tolerance: 1.0e-6 + loss_absolute_tolerance: 1.0e-6 + sentinels_match: null + expected_result: > + After the pending uninterrupted reference completes, this command must + restore optimizer, scheduler, data-order, and RNG state directly from + iter_0000050, begin at step 51, and finish at step 100 in the distinct + output directory. Step-51 and step-100 loss must satisfy the configured + absolute-plus-relative tolerance against the reference, with no skipped + or NaN iterations. diff --git a/scripts/conversion/arguments.py b/scripts/conversion/arguments.py index 9d13a525c3..732ac04db2 100644 --- a/scripts/conversion/arguments.py +++ b/scripts/conversion/arguments.py @@ -153,6 +153,10 @@ def _add_common_conversion_arguments(parser: argparse.ArgumentParser, *, include conversion = parser.add_argument_group("Conversion") conversion.add_argument("--hf-model", required=True, help="Hugging Face model ID or local path.") + conversion.add_argument( + "--hf-revision", + help="Immutable Hugging Face Hub revision to resolve before conversion (for example, a commit SHA).", + ) conversion.add_argument("--megatron-path", required=True, help="Megatron checkpoint path.") conversion.add_argument( "--torch-dtype", @@ -189,6 +193,10 @@ def _add_roundtrip_arguments(parser: argparse.ArgumentParser, *, include_executi dest="hf_model", help="Hugging Face model ID or local path.", ) + roundtrip.add_argument( + "--hf-revision", + help="Immutable Hugging Face Hub revision to resolve before conversion (for example, a commit SHA).", + ) roundtrip.add_argument( "--trust-remote-code", action="store_true", @@ -304,6 +312,8 @@ def conversion_worker_args(args: argparse.Namespace) -> list[str]: "--etp", str(args.etp), ] + if args.hf_revision is not None: + worker_args.extend(["--hf-revision", args.hf_revision]) if args.trust_remote_code: worker_args.append("--trust-remote-code") if args.distributed_timeout_minutes is not None: @@ -329,6 +339,8 @@ def conversion_worker_args(args: argparse.Namespace) -> list[str]: "--etp", str(args.etp), ] + if args.hf_revision is not None: + worker_args.extend(["--hf-revision", args.hf_revision]) if args.trust_remote_code: worker_args.append("--trust-remote-code") if args.overwrite: diff --git a/scripts/conversion/run_conversion.py b/scripts/conversion/run_conversion.py index 939e0d1244..8d6af4cc0a 100755 --- a/scripts/conversion/run_conversion.py +++ b/scripts/conversion/run_conversion.py @@ -20,6 +20,7 @@ import cpu_backend import gpu_backend from arguments import build_parser +from utils import resolve_hf_model_revision logger = logging.getLogger(__name__) @@ -120,6 +121,9 @@ def main(argv: list[str] | None = None) -> None: """Parse worker arguments and run checkpoint conversion.""" args = build_parser(include_execution=False).parse_args(argv) _validate_args(args) + if args.hf_revision is not None: + args.hf_model = resolve_hf_model_revision(args.hf_model, args.hf_revision) + logger.info("Resolved Hugging Face model at revision %s", args.hf_revision) logger.info("Selected %s backend for %s conversion", args.device.upper(), args.command) if args.command == "import": _run_import(args) diff --git a/scripts/conversion/setup_conversion.py b/scripts/conversion/setup_conversion.py index 3b4b228c72..54d17010b6 100755 --- a/scripts/conversion/setup_conversion.py +++ b/scripts/conversion/setup_conversion.py @@ -163,6 +163,7 @@ def _build_executor( executor = run.SlurmExecutor( account=args.account, partition=args.partition, + job_name_prefix=args.experiment_name, nodes=args.nodes, ntasks_per_node=task_count, mem=args.mem, diff --git a/scripts/conversion/utils.py b/scripts/conversion/utils.py index 04877d5d27..cfcb4d4260 100644 --- a/scripts/conversion/utils.py +++ b/scripts/conversion/utils.py @@ -27,6 +27,31 @@ } +def resolve_hf_model_revision(hf_model: str, hf_revision: str | None) -> str: + """Resolve a remote Hugging Face model revision to an immutable local snapshot. + + Args: + hf_model: Hugging Face model ID or local path. + hf_revision: Hub branch, tag, or commit to resolve. ``None`` preserves + the original model reference. + + Returns: + The original model reference when no revision is provided, otherwise + the local path of the resolved Hub snapshot. + + Raises: + ValueError: If a revision is paired with an existing local path. + """ + if hf_revision is None: + return hf_model + if Path(hf_model).expanduser().exists(): + raise ValueError("--hf-revision applies only to Hugging Face Hub model IDs, not local paths.") + + from huggingface_hub import snapshot_download + + return snapshot_download(repo_id=hf_model, revision=hf_revision) + + def parse_dtype(name: str) -> torch.dtype: """Resolve a CLI dtype name. diff --git a/scripts/training/setup_experiment.py b/scripts/training/setup_experiment.py index 4a85a74454..65a7b778ea 100755 --- a/scripts/training/setup_experiment.py +++ b/scripts/training/setup_experiment.py @@ -124,6 +124,11 @@ def _build_parser() -> argparse.ArgumentParser: dest="submission_dry_run", help="Render the Slurm submission without submitting it.", ) + execution.add_argument( + "--wait", + action="store_true", + help="Wait for the Slurm experiment to finish and stream its logs.", + ) return parser @@ -268,7 +273,7 @@ def main(argv: list[str] | None = None) -> None: if args.submission_dry_run: experiment.dryrun() return - experiment.run(detach=True) + experiment.run(detach=not args.wait, tail_logs=args.wait) if __name__ == "__main__": diff --git a/skills/create-model-verification-card/SKILL.md b/skills/create-model-verification-card/SKILL.md new file mode 100644 index 0000000000..e83b3fb732 --- /dev/null +++ b/skills/create-model-verification-card/SKILL.md @@ -0,0 +1,500 @@ +--- +name: create-model-verification-card +description: Create or update concise, agent-readable Megatron Bridge model verification cards. Use when adding a model support card, auditing cross-model convergence comparability or verification coverage, recording conversion, deterministic inference, training, checkpoint resume, post-SFT export, or performance results, or preparing a model-support PR. Enforce the required core inventory, convergence-versus-performance contracts, optional canonical performance item, public Slurm launcher commands, training metrics, important-feature allowlist, and a strict privacy boundary that excludes private runtime wiring, internal paths, credentials, and job metadata. +--- + +# Create Model Verification Card + +Create `model_cards//card.yaml` from public model facts and verified +commands. Keep the card small enough for an agent to scan without interpreting +logs or reconstructing the execution environment. + +## Use the repository resources + +- Validate the result with [scripts/validate_card.py](scripts/validate_card.py). +- Verify exact-length deterministic HF output with + [scripts/verify_hf_inference.py](scripts/verify_hf_inference.py). +- Use the inventory and field rules below as the format contract. Do not infer + model-specific settings from another family or variant. + +Do not add a README, evidence blobs, log excerpts, runtime setup, or scheduler +metadata to the skill or card. + +## Workflow + +### 1. Pull facts before drafting + +Read the model implementation, public HF config, conversion bridge, recipes, +tests, examples, and the exact revision being verified. Determine which modes +exist; do not infer support from a family name or another model size. + +Record the public HF model name in commands, its immutable revision in +`model.hf_revision`, and the minimum supported Transformers version in +`model.min_transformers_version`. Do not introduce an HF snapshot path. + +Record only two execution-environment facts: the public base container +identifier and the exact Bridge commit used for verification. Put them in +`verification_environment.base_container` and +`verification_environment.bridge_commit`. If the run mounted a checkout over +the container, record the mounted checkout commit. Never substitute a private +image path for the public base container identifier. + +The top-level Bridge commit is the default for every item. If one verified item +was run from a different clean checkout, put that exact 40-hex commit in the +item's optional `bridge_commit` field. Omit the field when it would repeat the +top-level value, and never use a commit field to disguise uncommitted runtime +changes. Items that are not verified must not carry a commit override. + +### 2. Create the core inventory and add performance when available + +Include these twelve required items, even when their status is `unsupported` or +`not_applicable`: + +1. CPU HF-to-Megatron conversion +2. GPU HF-to-Megatron conversion +3. CPU Megatron-to-HF conversion +4. GPU Megatron-to-HF conversion +5. Manual HF/Megatron forward-pass correlation +6. Deterministic Megatron inference +7. Pretraining +8. SFT +9. SFT checkpoint export and deterministic HF inference +10. Long-context SFT +11. PEFT +12. Checkpoint resume + +Add `pretrain_performance` as a thirteenth item only when the exact model +variant has a canonical public performance recipe. If no such recipe is +exported, omit the item instead of adding an unverified placeholder. Once a +canonical recipe exists, keep the item in the card even if its run is still +unverified. + +Use only `unverified`, `verified`, `unsupported`, or `not_applicable`. Do not +add `smoke` or an evidence field. + +For `unsupported` and `not_applicable`, leave `command` or `commands`, date, +precision, GPU type, and metrics null, then state the public product limitation +in `expected_result`. + +Put a scalar `precision` on every item. It describes the workload that was +actually verified, not every precision the model might support: + +- for conversion, record the imported or exported weight precision; +- for forward pass and inference, record the compute precision; +- for training, record the recipe's mixed-precision mode. + +Use `bf16` for BF16. Training items may instead use `fp8_mx` for MXFP8 or +`nvfp4` for NVFP4. Keep MXFP8 and NVFP4 training-only, and do not list either +until that exact item has completed in that mode. + +### 3. Use the public Slurm launchers + +Assume the caller supplies the account, partition, concrete runtime image, +credentials, and storage mappings outside the card. The top-level +`verification_environment.base_container` is provenance only; it is not +launcher configuration. Record the portable public launcher and the +model-verification workload: + +- use `scripts/conversion/convert.sh --executor slurm` for CPU and GPU + conversion, with portable node and GPU counts; +- use `scripts/training/train.sh --nodes ... --gpus-per-node ...` for every + training item; +- use `uv run python ...` only for inference helpers that do not yet have a + public Slurm executor; +- use short, ignored repository-relative logical paths under `work/...`; + prefer aliases such as `work/data/` and `work/cache/` over + reproducing a physical storage hierarchy. + +The public launchers may read their required generic Slurm configuration from +the caller's environment. Do not include `srun`, `sbatch`, concrete account or +partition values, container image arguments, `--mount`, `--env`, shell exports, +environment-variable references, cluster-specific `--srun-arg` values, or +launcher overlays. That wiring is personal to the verifier and does not belong +in a model support card. + +Never record or reproduce: + +- execution-environment names, hostnames, IPs, usernames, emails, or accounts; +- concrete partitions, reservations, node lists, private image locations, + mount sources, or environment forwarding; +- host/shared-storage paths, home-directory paths, log locations, or job IDs; +- tokens, token-loading commands, private URLs, or private registry references; +- environment-specific launcher overlays. + +Keep private run notes outside the tracked repository and public PR text. If a +private codename cannot be recognized generically, pass it to the validator via +`--deny-term` or an untracked file through `--denylist "$PRIVATE_DENYLIST"`. +The validator reports the match without printing the private term. + +### 4. Freeze convergence, execution, and benchmark contracts + +Before launching any training item, resolve the selected recipe and classify +its effective configuration into the three groups below. Do this from the +built `ConfigContainer`, not only from command-line overrides. Record the +resolved convergence and execution field/value fingerprints in the internal +per-model record; keep the public card concise. + +The **convergence contract** contains settings that change the training +objective, examples seen at an optimizer boundary, or numerical update rule: + +- starting checkpoint and trainable parameter set; +- dataset identity/revision, split or bounded selection, sample order, seeds, + tokenizer/chat template, truncation, masking, and packing semantics; +- sequence length, global batch size, global tokens per optimizer step, total + optimizer steps, and total token budget; +- objective and loss settings, including label masking, MoE auxiliary/router + losses, token dropping/capacity, natural versus forced routing, and loss + normalization; +- optimizer family, peak/minimum learning rate, schedule shape, warmup and + decay horizon, betas, epsilon, weight decay, gradient clipping, and dropout; +- model/gradient/optimizer-state precision and loss-scaling behavior; +- for PEFT, adapter type, targets, rank, alpha, dropout, and which base weights + are frozen. + +The **execution/performance contract** maps the frozen training semantics to +hardware. It may vary across models or be tuned for throughput: + +- node/GPU count and TP, PP, VP, CP, EP, ETP, DP, and sequence parallelism; +- activation recompute, activation/optimizer offload, distributed optimizer or + FSDP sharding, checkpoint I/O strategy, and garbage-collection policy; +- communication overlap, fused kernels, Transformer Engine implementation, + attention backend, CUDA graphs, and compilation; +- MoE transport/dispatcher backend such as all-to-all, DeepEP, or HybridEP, + provided routing, capacity, token dropping, and auxiliary losses are + unchanged. + +Treat micro batch size and gradient-accumulation layout as execution +fingerprints. They may be tuned while global batch size, global batch +membership/order, loss normalization, optimizer-step boundaries, and total +token budget remain intact. Because they can change accumulation order, +dropout RNG, MoE token grouping, and auxiliary-loss reduction, require fresh +loss sentinels for every layout and do not claim step-by-step numerical parity. + +Performance settings are **intended** to preserve training semantics, not +guaranteed to be bitwise neutral. Parallel reductions, fusions, recompute, and +dispatcher implementations can change floating-point order. After changing +them, require finite loss, no skipped iterations, and compatible loss sentinels +before calling the mapping verified. Anything that changes arithmetic +precision, forced router balancing, token dropping, packing, or effective batch +construction is a convergence change, even when introduced to improve speed. + +The **benchmark-only configuration** may deliberately change semantics to find +an upper throughput bound. It includes mock data, forced MoE load balancing, +changed batch/LR or timing-only schedules, and disabled NaN/large-gradient, +evaluation, or checkpoint checks. These settings may be valid for a canonical +`pretrain_performance` item, but their losses and checkpoints are never +convergence evidence. + +Use `qwen3_30b_a3b_convergence_v1` as the named default cross-model bounded +convergence cohort. It is the target contract derived from the resolved +Qwen3-30B-A3B H100 recipes; the name identifies settings, not evidence status. +Do not call a workload cohort-verified until its recipe owns this contract and +a clean-commit run passes the applicable gates below. Its 100 optimizer steps +test finite loss, short-horizon loss trend, checkpoint reload, and direct +resume; they do not establish full training convergence. + +Freeze this optimizer fingerprint for all three workloads: + +| Field | Value | +| --- | --- | +| Optimizer | Megatron distributed fused Adam | +| Betas / epsilon | `(0.9, 0.95)` / `1e-8` | +| Effective weight decay | `0.033`, constant | +| Gradient clipping | `1.0` | +| LR schedule | Cosine, starting from zero | +| Model / compute precision | BF16 | +| Optimizer master parameters, main gradients, moments | FP32 | + +Record the effective weight decay applied to optimizer parameter groups, not +only the nominal optimizer-config value. The Qwen anchor has a nominal +`optimizer.weight_decay=0.1`, but its constant scheduler applies `0.033`; use +`0.033` when reproducing this contract. Treat any other effective value +as a convergence-contract change. + +Freeze the following workload profiles: + +| Field | Pretrain | Full SFT | PEFT | +| --- | --- | --- | --- | +| Start / trainable set | Random initialization, no checkpoint load, full model | Exact immutable HF checkpoint revision, full model | Same immutable HF revision, frozen base model; LoRA on model-native attention Q/K/V and output projections, rank 8, alpha 16, dropout 0 | +| Data | Same bounded raw RP2 selection, revision, sample order, and seeds | Tulu 3 `train[:10000]`; same revision, order, chat template, label mask, truncation, and offline packing | Same as full SFT | +| Sequence / GBS | `4096 / 1024` | `2048 / 32` | `2048 / 32` | +| Reference MBS | `1` | `1` | `1` | +| Reference packing alignment | Not applicable | `pad_seq_to_mult=1` | `pad_seq_to_mult=4` | +| Token slots | `4,194,304` per step; `419,430,400` total | `65,536` per step; `6,553,600` total | `65,536` per step; `6,553,600` total | +| Peak / minimum LR | `3e-4 / 3e-5` | `5e-6 / 0` | `1e-4 / 0` | +| Horizon | 100 steps, 40 warmup steps, cosine decay through step 100, saves at steps 50 and 100 | 100 steps, 10 warmup steps, cosine decay through step 100, final checkpoint at step 100 | 100 steps, 10 warmup steps, cosine decay through step 100, final adapter checkpoint at step 100 | +| RNG | Model and dataset seed `1234` | Model RNG seed `5678`; data-order and packing seed `1234` | Model RNG seed `5678`; data-order and packing seed `1234` | +| Gradient path | BF16 gradient reduction; precision-aware optimizer enabled | FP32 gradient reduction; precision-aware optimizer disabled | FP32 gradient reduction; precision-aware optimizer disabled | + +Treat the reference packing alignments as frozen data semantics, not as +topology-derived performance defaults. Set them explicitly in the recipe and +build every compared run from a fresh packing output. If a topology or kernel +requires a larger alignment, record the resolved value, packing-manifest hash, +and actual supervised-token count, then classify that run as support +verification rather than evidence for this cross-model convergence cohort. +Never let runtime alignment synchronization silently change a cohort run. + +The public training runner applies `--dataset` after constructing the model +recipe and replaces the recipe's dataset object with the selected preset. +Consequently, a card command that uses `--dataset tulu3` must explicitly pin +the dataset revision, split, data-order/packing seed, offline-packing enablement, +and `+dataset.offline_packing_specs.pad_seq_to_mult`; recipe-level dataset +defaults alone do not freeze the resolved CLI workload. Use a fresh +`dataset.hf_output_root` (or force a deliberate rewrite) whenever any of these +fields changes, and audit the final `ConfigContainer` after all overrides and +runtime synchronization. + +Before launching training, create the parent directory named by +`logger.save_config_filepath` and require the post-setup file to persist. +Treat only that saved, post-synchronization `ConfigContainer` as runtime-config +evidence. YAML printed by the recipe runner before setup is launch-time +configuration and may differ after model finalization; never relabel it as +resolved runtime evidence or combine it with another run's logs. If the file +does not persist, fix the output path and rerun the workload from a fresh root. + +Treat the token counts above as token slots. For SFT and PEFT, also record the +actual supervised-token count after label masking; do not present padded or +masked token slots as supervised tokens. + +Use these accumulation layouts for the Qwen reference executions: + +| Workload | Reference topology | DP | Gradient accumulation | +| --- | --- | ---: | ---: | +| Pretrain | 16 GPUs, TP1/PP1/CP1/EP16 | 16 | 64 | +| Full SFT | 16 GPUs, TP1/PP1/CP1/EP16 | 16 | 2 | +| PEFT | 4 GPUs, TP4/PP1/CP1/EP4 | 1 | 32 | + +Topology, DP, micro batch size, and gradient accumulation are execution +fingerprints rather than convergence constraints. They may change while the +global batch size remains fixed. Every new execution layout must pass fresh +loss sentinels, and its step-by-step values are not strictly numerically +comparable with another layout. The current offline-packed SFT implementation +requires MBS=1; treat that as an implementation limit, not a convergence rule. + +For the Qwen anchor, record 128 experts, top-8 post-softmax-normalized routing, +auxiliary load-balancing loss coefficient `1e-3`, natural routing, and no +forced balancing or token dropping. These are model-native identity fields, +not universal cross-model overrides. Keep another model's native expert count, +top-k, router objective, auxiliary loss, capacity, and dropout values, record +them in its convergence fingerprint, and never alter them merely to imitate +Qwen. Require natural routing and prohibit benchmark-only forced balancing in +all convergence cohorts. + +The Qwen PEFT anchor names its fused attention projections `linear_qkv` and +`linear_proj`. Preserve the same semantic adapter scope on architectures that +split Q, K, and V: list every model-native attention projection explicitly and +record the names as a model-specific fingerprint. For Moonlight MLA this is +`linear_q_proj`, `linear_kv_down_proj`, `linear_kv_up_proj`, and +`linear_proj`. Never retain a nonexistent fused-module name merely to make the +textual configuration look identical; verify that every declared target +actually matches modules before accepting PEFT evidence. + +Use the same numerical value for global batch size within a comparison cohort. +If a model cannot use that value, change and validate its recipe separately, +record the exception, and treat the result as support verification rather than +an apples-to-apples convergence comparison. Compare progress at equal processed +token counts as well as equal optimizer steps. Different model architectures +and tokenizers make absolute cross-model loss values non-comparable; the shared +contract supports comparisons of stability and loss trend, not a ranking by +final loss. + +Pin the same raw-document selection across models, then tokenize it with each +model's verified tokenizer unless a deliberately shared tokenizer is part of +the cohort. Do not assume that one indexed token-ID prefix represents the same +text under different tokenizers. The Qwen anchor's pinned tokenizer revision +may reproduce the Qwen run, but do not silently reuse its token IDs for another +model family. Compare cross-model stability and trend, never absolute loss. + +Make every bounded convergence override explicit in the card command and apply +the same cohort values across models; never tune these values merely to improve +throughput. Library recipes own global and micro batch size. If a recipe batch +disagrees with the cohort contract, update and validate the recipe separately +instead of overriding it in the card. A canonical `perf_recipes` benchmark may +intentionally use mock data, forced balancing, or a different batch and is not +convergence evidence. + +Keep long-context SFT outside `qwen3_30b_a3b_convergence_v1`. Sequence length, +CP, packing, batch construction, LR, and horizon define a separate convergence +cohort even when the starting checkpoint and dataset are shared. + +### 5. Apply the verification gates + +Mark an item `verified` only after the workload represented by its command has +completed with the recorded model, recipe, data, and checkpoint arguments and +the concrete expected result has been checked. A successful detached Slurm +submission is not completion; wait for the submitted workload and inspect its +result. Private executor configuration stays outside the card. + +- **Conversion:** Test CPU and GPU import/export separately. Reload every + output. Require exact keys, shapes, dtypes, and values when the conversion is + expected to be lossless; otherwise state the numerical tolerance. Do not use + `--detach` or a dry-run flag in a verified conversion command. +- **Manual forward pass:** Compare Hugging Face and Megatron logits on the same + prompt with `examples/conversion/compare_hf_and_megatron/compare.py`. Record + whether the next token matches, the cosine similarity, and the maximum and + mean absolute logit differences. For new evidence, pass `--hf-revision` with + the exact `model.hf_revision` so the command itself is reproducibly pinned. + Historical evidence verified before 2026-07-20, when the helper gained + explicit revision pinning, may remain verified without a rerun when its + clean-run provenance is tied to the card's immutable `model.hf_revision`; + state this grandfathering explicitly in `expected_result`. Do not use the + exception for unverified items or evidence dated 2026-07-20 and later. + Mark the item verified when the next token matches and cosine similarity is + at least 0.99 (cosine distance at most 1%). Numeric maximum and mean absolute + differences are required diagnostic observations, but are report-only and + must not guard the item status. This gate establishes functional logit + correlation, not strict numerical equality. Keep this result separate from + generation. Choose a prompt whose tokenized length is divisible by TP so the + helper does not append padding before selecting the compared next-token + position. +- **Megatron inference:** Use deterministic greedy generation, specify an exact + token count, run twice, and record the literal completion including + whitespace. +- **Pretrain:** Use a bounded public dataset description and a stable schedule. + Save a middle and final checkpoint when resume is in scope. For expensive + workloads, a 100-step reference with checkpoints at steps 50 and 100 is a + suitable support-verification run when it crosses the peak learning rate and + completes the configured decay; resume directly from step 50 through step + 100. This verifies bounded training and resume behavior, not full convergence. +- **SFT and PEFT:** Prefer about 100 optimizer steps with warmup and full-horizon + decay. Use a public dataset name or preset, not its storage location. Save the + final full-SFT checkpoint when export verification is in scope. +- **SFT export and inference:** Depend on `sft`, export its final full-model + checkpoint to HF, reload the exported model with Transformers, and run greedy + generation twice. Store this item as an ordered `commands` list containing + exactly two strings: the synchronous Slurm export first and the `uv run` HF + inference second. Specify an exact new-token count and record the literal + byte-identical completion, including whitespace, in `expected_result`. +- **Long-context SFT:** Verify sequence packing and CP together. Record CP only + when its size is greater than one. +- **Checkpoint resume:** Depend on `pretrain`; load its middle checkpoint + directly, resume into a distinct new output root, load optimizer and RNG + state, and compare the first resumed and final steps with the uninterrupted + reference. For each declared loss sentinel, require this bound: + `abs(resumed - reference) <= 1e-6 + 0.01 * abs(reference)`. Tighter + model-specific tolerances are allowed. Do not repeat the pre-checkpoint + training segment. +- **Performance (when present):** Use the exact canonical public performance + recipe. Keep its bounded mock-data run separate from the real-data functional + run and state public hardware plus thresholds. + +Before adding checkpoint overrides, inspect the selected recipe and its +inherited checkpoint defaults. Keep only values that change the effective +behavior, such as an explicit load/save root, save interval, resume step, or +intentional strictness. For resume, the effective configuration must load and +save optimizer and RNG state and must not use finetune mode, but do not restate +those values when the recipe already inherits the required defaults. + +Submission is not success. Require the workload process to finish successfully, +the exact optimizer-step set to be present, finite losses and performance +values, zero skipped/NaN iterations, and complete reloadable artifacts. +For training items, also require the saved post-setup `ConfigContainer`; a +launch-time config dump or dry run cannot replace it. + +For ordered command lists, wait for the preceding workload to finish and verify +its artifact before starting the next command. Reference and resumed checkpoint +roots must be distinct; the resumed root must be new or empty. Use the same +Bridge commit, public base container, accelerator topology, and convergence +fingerprint as the reference run. Keep the execution fingerprint identical +when possible. Record any necessary execution-only deviation, explain why it +preserves semantics, and require the resume loss sentinels to pass. + +Use the batch sizes defined by the selected recipe. A model verification card +must not override global or micro batch size. If a recipe batch size is wrong, +change and validate the recipe separately instead of hiding the change in the +card command. + +### 6. Record training metrics consistently + +For every verified training item, record: + +- `initial_loss`: loss at the first optimizer step executed by that item; +- `final_loss`: loss at its final optimizer step; +- `last_10_steps_step_time_ms_avg`: arithmetic mean over the final 10 executed + optimizer steps; +- `last_10_steps_model_tflops_per_gpu_avg`: arithmetic mean over the same rows. + +Parse all fields from each complete keyed optimizer-step line. Do not collect +loss, time, and throughput independently and zip them. Reject missing, +duplicate, skipped, NaN, or non-finite rows rather than excluding them. +Every verified training item must execute at least 10 optimizer steps so this +window is complete. + +For a resume, use the first optimizer step after the selected checkpoint as +`initial_loss`, the final executed optimizer step as `final_loss`, and average +the final 10 resumed steps. For example, a step-50-to-100 continuation uses +step 51, step 100, and the average over steps 91-100. + +For Megatron-indexed data, command values are prefixes and omit `.bin` and +`.idx`. A bounded RedPajama2 prefix such as `head_01` is suitable for a +reproducible functional run; keep the physical dataset root private. + +### 7. Record only important enabled features + +Use `enabled_features` only on pretrain, SFT, long-context SFT, and PEFT. Keep +it empty when none of these are central to the verification. + +| Key | Allowed value | +| --- | --- | +| `sequence_packing` | `offline` or `in_batch` | +| `cuda_graph.implementation` | `local` or `transformer_engine` | +| `cuda_graph.scopes` | `full_iteration`, `attn`, `mlp`, `moe`, `moe_router`, `moe_preprocess`, or `mamba` | +| `context_parallel_size` | integer greater than one | +| `moe_dispatcher` | `deepep` or `hybridep` | + +Do not list routine TP/PP/DP sizes, Transformer Engine, fused loss, +distributed optimizer, ordinary communication overlap, LoRA, or DoRA. + +### 8. Validate before review + +Run: + +```bash +uv run --no-project --with pyyaml python \ + skills/create-model-verification-card/scripts/validate_card.py \ + model_cards//card.yaml +``` + +When private terminology is known only at runtime, add: + +```bash +--denylist "$PRIVATE_DENYLIST" +``` + +Then parse the YAML, run relevant targeted tests, run +`uv run pre-commit run --all-files`, and inspect the complete diff. Do not mark +an item verified merely to make validation pass. + +## Completion checklist + +- Keep all twelve core inventory items and use only the four statuses. Include + `pretrain_performance` only when the exact variant has a canonical public + performance recipe. +- Put the verified workload precision on every item; use `fp8_mx` and `nvfp4` + only for training items that ran in those modes. +- Pin a public immutable HF revision, minimum Transformers version, public base + container, and exact Bridge verification commit; use an item override only + for a verified workload run from a different clean commit. +- Use the public model name in commands. +- Include commands and concrete expected results for verified items. +- For manual forward pass, require a next-token match and cosine similarity of + at least 0.99, and record numeric max/mean absolute logit differences without + guarding on them. New evidence must pass the exact `model.hf_revision` + through `--hf-revision`; retain older unpinned evidence only under the + explicitly documented grandfathering rule above. +- Use `convert.sh --executor slurm` for conversion and `train.sh` for training. +- Keep private executor wiring out of commands: no mounts, environment + forwarding, concrete accounts/partitions/images, or remote-launch setup. +- Include GPU type and all four metrics for verified training items. +- Audit the resolved convergence contract before each training run; align it + with `qwen3_30b_a3b_convergence_v1` or record the exception and classify the + result as support verification rather than cross-model convergence evidence. +- Change only the execution/performance contract while tuning throughput, and + recheck loss sentinels after numerically non-bitwise changes. +- Leave recipe global and micro batch sizes unchanged in card commands. +- Save full SFT, export it to HF, and record an exact deterministic N-token HF + completion in a two-command ordered list. +- Keep resume as one direct continuation from the pretrain checkpoint. +- Keep enabled features within the four-family allowlist. +- Pass the bundled validator, including any caller-supplied denylist. +- Confirm the card, commit, and PR contain no private runtime information. diff --git a/skills/create-model-verification-card/scripts/validate_card.py b/skills/create-model-verification-card/scripts/validate_card.py new file mode 100644 index 0000000000..b1ea13469c --- /dev/null +++ b/skills/create-model-verification-card/scripts/validate_card.py @@ -0,0 +1,1415 @@ +#!/usr/bin/env python3 +# 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. + +"""Validate model verification card structure, claims, and privacy boundaries.""" + +from __future__ import annotations + +import argparse +import datetime as dt +import ipaddress +import logging +import math +import re +import shlex +import sys +from collections.abc import Iterable, Mapping +from pathlib import Path +from typing import Any, TypeGuard + +import yaml +from yaml.tokens import AliasToken, AnchorToken, TagToken + + +LOG = logging.getLogger(__name__) + +STATUSES = frozenset({"unverified", "verified", "unsupported", "not_applicable"}) +PRECISIONS = frozenset({"bf16", "fp8_mx", "nvfp4"}) +TRAINING_ONLY_PRECISIONS = PRECISIONS - {"bf16"} +REQUIRED_ITEM_NAMES = ( + "hf_to_megatron_cpu", + "hf_to_megatron_gpu", + "megatron_to_hf_cpu", + "megatron_to_hf_gpu", + "manual_forward_pass", + "inference", + "pretrain", + "sft", + "sft_export_inference", + "sft_long_context", + "peft", + "checkpoint_resume", +) +OPTIONAL_ITEM_NAMES = ("pretrain_performance",) +ITEM_NAMES = REQUIRED_ITEM_NAMES + OPTIONAL_ITEM_NAMES +TRAINING_ITEMS = frozenset( + {"pretrain", "sft", "sft_long_context", "peft", "checkpoint_resume", "pretrain_performance"} +) +CONVERSION_ITEMS = frozenset({"hf_to_megatron_cpu", "hf_to_megatron_gpu", "megatron_to_hf_cpu", "megatron_to_hf_gpu"}) +FEATURE_ITEMS = frozenset({"pretrain", "sft", "sft_long_context", "peft"}) +METRIC_NAMES = frozenset( + { + "initial_loss", + "final_loss", + "last_10_steps_step_time_ms_avg", + "last_10_steps_model_tflops_per_gpu_avg", + } +) +FEATURE_KEYS = frozenset({"sequence_packing", "cuda_graph", "context_parallel_size", "moe_dispatcher"}) +PACKING_VALUES = frozenset({"offline", "in_batch"}) +CUDA_GRAPH_IMPLEMENTATIONS = frozenset({"local", "transformer_engine"}) +CUDA_GRAPH_SCOPES = frozenset({"full_iteration", "attn", "mlp", "moe", "moe_router", "moe_preprocess", "mamba"}) +MOE_DISPATCHERS = frozenset({"deepep", "hybridep"}) +MANUAL_FORWARD_COSINE_THRESHOLD = 0.99 +MANUAL_FORWARD_REVISION_PINNING_DATE = dt.date(2026, 7, 20) + +TOP_LEVEL_KEYS = frozenset({"title", "model", "verification_environment", "summary", "items"}) +MODEL_KEYS = frozenset({"hf_id", "hf_revision", "architecture", "min_transformers_version"}) +ENVIRONMENT_KEYS = frozenset({"base_container", "bridge_commit"}) +ITEM_KEYS = frozenset( + { + "status", + "precision", + "bridge_commit", + "depends_on", + "command", + "commands", + "last_verified", + "expected_result", + "gpu_type", + "enabled_features", + "metrics", + "resume_comparison", + } +) +RESUME_KEYS = frozenset( + { + "reference_item", + "sentinel_steps", + "loss_relative_tolerance", + "loss_absolute_tolerance", + "sentinels_match", + } +) + +FORBIDDEN_KEY_FRAGMENTS = ( + "account", + "cluster", + "container", + "evidence", + "executor", + "host", + "job", + "launcher", + "mount", + "partition", +) + +PLACEHOLDER_RE = re.compile(r"\b(?:TODO|TBD|PLACEHOLDER)\b|<[^>]+>", re.IGNORECASE) +HF_ID_RE = re.compile(r"[A-Za-z0-9_.-]+/[A-Za-z0-9_.-]+") +REVISION_RE = re.compile(r"[0-9a-f]{40}") +VERSION_RE = re.compile(r"\d+\.\d+\.\d+") +PUBLIC_BASE_CONTAINER_RE = re.compile(r"nvcr\.io/nvidia/(?:nemo|pytorch):[A-Za-z0-9._-]+") +URL_RE = re.compile(r"\b[a-z][a-z0-9+.-]*://\S+", re.IGNORECASE) +EMAIL_RE = re.compile(r"\b[A-Za-z0-9._%+-]+@[A-Za-z0-9.-]+\.[A-Za-z]{2,}\b") +IPV4_RE = re.compile(r"(? dict[Any, Any]: + loader.flatten_mapping(node) + result: dict[Any, Any] = {} + for key_node, value_node in node.value: + key = loader.construct_object(key_node, deep=deep) + try: + duplicate = key in result + except TypeError as error: + raise yaml.constructor.ConstructorError( + "while constructing a mapping", + node.start_mark, + "found an unhashable key", + key_node.start_mark, + ) from error + if duplicate: + raise yaml.constructor.ConstructorError( + "while constructing a mapping", + node.start_mark, + f"found duplicate key {key!r}", + key_node.start_mark, + ) + result[key] = loader.construct_object(value_node, deep=deep) + return result + + +_UniqueKeyLoader.add_constructor(yaml.resolver.BaseResolver.DEFAULT_MAPPING_TAG, _construct_unique_mapping) + + +def _pointer(*parts: str) -> str: + if not parts: + return "/" + escaped = (part.replace("~", "~0").replace("/", "~1") for part in parts) + return "/" + "/".join(escaped) + + +def _check_keys( + value: Mapping[str, Any], + *, + allowed: frozenset[str], + required: frozenset[str], + path: tuple[str, ...], + errors: list[str], +) -> None: + non_string_keys = sorted(repr(key) for key in value if not isinstance(key, str)) + for key in non_string_keys: + errors.append(f"{_pointer(*path)}: non-string key {key} is not allowed") + keys = {key for key in value if isinstance(key, str)} + for key in sorted(keys - allowed): + errors.append(f"{_pointer(*path, str(key))}: unknown key") + for key in sorted(required - keys): + errors.append(f"{_pointer(*path, key)}: required key is missing") + + +def _as_mapping(value: Any, *, path: tuple[str, ...], errors: list[str]) -> Mapping[str, Any] | None: + if not isinstance(value, Mapping): + errors.append(f"{_pointer(*path)}: expected a mapping") + return None + if any(not isinstance(key, str) for key in value): + errors.append(f"{_pointer(*path)}: mapping keys must be strings") + return None + return value + + +def _is_iso_date(value: Any) -> bool: + if isinstance(value, dt.datetime): + return False + if isinstance(value, dt.date): + return True + if isinstance(value, str): + try: + dt.date.fromisoformat(value) + except ValueError: + return False + return True + return False + + +def _is_finite_number(value: object) -> TypeGuard[int | float]: + return isinstance(value, int | float) and not isinstance(value, bool) and math.isfinite(float(value)) + + +def _contains_ipv6(value: str) -> bool: + for match in IPV6_CANDIDATE_RE.finditer(value): + try: + address = ipaddress.ip_address(match.group(0)) + except ValueError: + continue + if address.version == 6: + return True + return False + + +def _command_value(item: Mapping[str, Any]) -> str | None: + command = item.get("command") + if isinstance(command, str) and command.strip(): + return command + return None + + +def _command_values(item: Mapping[str, Any]) -> list[str]: + command = item.get("command") + if isinstance(command, str) and command.strip(): + return [command] + commands = item.get("commands") + if isinstance(commands, list): + return [value for value in commands if isinstance(value, str) and value.strip()] + return [] + + +def _validate_enabled_features(value: Any, *, item_name: str, errors: list[str]) -> None: + path = ("items", item_name, "enabled_features") + features = _as_mapping(value, path=path, errors=errors) + if features is None: + return + _check_keys(features, allowed=FEATURE_KEYS, required=frozenset(), path=path, errors=errors) + + packing = features.get("sequence_packing") + if packing is not None and (not isinstance(packing, str) or packing not in PACKING_VALUES): + errors.append(f"{_pointer(*path, 'sequence_packing')}: expected offline or in_batch") + + cuda_graph_value = features.get("cuda_graph") + if cuda_graph_value is not None: + cuda_graph = _as_mapping(cuda_graph_value, path=(*path, "cuda_graph"), errors=errors) + if cuda_graph is not None: + cuda_path = (*path, "cuda_graph") + _check_keys( + cuda_graph, + allowed=frozenset({"implementation", "scopes"}), + required=frozenset({"implementation", "scopes"}), + path=cuda_path, + errors=errors, + ) + implementation = cuda_graph.get("implementation") + if not isinstance(implementation, str) or implementation not in CUDA_GRAPH_IMPLEMENTATIONS: + errors.append(f"{_pointer(*cuda_path, 'implementation')}: expected local or transformer_engine") + scopes = cuda_graph.get("scopes") + if not isinstance(scopes, list) or not scopes: + errors.append(f"{_pointer(*cuda_path, 'scopes')}: expected a non-empty list") + elif not all(isinstance(scope, str) for scope in scopes): + errors.append(f"{_pointer(*cuda_path, 'scopes')}: scopes must be strings") + else: + unknown_scopes = sorted(set(scopes) - CUDA_GRAPH_SCOPES) + if unknown_scopes: + errors.append(f"{_pointer(*cuda_path, 'scopes')}: unknown scopes {unknown_scopes}") + if len(scopes) != len(set(scopes)): + errors.append(f"{_pointer(*cuda_path, 'scopes')}: duplicate scopes are not allowed") + + cp_size = features.get("context_parallel_size") + if cp_size is not None and (not isinstance(cp_size, int) or isinstance(cp_size, bool) or cp_size <= 1): + errors.append(f"{_pointer(*path, 'context_parallel_size')}: expected an integer greater than one") + + dispatcher = features.get("moe_dispatcher") + if dispatcher is not None and (not isinstance(dispatcher, str) or dispatcher not in MOE_DISPATCHERS): + errors.append(f"{_pointer(*path, 'moe_dispatcher')}: expected deepep or hybridep") + + +def _validate_metrics(item: Mapping[str, Any], *, item_name: str, status: str, errors: list[str]) -> None: + path = ("items", item_name, "metrics") + metrics = _as_mapping(item.get("metrics"), path=path, errors=errors) + if metrics is None: + return + _check_keys(metrics, allowed=METRIC_NAMES, required=METRIC_NAMES, path=path, errors=errors) + + if status == "verified": + for name in sorted(METRIC_NAMES): + value = metrics.get(name) + if not _is_finite_number(value): + errors.append(f"{_pointer(*path, name)}: verified metrics must be finite numbers") + elif ( + name + in { + "last_10_steps_step_time_ms_avg", + "last_10_steps_model_tflops_per_gpu_avg", + } + and float(value) <= 0 + ): + errors.append(f"{_pointer(*path, name)}: verified performance metrics must be positive") + elif status in {"unsupported", "not_applicable"}: + for name in sorted(METRIC_NAMES): + if metrics.get(name) is not None: + errors.append(f"{_pointer(*path, name)}: must be null for status {status}") + + +def _argument_value(command: str, argument: str) -> str | None: + values = _argument_values(command, argument) + if not values: + return None + return values[0] + + +def _argument_values(command: str, argument: str) -> list[str]: + matches = re.findall(rf"(?:^|\s){re.escape(argument)}(?:=|\s+)(\"[^\"]+\"|'[^']+'|[^\s]+)", command) + return [match.strip("\"'") for match in matches] + + +def _config_override_values(command: str, field: str) -> list[str]: + """Return shell-parsed values for a trailing ConfigContainer override.""" + try: + tokens = shlex.split(command) + except ValueError: + return [] + prefix = f"{field}=" + normalized_tokens = (token.lstrip("+") for token in tokens) + return [token[len(prefix) :] for token in normalized_tokens if token.startswith(prefix)] + + +def _resume_reference_settings(command: str) -> list[tuple[str, str, str | None]] | None: + """Return order-independent settings that must match a reference run.""" + try: + tokens = shlex.split(command) + except ValueError: + return None + + ignored_arguments = { + "--load-dir", + "--load_dir", + "--save-dir", + "--save-interval", + "--save_dir", + "--save_interval", + } + ignored_checkpoint_overrides = { + "checkpoint.ckpt_step", + "checkpoint.finetune", + "checkpoint.load", + "checkpoint.load_optim", + "checkpoint.load_rng", + "checkpoint.save", + "checkpoint.save_optim", + "checkpoint.save_rng", + } + ignored_runtime_overrides = {"train.empty_unused_memory_level"} + settings: list[tuple[str, str, str | None]] = [] + index = 1 # Both commands are validated separately as train.sh invocations. + while index < len(tokens): + token = tokens[index] + normalized = token.lstrip("+") + + if token.startswith("-"): + if "=" in token: + argument, value = token.split("=", maxsplit=1) + else: + argument = token + value = None + if index + 1 < len(tokens): + candidate = tokens[index + 1] + if not candidate.startswith("-") and "=" not in candidate: + value = candidate + index += 1 + if argument not in ignored_arguments: + settings.append(("argument", argument, value)) + elif "=" in normalized: + field, value = normalized.split("=", maxsplit=1) + if field not in ignored_checkpoint_overrides and field not in ignored_runtime_overrides: + settings.append(("override", field, value)) + else: + settings.append(("positional", normalized, None)) + index += 1 + + return sorted(settings, key=lambda setting: (setting[0], setting[1], setting[2] or "")) + + +def _resume_setting_names(settings: list[tuple[str, str, str | None]]) -> str: + """Describe mismatched settings without reproducing their possibly private values.""" + names = sorted({"" if kind == "positional" else name for kind, name, _ in settings}) + return ", ".join(names) + + +def _has_batch_size_override(command: str) -> bool: + try: + tokens = shlex.split(command) + except ValueError: + return False + option_names = { + "-gb", + "-mb", + "--global-batch-size", + "--global_batch_size", + "--micro-batch-size", + "--micro_batch_size", + } + config_names = { + "global_batch_size", + "micro_batch_size", + "train.global_batch_size", + "train.micro_batch_size", + } + for token in tokens: + normalized = token.lstrip("+") + name = normalized.split("=", 1)[0] + if name in option_names or name in config_names: + return True + return False + + +def _require_positive_integer_argument( + command: str, + argument: str, + *, + path: tuple[str, ...], + errors: list[str], +) -> None: + values = _argument_values(command, argument) + if len(values) != 1 or not values[0].isdigit() or int(values[0]) < 1: + errors.append(f"{_pointer(*path)}: requires exactly one positive integer {argument}") + + +def _validate_conversion_launcher( + command: str, + *, + operation: str, + device: str, + path: tuple[str, ...], + errors: list[str], +) -> None: + try: + tokens = shlex.split(command) + except ValueError: + return + if len(tokens) < 2 or tokens[:2] != ["./scripts/conversion/convert.sh", operation]: + errors.append(f"{_pointer(*path)}: conversion must use convert.sh {operation}") + executors = _argument_values(command, "--executor") + if executors != ["slurm"]: + errors.append(f"{_pointer(*path)}: conversion must specify --executor slurm exactly once") + devices = _argument_values(command, "--device") + if devices != [device]: + errors.append(f"{_pointer(*path)}: conversion must specify --device {device} exactly once") + _require_positive_integer_argument(command, "--nodes", path=path, errors=errors) + gpu_counts = _argument_values(command, "--gpus-per-node") + if device == "gpu": + _require_positive_integer_argument(command, "--gpus-per-node", path=path, errors=errors) + elif gpu_counts: + errors.append(f"{_pointer(*path)}: CPU conversion must not request GPUs") + if any(option in tokens for option in ("--detach", "--dry-run", "--submission-dry-run")): + errors.append(f"{_pointer(*path)}: verified conversion must wait for completion") + + +def _validate_training_launcher(command: str, *, item_name: str, errors: list[str]) -> None: + path = ("items", item_name, "command") + try: + tokens = shlex.split(command) + except ValueError: + return + if not tokens or tokens[0] != "./scripts/training/train.sh": + errors.append(f"{_pointer(*path)}: training must use ./scripts/training/train.sh") + _require_positive_integer_argument(command, "--nodes", path=path, errors=errors) + _require_positive_integer_argument(command, "--gpus-per-node", path=path, errors=errors) + if any(option in tokens for option in ("--dry-run", "--submission-dry-run")): + errors.append(f"{_pointer(*path)}: verified training command must submit the workload") + + +def _validate_command_text(command: str, *, path: tuple[str, ...], errors: list[str]) -> None: + try: + lexer = shlex.shlex(command, posix=True, punctuation_chars=";&|") + lexer.whitespace_split = True + lexer.commenters = "" + tokens = list(lexer) + except ValueError: + errors.append(f"{_pointer(*path)}: command must be shell-parseable") + return + if any(token in {"&", "&&", "|", "||", ";"} for token in tokens): + errors.append(f"{_pointer(*path)}: each entry must contain exactly one command") + if _has_batch_size_override(command): + errors.append(f"{_pointer(*path)}: card commands must use recipe batch sizes") + + +def _validate_resume(item: Mapping[str, Any], *, status: str, errors: list[str]) -> None: + path = ("items", "checkpoint_resume") + if item.get("depends_on") != "pretrain": + errors.append(f"{_pointer(*path, 'depends_on')}: must be pretrain") + + comparison = _as_mapping(item.get("resume_comparison"), path=(*path, "resume_comparison"), errors=errors) + if comparison is None: + return + comparison_path = (*path, "resume_comparison") + _check_keys( + comparison, + allowed=RESUME_KEYS, + required=RESUME_KEYS, + path=comparison_path, + errors=errors, + ) + if comparison.get("reference_item") != "pretrain": + errors.append(f"{_pointer(*comparison_path, 'reference_item')}: must be pretrain") + + for name in ("loss_relative_tolerance", "loss_absolute_tolerance"): + value = comparison.get(name) + if not _is_finite_number(value) or float(value) <= 0: + errors.append(f"{_pointer(*comparison_path, name)}: expected a positive finite number") + + relative_tolerance = comparison.get("loss_relative_tolerance") + if _is_finite_number(relative_tolerance) and float(relative_tolerance) > 1.0e-2: + errors.append(f"{_pointer(*comparison_path, 'loss_relative_tolerance')}: must not exceed 1%") + + if status != "verified": + return + + if not isinstance(item.get("command"), str): + errors.append(f"{_pointer(*path, 'command')}: verified resume must use one direct command") + return + command = item["command"] + required_fragments = ("checkpoint.ckpt_step=", "--load_dir", "--save_dir") + for fragment in required_fragments: + if fragment not in command: + errors.append(f"{_pointer(*path, 'command')}: missing {fragment}") + + # The inherited checkpoint defaults already provide full-state resume. The + # command may omit them for readability, but an explicit override must not + # disable the state needed for a valid continuation. + required_effective_values = { + "checkpoint.finetune": "false", + "checkpoint.load_optim": "true", + "checkpoint.load_rng": "true", + "checkpoint.save_optim": "true", + "checkpoint.save_rng": "true", + } + for field, expected_value in required_effective_values.items(): + values = _config_override_values(command, field) + if len(values) > 1: + errors.append(f"{_pointer(*path, 'command')}: {field} must not be overridden more than once") + elif values and values[0].lower() != expected_value: + errors.append(f"{_pointer(*path, 'command')}: {field} must remain {expected_value}") + + for field in ("checkpoint.load", "checkpoint.save"): + if _config_override_values(command, field): + errors.append( + f"{_pointer(*path, 'command')}: use the canonical --load_dir and --save_dir arguments, not {field}" + ) + + load_dirs = _argument_values(command, "--load_dir") + save_dirs = _argument_values(command, "--save_dir") + if len(load_dirs) != 1: + errors.append(f"{_pointer(*path, 'command')}: verified resume requires exactly one --load_dir") + if len(save_dirs) != 1: + errors.append(f"{_pointer(*path, 'command')}: verified resume requires exactly one --save_dir") + if len(load_dirs) == 1 and len(save_dirs) == 1 and load_dirs[0].rstrip("/") == save_dirs[0].rstrip("/"): + errors.append(f"{_pointer(*path, 'command')}: load and save directories must be distinct") + + ckpt_match = re.search(r"checkpoint\.ckpt_step=(\d+)", command) + max_steps_match = re.search(r"--max_steps(?:=|\s+)(\d+)", command) + sentinel_steps = comparison.get("sentinel_steps") + if ( + not isinstance(sentinel_steps, list) + or len(sentinel_steps) != 2 + or not all(isinstance(step, int) and not isinstance(step, bool) for step in sentinel_steps) + ): + errors.append(f"{_pointer(*comparison_path, 'sentinel_steps')}: expected two integer steps") + else: + if sentinel_steps != sorted(set(sentinel_steps)): + errors.append(f"{_pointer(*comparison_path, 'sentinel_steps')}: steps must be unique and sorted") + if ckpt_match is not None and sentinel_steps[0] != int(ckpt_match.group(1)) + 1: + errors.append( + f"{_pointer(*comparison_path, 'sentinel_steps')}: first sentinel must be the first resumed step" + ) + if max_steps_match is not None and sentinel_steps[-1] != int(max_steps_match.group(1)): + errors.append(f"{_pointer(*comparison_path, 'sentinel_steps')}: final sentinel must equal max_steps") + if comparison.get("sentinels_match") is not True: + errors.append(f"{_pointer(*comparison_path, 'sentinels_match')}: must be true when verified") + + +def _validate_resume_against_pretrain( + item: Mapping[str, Any], + pretrain_item: Mapping[str, Any], + *, + default_bridge_commit: str | None, + errors: list[str], +) -> None: + """Require a direct resume to use its reference run's checkpoint and settings.""" + if item.get("status") == "verified": + if pretrain_item.get("status") != "verified": + errors.append(f"{_pointer('items', 'checkpoint_resume', 'depends_on')}: pretrain must be verified first") + reference_commit = pretrain_item.get("bridge_commit") or default_bridge_commit + resume_commit = item.get("bridge_commit") or default_bridge_commit + if reference_commit != resume_commit: + errors.append( + f"{_pointer('items', 'checkpoint_resume', 'bridge_commit')}: " + "resume and pretrain must use the same Bridge commit" + ) + + resume_command = _command_value(item) + pretrain_command = _command_value(pretrain_item) + if resume_command is None or pretrain_command is None: + return + + resume_path = ("items", "checkpoint_resume", "command") + pretrain_path = ("items", "pretrain", "command") + reference_save_dirs = _argument_values(pretrain_command, "--save_dir") + resume_load_dirs = _argument_values(resume_command, "--load_dir") + if len(reference_save_dirs) != 1: + errors.append(f"{_pointer(*pretrain_path)}: checkpoint resume requires exactly one reference --save_dir") + if len(resume_load_dirs) != 1: + errors.append(f"{_pointer(*resume_path)}: checkpoint resume requires exactly one --load_dir") + if len(reference_save_dirs) == 1 and len(resume_load_dirs) == 1: + if reference_save_dirs[0].rstrip("/") != resume_load_dirs[0].rstrip("/"): + errors.append(f"{_pointer(*resume_path)}: --load_dir must equal the pretrain --save_dir") + + if _argument_values(pretrain_command, "--load_dir"): + errors.append(f"{_pointer(*pretrain_path)}: uninterrupted reference pretrain must not use --load_dir") + if _config_override_values(pretrain_command, "checkpoint.ckpt_step"): + errors.append( + f"{_pointer(*pretrain_path)}: uninterrupted reference pretrain must not set checkpoint.ckpt_step" + ) + reference_load_values = _config_override_values(pretrain_command, "checkpoint.load") + if len(reference_load_values) > 1: + errors.append(f"{_pointer(*pretrain_path)}: checkpoint.load must not be overridden more than once") + elif reference_load_values and reference_load_values[0].lower() not in {"none", "null"}: + errors.append(f"{_pointer(*pretrain_path)}: uninterrupted reference checkpoint.load must be null") + if _config_override_values(pretrain_command, "checkpoint.save"): + errors.append(f"{_pointer(*pretrain_path)}: use the canonical --save_dir argument, not checkpoint.save") + + required_reference_values = { + "checkpoint.finetune": "false", + "checkpoint.save_optim": "true", + "checkpoint.save_rng": "true", + } + for field, expected_value in required_reference_values.items(): + values = _config_override_values(pretrain_command, field) + if len(values) > 1: + errors.append(f"{_pointer(*pretrain_path)}: {field} must not be overridden more than once") + elif values and values[0].lower() != expected_value: + errors.append(f"{_pointer(*pretrain_path)}: {field} must remain {expected_value} for checkpoint resume") + + reference_settings = _resume_reference_settings(pretrain_command) + resume_settings = _resume_reference_settings(resume_command) + if reference_settings is None or resume_settings is None or reference_settings == resume_settings: + return + + differing_settings = [ + setting + for setting in {*reference_settings, *resume_settings} + if reference_settings.count(setting) != resume_settings.count(setting) + ] + errors.append( + f"{_pointer(*resume_path)}: reference and resume launch settings must match; " + f"differences found in {_resume_setting_names(differing_settings)}" + ) + + +def _validate_inference( + item: Mapping[str, Any], + *, + item_name: str, + status: str, + errors: list[str], + command_override: str | None = None, + command_path: tuple[str, ...] | None = None, +) -> None: + if status != "verified": + return + path = ("items", item_name) + resolved_command_path = command_path or (*path, "command") + command = command_override if command_override is not None else _command_value(item) + expected = item.get("expected_result") + if command is None or not isinstance(expected, str): + return + try: + command_tokens = shlex.split(command) + except ValueError: + command_tokens = [] + if command_tokens[:2] != ["uv", "run"]: + errors.append(f"{_pointer(*resolved_command_path)}: inference must use uv run") + prompts = _argument_values(command, "--prompt") + if len(prompts) != 1: + errors.append(f"{_pointer(*resolved_command_path)}: specify --prompt exactly once") + token_matches = re.findall(r"--max[_-]new[_-]tokens(?:=|\s+)(\d+)", command) + if not token_matches: + errors.append(f"{_pointer(*resolved_command_path)}: specify an exact max_new_tokens value") + return + if len(token_matches) != 1: + errors.append(f"{_pointer(*resolved_command_path)}: specify max_new_tokens exactly once") + return + token_count_text = token_matches[0] + token_count = int(token_count_text) + if token_count <= 0: + errors.append(f"{_pointer(*resolved_command_path)}: max_new_tokens must be positive") + if "exact" not in expected.lower() or token_count_text not in expected: + errors.append(f"{_pointer(*path, 'expected_result')}: state the exact {token_count_text}-token result") + literals = [ + left or right + for left, right in re.findall( + r"\bcompletion\b[^\"']{0,200}(?:\"([^\"]+)\"|'([^']+)')", + expected, + re.IGNORECASE | re.DOTALL, + ) + ] + if not literals: + errors.append(f"{_pointer(*path, 'expected_result')}: quote the literal completion after the word completion") + else: + literal = max(literals, key=len) + if token_count > 0 and len(literal.encode()) < token_count: + errors.append( + f"{_pointer(*path, 'expected_result')}: literal completion is too short for {token_count} tokens" + ) + if len(prompts) == 1 and literal.strip() == prompts[0].strip(): + errors.append(f"{_pointer(*path, 'expected_result')}: literal completion must not repeat the prompt") + repeated = re.search(r"\b(?:twice|two\s+(?:independent\s+)?(?:runs|executions))\b", expected, re.IGNORECASE) + matched = re.search( + r"\b(?:byte[- ](?:for[- ])?byte|byte[- ]identical|identical|match(?:es|ed)?)\b", expected, re.IGNORECASE + ) + if repeated is None or matched is None: + errors.append( + f"{_pointer(*path, 'expected_result')}: state that two independent runs produced byte-identical output" + ) + if re.search(r"\bdo[_-]sample(?:=|\s+)(?:true|1)\b", command, re.IGNORECASE): + errors.append(f"{_pointer(*resolved_command_path)}: inference verification must be deterministic") + + +def _validate_manual_forward_pass( + item: Mapping[str, Any], + *, + status: str, + model_revision: str | None, + errors: list[str], +) -> None: + path = ("items", "manual_forward_pass") + command = _command_value(item) + expected = item.get("expected_result") + if command is None: + return + try: + tokens = shlex.split(command) + except ValueError: + tokens = [] + revision_values = _argument_values(command, "--hf-revision") + revision_flags = [token for token in tokens if token == "--hf-revision" or token.startswith("--hf-revision=")] + if len(revision_flags) != len(revision_values) or len(revision_values) > 1: + errors.append(f"{_pointer(*path, 'command')}: specify --hf-revision at most once with a value") + elif revision_values and revision_values[0] != model_revision: + errors.append(f"{_pointer(*path, 'command')}: --hf-revision must equal model.hf_revision") + elif not revision_values: + last_verified = item.get("last_verified") + if isinstance(last_verified, dt.datetime): + last_verified_date = None + elif isinstance(last_verified, dt.date): + last_verified_date = last_verified + elif isinstance(last_verified, str) and _is_iso_date(last_verified): + last_verified_date = dt.date.fromisoformat(last_verified) + else: + last_verified_date = None + historical_date = ( + status == "verified" + and last_verified_date is not None + and last_verified_date < MANUAL_FORWARD_REVISION_PINNING_DATE + ) + historical = ( + re.search(r"\b(?:grandfathered|historical|predates)\b", expected, re.IGNORECASE) + if isinstance(expected, str) + else None + ) + recorded_revision = ( + re.search(r"\brecorded immutable HF revision\b", expected, re.IGNORECASE) + if isinstance(expected, str) + else None + ) + if not historical_date or historical is None or recorded_revision is None: + errors.append( + f"{_pointer(*path, 'command')}: missing --hf-revision is allowed only for verified historical " + "evidence dated before 2026-07-20 and tied to the recorded immutable HF revision" + ) + + if status != "verified" or not isinstance(expected, str): + return + + prefix = ["uv", "run", "python", "-m", "torch.distributed.run"] + if tokens[:5] != prefix: + errors.append(f"{_pointer(*path, 'command')}: manual forward pass must use uv distributed run") + if "examples/conversion/compare_hf_and_megatron/compare.py" not in tokens: + errors.append(f"{_pointer(*path, 'command')}: use the HF/Megatron comparison helper") + for argument in ("--hf_model_path", "--megatron_model_path", "--prompt"): + if len(_argument_values(command, argument)) != 1: + errors.append(f"{_pointer(*path, 'command')}: specify {argument} exactly once") + + token_match = re.search( + r"\b(?:token match:\s*true|next[- ]token\b[^.]*\bmatch(?:es|ed)?)", expected, re.IGNORECASE + ) + token_mismatch = re.search( + r"\b(?:token match:\s*false|next[- ]token\b[^.]*\b(?:" + r"(?:not|never)\s+match|(?:doesn|didn)['’]t\s+match|fail(?:s|ed)?\s+to\s+match|mismatch))", + expected, + re.IGNORECASE, + ) + if token_match is None or token_mismatch is not None: + errors.append(f"{_pointer(*path, 'expected_result')}: record that the next token matches") + similarity_match = re.search( + r"cosine similarity(?:\s+is|\s*:)\s*(0(?:\.\d+)?|1(?:\.0+)?)", + expected, + re.IGNORECASE, + ) + if similarity_match is None: + errors.append(f"{_pointer(*path, 'expected_result')}: record the cosine similarity") + elif float(similarity_match.group(1)) < MANUAL_FORWARD_COSINE_THRESHOLD: + errors.append( + f"{_pointer(*path, 'expected_result')}: cosine similarity must be at least " + f"{MANUAL_FORWARD_COSINE_THRESHOLD:.2f}" + ) + number = r"(?:\d+(?:\.\d*)?|\.\d+)(?:[eE][-+]?\d+)?" + combined_differences = re.search( + rf"\bmax(?:imum)?\s+and\s+mean\s+(?:absolute\s+)?logit\s+differences?\s*(?:are|:|=)\s*" + rf"{number}\s+(?:and|,)\s*{number}", + expected, + re.IGNORECASE, + ) + helper_differences = re.search( + rf"\blogits?\s+diff(?:erences?)?\s*-?\s*max(?:imum)?\s*(?::|=|is|was)\s*{number}\s*,\s*" + rf"mean\s*(?::|=|is|was)\s*{number}", + expected, + re.IGNORECASE, + ) + for label in ("max", "mean"): + narrative_difference = re.search( + rf"\b{label}(?:imum)?\b[^.]*\b(?:absolute\s+)?logit\s+differences?\s*" + rf"(?:is|are|was|were|:|=)\s*{number}", + expected, + re.IGNORECASE, + ) + if combined_differences is None and helper_differences is None and narrative_difference is None: + errors.append( + f"{_pointer(*path, 'expected_result')}: record the numeric {label} absolute logit difference" + ) + + +def _validate_sft_export_inference(item: Mapping[str, Any], sft_item: Mapping[str, Any], *, errors: list[str]) -> None: + path = ("items", "sft_export_inference") + if item.get("depends_on") != "sft": + errors.append(f"{_pointer(*path, 'depends_on')}: must be sft") + status = item.get("status") + if status != "verified": + return + if sft_item.get("status") != "verified": + errors.append(f"{_pointer(*path, 'depends_on')}: sft must be verified first") + + commands = item.get("commands") + expected = item.get("expected_result") + if ( + not isinstance(commands, list) + or len(commands) != 2 + or not all(isinstance(command, str) and command.strip() for command in commands) + or not isinstance(expected, str) + ): + return + export_command, inference_command = commands + export_path = (*path, "commands", "0") + inference_path = (*path, "commands", "1") + + for fragment in ( + "scripts/conversion/convert.sh export", + "--megatron-path", + "--hf-path", + ): + if fragment not in export_command: + errors.append(f"{_pointer(*export_path)}: missing {fragment}") + for fragment in ("uv run", "scripts/verify_hf_inference.py"): + if fragment not in inference_command: + errors.append(f"{_pointer(*inference_path)}: missing {fragment}") + try: + inference_tokens = shlex.split(inference_command) + except ValueError: + inference_tokens = [] + expected_inference_prefix = [ + "uv", + "run", + "python", + "skills/create-model-verification-card/scripts/verify_hf_inference.py", + ] + if inference_tokens[:4] != expected_inference_prefix: + errors.append(f"{_pointer(*inference_path)}: must directly run the HF inference verifier with uv") + + export_devices = _argument_values(export_command, "--device") + if len(export_devices) == 1 and export_devices[0] in {"cpu", "gpu"}: + _validate_conversion_launcher( + export_command, + operation="export", + device=export_devices[0], + path=export_path, + errors=errors, + ) + else: + errors.append(f"{_pointer(*export_path)}: SFT export must specify --device cpu or gpu exactly once") + _validate_inference( + item, + item_name="sft_export_inference", + status=status, + errors=errors, + command_override=inference_command, + command_path=inference_path, + ) + + sft_command = _command_value(sft_item) + megatron_path = _argument_value(export_command, "--megatron-path") + if sft_command is not None: + save_dir = _argument_value(sft_command, "--save_dir") + max_steps = _argument_value(sft_command, "--max_steps") + if save_dir is None or max_steps is None or not max_steps.isdigit(): + errors.append( + f"{_pointer('items', 'sft', 'command')}: verified SFT export requires a final save directory" + ) + else: + expected_megatron_path = f"{save_dir.rstrip('/')}/iter_{int(max_steps):07d}" + if megatron_path != expected_megatron_path: + errors.append(f"{_pointer(*export_path)}: export must consume the final SFT checkpoint") + + hf_path = _argument_value(export_command, "--hf-path") + inference_hf_models = _argument_values(inference_command, "--hf-model") + if len(inference_hf_models) != 1 or hf_path is None or inference_hf_models[0] != hf_path: + errors.append(f"{_pointer(*inference_path)}: HF inference must reload the export destination") + if not re.search(r"\b(?:reload|reloaded|reloads|from_pretrained)\b", expected, re.IGNORECASE): + errors.append(f"{_pointer(*path, 'expected_result')}: state that the HF export reloads") + + +def _validate_training_window(item: Mapping[str, Any], *, item_name: str, status: str, errors: list[str]) -> None: + if status != "verified": + return + path = ("items", item_name, "command") + command = _command_value(item) + if command is None: + return + max_steps_values = _argument_values(command, "--max_steps") + if len(max_steps_values) != 1 or not max_steps_values[0].isdigit(): + errors.append(f"{_pointer(*path)}: verified training requires exactly one integer --max_steps") + return + final_step = int(max_steps_values[0]) + first_step = 1 + if item_name == "checkpoint_resume": + checkpoint_steps = re.findall(r"checkpoint\.ckpt_step=(\d+)", command) + if len(checkpoint_steps) != 1: + errors.append(f"{_pointer(*path)}: verified resume requires exactly one checkpoint.ckpt_step") + return + first_step = int(checkpoint_steps[0]) + 1 + if final_step - first_step + 1 < 10: + errors.append(f"{_pointer(*path)}: last-10 metrics require at least 10 executed optimizer steps") + + +def _validate_item(item_name: str, value: Any, errors: list[str], *, model_revision: str | None) -> None: + path = ("items", item_name) + item = _as_mapping(value, path=path, errors=errors) + if item is None: + return + required = frozenset({"status", "precision", "last_verified", "expected_result"}) + if item_name == "sft_export_inference": + required |= frozenset({"commands"}) + else: + required |= frozenset({"command"}) + if item_name in TRAINING_ITEMS: + required |= frozenset({"gpu_type", "metrics"}) + if item_name in FEATURE_ITEMS: + required |= frozenset({"enabled_features"}) + if item_name == "checkpoint_resume": + required |= frozenset({"depends_on", "resume_comparison"}) + if item_name == "sft_export_inference": + required |= frozenset({"depends_on"}) + _check_keys(item, allowed=ITEM_KEYS, required=required, path=path, errors=errors) + + status = item.get("status") + if not isinstance(status, str) or status not in STATUSES: + errors.append(f"{_pointer(*path, 'status')}: expected one of {sorted(STATUSES)}") + return + + bridge_commit = item.get("bridge_commit") + if "bridge_commit" in item: + if status != "verified": + errors.append(f"{_pointer(*path, 'bridge_commit')}: item overrides are allowed only when verified") + if not isinstance(bridge_commit, str) or REVISION_RE.fullmatch(bridge_commit) is None: + errors.append(f"{_pointer(*path, 'bridge_commit')}: expected an immutable 40-hex commit") + + precision = item.get("precision") + if precision is not None and (not isinstance(precision, str) or precision not in PRECISIONS): + errors.append(f"{_pointer(*path, 'precision')}: expected one of {sorted(PRECISIONS)} or null") + elif isinstance(precision, str) and item_name not in TRAINING_ITEMS and precision in TRAINING_ONLY_PRECISIONS: + errors.append(f"{_pointer(*path, 'precision')}: {precision} is supported only on training items") + if status == "verified" and precision is None: + errors.append(f"{_pointer(*path, 'precision')}: verified items require a concrete precision") + elif status in {"unsupported", "not_applicable"} and precision is not None: + errors.append(f"{_pointer(*path, 'precision')}: must be null for status {status}") + + command = item.get("command") + commands = item.get("commands") + command_entries: list[tuple[tuple[str, ...], str]] = [] + command_field = "commands" if item_name == "sft_export_inference" else "command" + if item_name == "sft_export_inference": + if "command" in item: + errors.append(f"{_pointer(*path, 'command')}: use the ordered commands list for this item") + if commands is not None: + if not isinstance(commands, list): + errors.append(f"{_pointer(*path, 'commands')}: expected a two-string list or null") + else: + if len(commands) != 2: + errors.append(f"{_pointer(*path, 'commands')}: expected exactly two ordered commands") + for index, value in enumerate(commands): + entry_path = (*path, "commands", str(index)) + if not isinstance(value, str) or not value.strip(): + errors.append(f"{_pointer(*entry_path)}: expected a non-empty command string") + else: + command_entries.append((entry_path, value)) + else: + if "commands" in item: + errors.append(f"{_pointer(*path, 'commands')}: only sft_export_inference uses a command list") + if command is not None and not isinstance(command, str): + errors.append(f"{_pointer(*path, 'command')}: expected a string or null") + elif isinstance(command, str): + command_entries.append(((*path, "command"), command)) + + for entry_path, entry in command_entries: + _validate_command_text(entry, path=entry_path, errors=errors) + + expected = item.get("expected_result") + if expected is not None and not isinstance(expected, str): + errors.append(f"{_pointer(*path, 'expected_result')}: expected a string or null") + + if status == "verified": + complete_commands = _command_values(item) + expected_count = 2 if item_name == "sft_export_inference" else 1 + if len(complete_commands) != expected_count: + errors.append(f"{_pointer(*path, command_field)}: verified item requires {expected_count} command(s)") + for entry_path, entry in command_entries: + if PLACEHOLDER_RE.search(entry): + errors.append(f"{_pointer(*entry_path)}: verified command contains a placeholder") + if not isinstance(expected, str) or not expected.strip(): + errors.append(f"{_pointer(*path, 'expected_result')}: verified items require a concrete result") + elif PLACEHOLDER_RE.search(expected): + errors.append(f"{_pointer(*path, 'expected_result')}: verified result contains a placeholder") + if not _is_iso_date(item.get("last_verified")): + errors.append(f"{_pointer(*path, 'last_verified')}: verified items require an ISO date") + elif status in {"unsupported", "not_applicable"}: + if item.get(command_field) is not None: + errors.append(f"{_pointer(*path, command_field)}: must be null for status {status}") + if item.get("last_verified") is not None: + errors.append(f"{_pointer(*path, 'last_verified')}: must be null for status {status}") + if not isinstance(expected, str) or not expected.strip() or PLACEHOLDER_RE.search(expected): + errors.append(f"{_pointer(*path, 'expected_result')}: explain the public limitation") + + if status == "verified" and item_name in CONVERSION_ITEMS: + complete_command = _command_value(item) + if complete_command is not None: + operation = "import" if item_name.startswith("hf_to_megatron") else "export" + device = "cpu" if item_name.endswith("cpu") else "gpu" + _validate_conversion_launcher( + complete_command, + operation=operation, + device=device, + path=(*path, "command"), + errors=errors, + ) + + if item_name in TRAINING_ITEMS: + gpu_type = item.get("gpu_type") + if status == "verified" and (not isinstance(gpu_type, str) or not gpu_type.strip()): + errors.append(f"{_pointer(*path, 'gpu_type')}: verified training requires a public GPU type") + if status in {"unsupported", "not_applicable"} and gpu_type is not None: + errors.append(f"{_pointer(*path, 'gpu_type')}: must be null for status {status}") + _validate_metrics(item, item_name=item_name, status=status, errors=errors) + _validate_training_window(item, item_name=item_name, status=status, errors=errors) + if status == "verified": + complete_command = _command_value(item) + if complete_command is not None: + _validate_training_launcher(complete_command, item_name=item_name, errors=errors) + elif item.get("metrics") is not None or item.get("gpu_type") is not None: + errors.append(f"{_pointer(*path)}: metrics and gpu_type are training-only fields") + + if item_name in FEATURE_ITEMS: + _validate_enabled_features(item.get("enabled_features"), item_name=item_name, errors=errors) + if item_name == "sft_long_context" and status == "verified": + features = item.get("enabled_features") + if isinstance(features, Mapping): + if features.get("sequence_packing") not in PACKING_VALUES: + errors.append( + f"{_pointer(*path, 'enabled_features', 'sequence_packing')}: " + "verified long-context SFT requires sequence packing" + ) + cp_size = features.get("context_parallel_size") + if not isinstance(cp_size, int) or isinstance(cp_size, bool) or cp_size <= 1: + errors.append( + f"{_pointer(*path, 'enabled_features', 'context_parallel_size')}: " + "verified long-context SFT requires CP greater than one" + ) + elif item.get("enabled_features") is not None: + errors.append(f"{_pointer(*path, 'enabled_features')}: field is not allowed on this item") + + if item_name == "checkpoint_resume": + _validate_resume(item, status=status, errors=errors) + if item_name == "manual_forward_pass": + _validate_manual_forward_pass(item, status=status, model_revision=model_revision, errors=errors) + if item_name == "inference": + _validate_inference(item, item_name=item_name, status=status, errors=errors) + + +def _walk_keys(value: Any, path: tuple[str, ...] = ()) -> Iterable[tuple[tuple[str, ...], str]]: + if isinstance(value, Mapping): + for key, child in value.items(): + key_text = str(key) + yield path, key_text + yield from _walk_keys(child, (*path, key_text)) + elif isinstance(value, list): + for index, child in enumerate(value): + yield from _walk_keys(child, (*path, str(index))) + + +def _validate_privacy(raw: str, card: Mapping[str, Any], deny_terms: tuple[str, ...], errors: list[str]) -> None: + for path, key in _walk_keys(card): + if path == ("verification_environment",) and key == "base_container": + continue + normalized = key.lower().replace("-", "_") + if any(fragment in normalized for fragment in FORBIDDEN_KEY_FRAGMENTS): + errors.append(f"{_pointer(*path, key)}: runtime identity or evidence fields are forbidden") + + privacy_raw = raw + environment = card.get("verification_environment") + if isinstance(environment, Mapping): + base_container = environment.get("base_container") + if isinstance(base_container, str) and PUBLIC_BASE_CONTAINER_RE.fullmatch(base_container) is not None: + privacy_raw = privacy_raw.replace(base_container, "") + + if re.search(r"\bcluster\b", privacy_raw, re.IGNORECASE): + errors.append("/: execution-environment names do not belong in the card") + if URL_RE.search(privacy_raw): + errors.append("/: URLs are forbidden; use a public model name or repository-relative path") + if EMAIL_RE.search(privacy_raw) or IPV4_RE.search(privacy_raw) or _contains_ipv6(privacy_raw): + errors.append("/: email or IP address detected") + if REMOTE_COMMAND_RE.search(privacy_raw) or REMOTE_COPY_RE.search(privacy_raw): + errors.append("/: remote host commands are forbidden") + if JOB_ID_RE.search(privacy_raw): + errors.append("/: scheduler job metadata is forbidden") + if SECRET_ASSIGNMENT_RE.search(privacy_raw): + errors.append("/: secret assignment detected") + if HOME_PATH_RE.search(privacy_raw): + errors.append("/: home-directory paths are forbidden") + if REGISTRY_REFERENCE_RE.search(privacy_raw): + errors.append("/: concrete registry references are forbidden") + commands: list[str] = [] + items = card.get("items") + if isinstance(items, Mapping): + for item in items.values(): + if isinstance(item, Mapping): + command = item.get("command") + if isinstance(command, str): + commands.append(command) + command_list = item.get("commands") + if isinstance(command_list, list): + commands.extend(value for value in command_list if isinstance(value, str)) + if any("$(" in command or "`" in command for command in commands): + errors.append("/: shell command substitution is forbidden in cards") + if any(SECRET_FLAG_RE.search(command) for command in commands): + errors.append("/: credential flags are forbidden in cards") + if any(ENVIRONMENT_ASSIGNMENT_RE.search(command) for command in commands): + errors.append("/: environment assignments are forbidden in cards") + if any(RUNTIME_COMMAND_RE.search(command) for command in commands): + errors.append("/: scheduler and container commands are forbidden in cards") + if any(RUNTIME_ENTRYPOINT_RE.search(command) for command in commands): + errors.append("/: use the public launcher rather than its setup implementation") + if any(RUNTIME_FLAG_RE.search(command) for command in commands): + errors.append("/: private runtime orchestration flags are forbidden in cards") + if any(SHELL_SETUP_RE.search(command) for command in commands): + errors.append("/: shell environment setup is forbidden in cards") + if "../" in privacy_raw: + errors.append("/: parent-directory traversal is forbidden in cards") + + raw_without_urls = URL_RE.sub("", privacy_raw) + for match in ABSOLUTE_PATH_RE.finditer(raw_without_urls): + errors.append(f"/: absolute path detected at character {match.start()}") + + environment_references = set(ENVIRONMENT_REFERENCE_RE.findall(privacy_raw)) + if environment_references: + errors.append( + f"/: {len(environment_references)} environment reference(s) detected; runtime wiring belongs outside the card" + ) + + lowered = privacy_raw.casefold() + for index, term in enumerate((term for term in deny_terms if term), start=1): + if term.casefold() in lowered: + errors.append(f"/: matched caller-supplied deny term #{index}") + + +def _validate_card(card: Mapping[str, Any], raw: str, deny_terms: tuple[str, ...]) -> list[str]: + errors: list[str] = [] + _check_keys(card, allowed=TOP_LEVEL_KEYS, required=TOP_LEVEL_KEYS, path=(), errors=errors) + + for name in ("title", "summary"): + value = card.get(name) + if not isinstance(value, str) or not value.strip(): + errors.append(f"{_pointer(name)}: expected a non-empty string") + + model = _as_mapping(card.get("model"), path=("model",), errors=errors) + if model is not None: + _check_keys(model, allowed=MODEL_KEYS, required=MODEL_KEYS, path=("model",), errors=errors) + if not isinstance(model.get("architecture"), str) or not model.get("architecture", "").strip(): + errors.append(f"{_pointer('model', 'architecture')}: expected a non-empty string") + min_version = model.get("min_transformers_version") + if not isinstance(min_version, str) or VERSION_RE.fullmatch(min_version) is None: + errors.append(f"{_pointer('model', 'min_transformers_version')}: expected MAJOR.MINOR.PATCH") + + environment = _as_mapping( + card.get("verification_environment"), + path=("verification_environment",), + errors=errors, + ) + if environment is not None: + _check_keys( + environment, + allowed=ENVIRONMENT_KEYS, + required=ENVIRONMENT_KEYS, + path=("verification_environment",), + errors=errors, + ) + base_container = environment.get("base_container") + if not isinstance(base_container, str) or PUBLIC_BASE_CONTAINER_RE.fullmatch(base_container) is None: + errors.append( + f"{_pointer('verification_environment', 'base_container')}: " + "expected a public NVIDIA NeMo or PyTorch container" + ) + bridge_commit = environment.get("bridge_commit") + if not isinstance(bridge_commit, str) or REVISION_RE.fullmatch(bridge_commit) is None: + errors.append( + f"{_pointer('verification_environment', 'bridge_commit')}: expected an immutable 40-hex commit" + ) + + items = _as_mapping(card.get("items"), path=("items",), errors=errors) + if items is not None: + item_names = set(items) + for name in sorted(item_names - set(ITEM_NAMES)): + errors.append(f"{_pointer('items', name)}: unknown verification item") + for name in sorted(set(REQUIRED_ITEM_NAMES) - item_names): + errors.append(f"{_pointer('items', name)}: required verification item is missing") + model_revision = model.get("hf_revision") if model is not None else None + if not isinstance(model_revision, str): + model_revision = None + for name in ITEM_NAMES: + if name in items: + _validate_item(name, items[name], errors, model_revision=model_revision) + sft_item = items.get("sft") + sft_export_item = items.get("sft_export_inference") + if isinstance(sft_item, Mapping) and isinstance(sft_export_item, Mapping): + _validate_sft_export_inference(sft_export_item, sft_item, errors=errors) + pretrain_item = items.get("pretrain") + resume_item = items.get("checkpoint_resume") + if isinstance(pretrain_item, Mapping) and isinstance(resume_item, Mapping): + default_bridge_commit = environment.get("bridge_commit") if environment is not None else None + _validate_resume_against_pretrain( + resume_item, + pretrain_item, + default_bridge_commit=default_bridge_commit if isinstance(default_bridge_commit, str) else None, + errors=errors, + ) + + if environment is not None: + default_bridge_commit = environment.get("bridge_commit") + for name, item in items.items(): + if ( + isinstance(default_bridge_commit, str) + and isinstance(item, Mapping) + and "bridge_commit" in item + and item.get("bridge_commit") == default_bridge_commit + ): + errors.append( + f"{_pointer('items', str(name), 'bridge_commit')}: " + "omit a redundant override of verification_environment.bridge_commit" + ) + + any_verified = bool( + items and any(isinstance(item, Mapping) and item.get("status") == "verified" for item in items.values()) + ) + if model is not None and any_verified: + hf_id = model.get("hf_id") + revision = model.get("hf_revision") + if not isinstance(hf_id, str) or HF_ID_RE.fullmatch(hf_id) is None: + errors.append(f"{_pointer('model', 'hf_id')}: verified cards require a public ORG/MODEL name") + elif hf_id == "ORG/MODEL": + errors.append(f"{_pointer('model', 'hf_id')}: replace the placeholder model name") + if not isinstance(revision, str) or REVISION_RE.fullmatch(revision) is None: + errors.append(f"{_pointer('model', 'hf_revision')}: verified cards require an immutable 40-hex revision") + if model.get("architecture") == "MODEL_ARCHITECTURE": + errors.append(f"{_pointer('model', 'architecture')}: replace the placeholder architecture") + if card.get("title") == "model_variant": + errors.append(f"{_pointer('title')}: replace the placeholder title") + + _validate_privacy(raw, card, deny_terms, errors) + return sorted(set(errors)) + + +def _load_card(raw: str) -> tuple[Mapping[str, Any] | None, list[str]]: + errors: list[str] = [] + try: + for token in yaml.scan(raw): + if isinstance(token, AnchorToken | AliasToken | TagToken): + errors.append("/: YAML anchors, aliases, and tags are forbidden") + except yaml.YAMLError as error: + return None, [f"/: invalid YAML token stream: {error}"] + if errors: + return None, sorted(set(errors)) + + try: + card = yaml.load(raw, Loader=_UniqueKeyLoader) + except yaml.YAMLError as error: + return None, sorted(set([*errors, f"/: invalid YAML: {error}"])) + if not isinstance(card, Mapping): + errors.append("/: card root must be a mapping") + return None, sorted(set(errors)) + return card, sorted(set(errors)) + + +def _read_denylist(path: Path | None) -> tuple[str, ...]: + if path is None: + return () + terms = [] + for line in path.read_text(encoding="utf-8").splitlines(): + term = line.strip() + if term and not term.startswith("#"): + terms.append(term) + return tuple(terms) + + +def _parse_args() -> argparse.Namespace: + parser = argparse.ArgumentParser(description=__doc__) + parser.add_argument("cards", nargs="+", help="Card paths, or - for standard input") + parser.add_argument( + "--deny-term", + action="append", + default=[], + help="Private term to reject without printing it; may be repeated", + ) + parser.add_argument("--denylist", type=Path, help="Untracked file containing private terms") + return parser.parse_args() + + +def main() -> int: + """Validate all requested cards and return a process exit code.""" + args = _parse_args() + try: + deny_terms = (*_read_denylist(args.denylist), *args.deny_term) + except OSError as error: + LOG.error("Unable to read denylist: %s", error) + return 2 + + failed = False + for card_path in args.cards: + source = "" if card_path == "-" else card_path + try: + raw = sys.stdin.read() if card_path == "-" else Path(card_path).read_text(encoding="utf-8") + except OSError as error: + LOG.error("%s: unable to read card: %s", source, error) + failed = True + continue + + card, load_errors = _load_card(raw) + errors = load_errors if card is None else [*load_errors, *_validate_card(card, raw, deny_terms)] + if errors: + failed = True + for validation_error in sorted(set(errors)): + LOG.error("%s%s", source, validation_error) + else: + LOG.info("%s: valid model verification card", source) + return 1 if failed else 0 + + +if __name__ == "__main__": + logging.basicConfig(level=logging.INFO, format="%(levelname)s: %(message)s") + raise SystemExit(main()) diff --git a/skills/create-model-verification-card/scripts/verify_hf_inference.py b/skills/create-model-verification-card/scripts/verify_hf_inference.py new file mode 100644 index 0000000000..e7ee4b85d8 --- /dev/null +++ b/skills/create-model-verification-card/scripts/verify_hf_inference.py @@ -0,0 +1,128 @@ +#!/usr/bin/env python3 +# 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. + +"""Verify deterministic, exact-length inference from an exported HF checkpoint.""" + +from __future__ import annotations + +import argparse +import json +import logging +from typing import Any + + +LOG = logging.getLogger(__name__) + + +def _parse_args() -> argparse.Namespace: + parser = argparse.ArgumentParser(description=__doc__) + parser.add_argument("--hf-model", required=True, help="Exported HF model directory.") + parser.add_argument("--prompt", required=True, help="Prompt to generate from.") + parser.add_argument("--max-new-tokens", required=True, type=int, help="Exact number of tokens to generate.") + parser.add_argument("--chat-template", action="store_true", help="Format the prompt as a user chat turn.") + parser.add_argument( + "--trust-remote-code", + action="store_true", + help="Allow custom model and tokenizer code from the selected Hugging Face repository.", + ) + parser.add_argument( + "--disable-thinking", + action="store_true", + help="Pass enable_thinking=False to the tokenizer chat template.", + ) + parser.add_argument("--device", default="cuda", help="Torch device used for inference.") + parser.add_argument( + "--dtype", + choices=("bfloat16", "float16", "float32"), + default="bfloat16", + help="Model loading dtype.", + ) + args = parser.parse_args() + if args.max_new_tokens <= 0: + parser.error("--max-new-tokens must be positive") + if args.disable_thinking and not args.chat_template: + parser.error("--disable-thinking requires --chat-template") + return args + + +def _format_prompt(tokenizer: Any, prompt: str, *, chat_template: bool, disable_thinking: bool) -> str: + if not chat_template: + return prompt + template_options = {"enable_thinking": False} if disable_thinking else {} + return tokenizer.apply_chat_template( + [{"role": "user", "content": prompt}], + tokenize=False, + add_generation_prompt=True, + **template_options, + ) + + +def main() -> int: + """Run two exact-length greedy generations and print their shared completion.""" + args = _parse_args() + + import torch + from transformers import AutoModelForCausalLM, AutoTokenizer + + dtype = getattr(torch, args.dtype) + tokenizer = AutoTokenizer.from_pretrained(args.hf_model, trust_remote_code=args.trust_remote_code) + model = ( + AutoModelForCausalLM.from_pretrained( + args.hf_model, + dtype=dtype, + trust_remote_code=args.trust_remote_code, + ) + .to(args.device) + .eval() + ) + formatted_prompt = _format_prompt( + tokenizer, + args.prompt, + chat_template=args.chat_template, + disable_thinking=args.disable_thinking, + ) + inputs = tokenizer(formatted_prompt, return_tensors="pt").to(model.device) + prompt_length = inputs["input_ids"].shape[1] + pad_token_id = tokenizer.pad_token_id if tokenizer.pad_token_id is not None else tokenizer.eos_token_id + + outputs = [] + with torch.inference_mode(): + for _ in range(2): + outputs.append( + model.generate( + **inputs, + do_sample=False, + min_new_tokens=args.max_new_tokens, + max_new_tokens=args.max_new_tokens, + pad_token_id=pad_token_id, + ) + ) + + expected_length = prompt_length + args.max_new_tokens + if any(output.shape != (1, expected_length) for output in outputs): + lengths = [output.shape[1] - prompt_length for output in outputs] + raise RuntimeError(f"Expected exactly {args.max_new_tokens} generated tokens; observed {lengths}") + if not torch.equal(outputs[0], outputs[1]): + raise RuntimeError("Two greedy HF inference runs produced different token IDs") + + completion_ids = outputs[0][0, prompt_length:].tolist() + completion = tokenizer.decode(completion_ids, skip_special_tokens=True) + LOG.info("HF completion (%d tokens): %s", args.max_new_tokens, json.dumps(completion, ensure_ascii=False)) + return 0 + + +if __name__ == "__main__": + logging.basicConfig(level=logging.INFO, format="%(levelname)s: %(message)s") + raise SystemExit(main()) diff --git a/skills/nemo-mbridge-recipe-recommender/SKILL.md b/skills/nemo-mbridge-recipe-recommender/SKILL.md index dfdd2c66fa..df8e6a0cef 100644 --- a/skills/nemo-mbridge-recipe-recommender/SKILL.md +++ b/skills/nemo-mbridge-recipe-recommender/SKILL.md @@ -1,8 +1,7 @@ --- name: nemo-mbridge-recipe-recommender license: Apache-2.0 -description: Recommend and customize Megatron Bridge recipes for a user's model, GPU count, and training goal. Indexes library recipes (pretrain/SFT/PEFT) and performance recipes. -when_to_use: User wants a starting recipe or training config; 'which recipe', 'recommend recipe', 'how to train Llama', 'starting config for X GPUs', 'what recipe for SFT'. +description: Recommend and customize Megatron Bridge recipes for a user's model, GPU count, and training goal. Indexes library recipes (pretrain/SFT/PEFT) and performance recipes, and distinguishes convergence changes, semantics-preserving execution tuning, and benchmark-only shortcuts. --- # Auto Recipe — Recipe Index & Recommendation @@ -37,6 +36,58 @@ index details: divide `num_key_value_heads`, keep TP within one node unless using NVL72-class interconnect, enable SP when TP > 1, configure CP for long context, DP is implicit, and reduce `micro_batch_size` first on OOM. +6. State whether each proposed override changes the convergence contract or + only the execution/performance mapping. Do not trade convergence semantics + for throughput without calling it a new experiment. + +## Configuration Layers and Change Control + +Separate training semantics from their hardware mapping before recommending or +tuning a recipe. + +**Convergence configuration** includes the starting checkpoint and trainable +parameters; dataset/revision/split/order/seeds; tokenizer, masking, truncation, +and packing; sequence length; global batch and token budget; objective and loss +coefficients; natural or forced MoE routing and token-dropping policy; +optimizer, LR, schedule, warmup, betas, epsilon, weight decay, clipping, and +dropout; arithmetic and optimizer-state precision; and PEFT adapter settings. +Changing one of these creates a new convergence experiment. + +**Execution/performance configuration** includes hardware count and topology; +TP/PP/VP/CP/EP/ETP/DP/SP; recompute and offload; distributed optimizer/FSDP; +communication overlap; fusions and attention backends; CUDA graphs and +compilation; checkpoint I/O; and MoE transport through all-to-all, DeepEP, or +HybridEP when the routing policy is unchanged. These settings should preserve +the objective and effective updates, although floating-point reduction order +can produce small numerical drift that still needs validation. + +Treat micro batch size and gradient accumulation as execution fingerprints. +Tune them only with fixed global batch size, global batch membership/order, +normalization, optimizer boundaries, and token budget, and validate fresh loss +sentinels for each layout. Packing, precision, forced MoE load balancing, token +dropping/capacity, and router/auxiliary loss changes are never performance-only +knobs. + +Treat mock data, forced balancing, disabled correctness checks, and timing-only +schedules as benchmark-only shortcuts. They may be appropriate in +`perf_recipes`, but their losses and checkpoints are not convergence evidence. + +For comparable model-verification recipes, choose a cohort-wide convergence +contract before tuning performance. Keep the same bounded data selection, +preprocessing, sequence length, global batch, optimizer/schedule, precision, +seeds, routing policy, optimizer-step horizon, and processed-token checkpoints +where the architectures permit. Record any necessary model-specific deviation +and do not present that result as apples-to-apples convergence evidence. +Absolute losses from different architectures or tokenizers are not directly +rankable; compare stability and trend at equal token counts. + +When a recipe's batch disagrees with the chosen convergence contract, modify +and validate the library recipe separately. A declared bounded-verification +protocol may explicitly apply the same LR, schedule, sequence, and data +overrides across a cohort, but do not make one-off convergence changes merely +to improve throughput. Conversely, first try TP/PP/CP/EP, recompute/offload, +dispatcher transport, overlap, fusion, and CUDA graphs when optimizing fit or +throughput. --- @@ -132,20 +183,22 @@ All recipes live under `src/megatron/bridge/recipes/`. Each function returns a ### Qwen3 (Dense) -| Recipe | Mode | TP | PP | CP | Sizes | -|--------|------|----|----|-----|-------| -| `qwen3_*_pretrain_config` | Pretrain | 1–8 | 1–2 | — | 600M–32B | -| `qwen3_*_sft_config` | SFT | 1–8 | 1–2 | — | 600M–32B | -| `qwen3_600m_sft_128k_config` | SFT | 1 | 1 | 8 | 600M (128K seq) | -| `qwen3_*_peft_config` | PEFT | 1 | 1 | — | 600M–32B | +| Recipe | Mode | TP | PP | CP | GPUs | Sizes / notes | +|--------|------|----|----|----|------|---------------| +| `qwen3_8b_pretrain_config` | Pretrain | 1 | 1 | — | 16 | Bounded convergence cohort | +| `qwen3_8b_sft_config` | SFT | 4 | 1 | — | 4 | 2K bounded convergence cohort | +| `qwen3_8b_sft_32k_config` | SFT | 4 | 1 | 2 | 8 | Separate 32K long-context cohort | +| `qwen3_8b_peft_config` | PEFT | 1 | 1 | — | 1 | Bounded LoRA/DoRA cohort | +| `qwen3_*_{pretrain,sft,peft}_config` | All | 1–8 | 1–2 | — | varies | Other dense sizes, 600M–32B | +| `qwen3_600m_sft_128k_config` | SFT | 1 | 1 | 8 | 8 | 600M, 128K sequence | ### Qwen3 MoE | Recipe | Mode | TP | PP | EP | CP | GPUs | |--------|------|----|----|----|----|------| -| `qwen3_30b_a3b_pretrain_config` | Pretrain | 1 | 1 | 8 | — | 8 | -| `qwen3_30b_a3b_sft_config` | SFT | 1 | 1 | 8 | — | 8 | -| `qwen3_30b_a3b_peft_config` | PEFT | 1 | 1 | 1 | — | 1 | +| `qwen3_30b_a3b_pretrain_config` | Pretrain | 1 | 1 | 16 | — | 16 | +| `qwen3_30b_a3b_sft_config` | SFT | 1 | 1 | 16 | — | 16 | +| `qwen3_30b_a3b_peft_config` | PEFT | 4 | 1 | 4 | — | 4 | | `qwen3_235b_a22b_pretrain_config` | Pretrain | 4 | 16 | 8 | 2 | 512+ | | `qwen3_235b_a22b_sft_config` | SFT | 4 | 8 | 8 | — | 256 | | `qwen3_235b_a22b_peft_config` | PEFT | 1 | 4 | 4 | — | 16 | @@ -198,7 +251,10 @@ All recipes live under `src/megatron/bridge/recipes/`. Each function returns a | Recipe | Mode | Notes | |--------|------|-------| -| `moonlight_16b_{pretrain,sft,peft}_config` | All | MoE EP=8 | +| `moonlight_16b_pretrain_config` | Pretrain | 16 GPUs, TP1/PP1/EP16; bounded convergence cohort | +| `moonlight_16b_sft_config` | SFT | 8 GPUs, TP4/PP2/EP4; 2K bounded convergence cohort | +| `moonlight_16b_peft_config` | PEFT | 4 GPUs, TP4/PP1/EP4; bounded LoRA/DoRA cohort | +| `moonlight_16b_sft_8k_config` | SFT | 8 GPUs, TP2/PP1/CP2/EP8; separate 8K cohort | | `olmoe_7b_{pretrain,sft,peft}_config` | All | MoE EP=8 | | `ministral3_{3b,8b,14b}_{sft,peft}_config` | SFT/PEFT | Dense | | `gpt_oss_20b_*_config` | All | MoE + FP8/MXFP8 variants | @@ -346,8 +402,11 @@ When the user's GPU count differs from the recipe default: copies and is essentially free. 6. **CP requires all-to-all or ring attention.** Check `cp_comm_type`. For GQA models, `a2a+p2p` hierarchical CP allows CP > num_kv_heads. -7. **world_size = DP × TP × PP × CP × EP.** DP is implicit. Make sure the - product of explicit parallelisms divides your total GPU count. +7. **Dense and expert meshes overlap.** Do not multiply TP and EP together. + The minimum MoE world size is `PP × max(TP × CP, EP × ETP)`. Dense DP is + `world_size / (TP × PP × CP)` and expert EDP is + `world_size / (PP × EP × ETP)`; both quotients must be integral, and the + expert count must be divisible by EP. ### Batch Size Tuning @@ -380,11 +439,10 @@ uv run python -m torch.distributed.run --nproc_per_node=8 scripts/training/run_r --recipe llama3_8b_pretrain_config \ --dataset llm-pretrain-mock -# Reduce parallelism for Qwen3-MoE 30B to fit on 4 GPUs +# Run the native 4-GPU Qwen3-MoE 30B PEFT topology uv run python -m torch.distributed.run --nproc_per_node=4 scripts/training/run_recipe.py \ - --recipe qwen3_30b_a3b_sft_config \ - --dataset llm-finetune \ - 'model.expert_model_parallel_size=4' + --recipe qwen3_30b_a3b_peft_config \ + --dataset llm-finetune # Add long context to an existing recipe uv run python -m torch.distributed.run --nproc_per_node=8 scripts/training/run_recipe.py \ @@ -410,10 +468,10 @@ uv run python -m torch.distributed.run --nproc_per_node=8 scripts/training/run_r | I want to... | Start with | GPUs needed | |---|---|---| | Try Bridge for the first time | `llama3_8b_sft_config` + mock data | 2 | -| Fine-tune a 7-8B model | `llama3_8b_sft_config` or `qwen3_8b_sft_config` | 2–8 | +| Fine-tune a 7-8B model | `llama3_8b_sft_config` or `qwen3_8b_sft_config` | 2–4 | | LoRA on 1 GPU | `llama3_8b_peft_config` or `qwen3_8b_peft_config` | 1 | | Pretrain a dense 70B | `llama3_70b_pretrain_config` | 32–64 | -| Train a small MoE | `qwen3_30b_a3b_pretrain_config` | 8 | +| Train a small MoE | `qwen3_30b_a3b_pretrain_config` | 16 | | Train a large MoE (235B+) | `qwen3_235b_a22b_pretrain_config` | 256–512 | | Benchmark throughput | Perf recipes via `run_script.py` | Varies | | Long-context training | `llama3_8b_128k_pretrain_config` or add CP override | 16+ | diff --git a/skills/nemo-mbridge-recipe-recommender/evals/evals.json b/skills/nemo-mbridge-recipe-recommender/evals/evals.json index 404392e135..f24864228e 100644 --- a/skills/nemo-mbridge-recipe-recommender/evals/evals.json +++ b/skills/nemo-mbridge-recipe-recommender/evals/evals.json @@ -1,10 +1,10 @@ [ { "id": "recipe-recommender-positive-sft-peft-smoke", - "question": "Use the nemo-mbridge-recipe-recommender skill. Recommend recipes for these exact Megatron Bridge cases: Qwen3 30B-A3B SFT on 8 GPUs, Qwen3 235B-A22B PEFT on 16 GPUs, Llama3 8B 128K pretrain, and first-time Bridge tryout. Include the entry point, datasets, library-vs-performance recipe distinction, and key adjustment rules.", + "question": "Use the nemo-mbridge-recipe-recommender skill. Recommend recipes for these exact Megatron Bridge cases: Qwen3 30B-A3B SFT on 16 GPUs, Qwen3 235B-A22B PEFT on 16 GPUs, Llama3 8B 128K pretrain, and first-time Bridge tryout. Include the entry point, datasets, library-vs-performance recipe distinction, and key adjustment rules.", "expected_skill": "nemo-mbridge-recipe-recommender", "expected_script": null, - "ground_truth": "The answer should use the recipe recommender skill. It should recommend qwen3_30b_a3b_sft_config for Qwen3 30B-A3B SFT on 8 GPUs, qwen3_235b_a22b_peft_config for Qwen3 235B-A22B PEFT on 16 GPUs, llama3_8b_128k_pretrain_config for Llama3 8B 128K pretrain, and llama3_8b_sft_config with mock data as the first-time Bridge tryout. It should name scripts/training/run_recipe.py with uv run python -m torch.distributed.run for library recipes, use llm-finetune for SFT, use llm-pretrain-mock for pretrain and the first-time mock tryout, and warn that performance recipes under scripts/performance are for upper-bound mock-data throughput rather than production training. It should include adjustment rules: TP must divide num_key_value_heads, TP should stay within a node unless using NVL72-style interconnect, SP should be true whenever TP>1, CP needs cp_comm_type and long-context variants/overrides, DP is implicit from the product of explicit parallelisms, and micro_batch_size should be reduced first on OOM.", + "ground_truth": "The answer should use the recipe recommender skill. It should recommend qwen3_30b_a3b_sft_config for Qwen3 30B-A3B SFT on 16 GPUs, qwen3_235b_a22b_peft_config for Qwen3 235B-A22B PEFT on 16 GPUs, llama3_8b_128k_pretrain_config for Llama3 8B 128K pretrain, and llama3_8b_sft_config with mock data as the first-time Bridge tryout. It should name scripts/training/run_recipe.py with uv run python -m torch.distributed.run for library recipes, use llm-finetune for SFT, use llm-pretrain-mock for pretrain and the first-time mock tryout, and warn that performance recipes under scripts/performance are for upper-bound mock-data throughput rather than production training. It should include adjustment rules: TP must divide num_key_value_heads, TP should stay within a node unless using NVL72-style interconnect, SP should be true whenever TP>1, CP needs cp_comm_type and long-context variants/overrides, world size is DP times TP times PP times CP while EP has separate expert-group divisibility constraints, and micro_batch_size should be reduced first on OOM.", "expected_behavior": [ "Read the nemo-mbridge-recipe-recommender skill before answering.", "Identify the task as recipe selection or customization.", diff --git a/skills/parity-testing/SKILL.md b/skills/parity-testing/SKILL.md index a5b9b9c37a..a72c01ce4b 100644 --- a/skills/parity-testing/SKILL.md +++ b/skills/parity-testing/SKILL.md @@ -1,7 +1,6 @@ --- name: parity-testing -description: Structured framework for verifying numerical parity of HF<->MCore weight conversions. References existing tools and the add-model-support skill. -when_to_use: Debugging weight mismatches, verifying HF↔MCore checkpoint round-trips, choosing verification tools, or investigating a commit that changed weight conversion and caused parity failures; 'weights don't match', 'parity test', 'roundtrip check', 'logit equivalence'. +description: Structured framework for exact HF↔MCore weight verification, forward-pass logit correlation, and optional strict numerical diagnostics. Use when debugging weight mismatches, verifying HF↔MCore checkpoint round-trips, choosing verification tools, or investigating conversion commits that caused parity failures. References existing tools and the add-model-support skill. --- # Parity Testing for Megatron Bridge @@ -17,8 +16,8 @@ workflow (which includes parity testing as milestones 1 and 2), see the |---|---|---|---| | All weights round-trip exactly (single GPU) | `hf_megatron_roundtrip.py` | No | First check after writing a bridge | | Weights round-trip with TP/PP/EP | `hf_megatron_roundtrip_multi_gpu.py` | Yes | After single-GPU passes | -| Forward-pass logit equivalence | `compare_hf_and_megatron/compare.py` | Yes | After round-trip passes | -| Text generation sanity | `hf_to_megatron_generate_text.py` | Yes | Large models that OOM compare.py | +| Forward-pass logit correlation | `compare_hf_and_megatron/compare.py` | Yes | After round-trip passes | +| Text generation sanity | `hf_to_megatron_generate_text.py` | Yes | Separate inference evidence | | Programmatic weight check | `weights_verification_table()` | Yes | Inside Python scripts | | VLM generation sanity | `hf_to_megatron_generate_vlm.py` | Yes | VLM models | @@ -60,10 +59,11 @@ from megatron.bridge.models.conversion.utils import weights_verification_table weights_verification_table(bridge, hf_pretrained, megatron_model) ``` -### Level 2: Forward-Pass Parity (GPU / bfloat16) +### Level 2: Forward-Pass Correlation (GPU / bfloat16) -After round-trip passes, verify that converted weights produce identical -forward-pass output. +After round-trip passes, verify that the converted model produces strongly +correlated logits and the same next-token prediction. This is a functional +correlation gate, not a claim of bitwise-identical floating-point arithmetic. ```bash # Compare logits (loads both HF and Megatron models) @@ -73,10 +73,15 @@ uv run python -m torch.distributed.run --nproc_per_node=2 \ --prompt "The capital of France is" ``` -**Expected:** Cosine similarity > 99.99%, matching next-token predictions. +**Expected:** Matching next-token predictions and cosine similarity at least +0.99. Equivalently, cosine distance (`1 - cosine_similarity`) must be at most +1%. Record numeric maximum and mean absolute logit differences for diagnosis, +but do not use either absolute difference as a pass/fail guard. A fixed +absolute threshold is not scale- or ULP-aware across BF16 logits and +implementations. For large models that OOM `compare.py` (which loads both models), use text -generation instead: +generation as a separate inference sanity check: ```bash uv run python -m torch.distributed.run --nproc_per_node=2 \ @@ -85,6 +90,9 @@ uv run python -m torch.distributed.run --nproc_per_node=2 \ --prompt "The capital of France is" --max_new_tokens 50 ``` +Generation does not measure logit correlation and therefore cannot by itself +verify the manual forward-pass item. + ### Level 3: Training Parity (optional) Verify that a few training steps produce decreasing loss. This catches @@ -94,11 +102,15 @@ with 2 layers and small dimensions. See the functional test pattern in the ## Tolerance Table -| Test Level | Dtype | Device | Max Diff | Cosine Sim | +| Test Level | Dtype | Device | Required Gate | Report Only | |---|---|---|---|---| -| Round-trip | float32 | CPU | 0.0 (exact) | 1.0 (exact) | -| Forward pass | bfloat16 | GPU | < 1e-2 | > 0.9999 | -| Forward pass | float16 | GPU | < 1e-3 | > 0.99999 | +| Round-trip | float32 | CPU | `max_diff == 0.0` and cosine `== 1.0` | None | +| Forward pass | bfloat16 | GPU | next token matches and cosine `>= 0.99` | max/mean absolute logit difference | +| Forward pass | float16 | GPU | next token matches and cosine `>= 0.99` | max/mean absolute logit difference | + +Keep stricter cosine or absolute-difference thresholds as optional diagnostic +targets when investigating arithmetic differences. Do not use them to downgrade +a model-support card after the correlation gate passes. ## Comparison Utilities @@ -151,8 +163,10 @@ When a parity test fails, follow this sequence: is wrong. Compare the TP=1 result against each TP shard. See the `nccl-contiguous-tensors` skill for NCCL-specific issues. -3. **If round-trip passes but forward pass fails** — weights loaded - correctly but the model architecture differs. Check `provider_bridge()` +3. **If round-trip passes but forward correlation fails** — the weights loaded + correctly, but the runtime architectures may differ. A failure means the + next token differs or cosine similarity is below 0.99; a large absolute + difference alone is diagnostic, not a failure. Check `provider_bridge()` config mapping (normalization, activation, RoPE, etc.). 4. **Use the debugging script template** from the `add-model-support` skill diff --git a/src/megatron/bridge/data/conversation_processing.py b/src/megatron/bridge/data/conversation_processing.py index f3c4c0ccf2..10c1194183 100644 --- a/src/megatron/bridge/data/conversation_processing.py +++ b/src/megatron/bridge/data/conversation_processing.py @@ -956,6 +956,16 @@ def infer_assistant_mask_boundary_config(processor: Any) -> AssistantMaskBoundar ("<|im_end|>",), {role: f"<|im_start|>{role}\n" for role in ("system", "developer", "user", "tool")}, ), + ( + ("<|im_assistant|>", "<|im_user|>", "<|im_system|>", "<|im_middle|>", "<|im_end|>"), + "<|im_assistant|>assistant<|im_middle|>", + "<|im_end|>", + (), + { + "system": "<|im_system|>system<|im_middle|>", + "user": "<|im_user|>user<|im_middle|>", + }, + ), (("<|turn>model", ""), "<|turn>model\n", "", (), {}), (("model", ""), "model\n", "", (), {}), ( diff --git a/src/megatron/bridge/data/loaders.py b/src/megatron/bridge/data/loaders.py index 2a31a442d0..22df5f8513 100644 --- a/src/megatron/bridge/data/loaders.py +++ b/src/megatron/bridge/data/loaders.py @@ -308,6 +308,14 @@ def worker_init_fn(_): eval_dp_group = dp_group eval_dp_rank = torch.distributed.get_rank(group=eval_dp_group) eval_dp_size = torch.distributed.get_world_size(group=eval_dp_group) + # Text SFT configs call this field ``seed`` while Megatron GPT configs call + # it ``random_seed``. Fall back to the unoffset config RNG seed so batch + # sampling never depends on the pipeline-rank-specific torch global seed. + sampler_seed = getattr(cfg.dataset, "seed", None) + if sampler_seed is None: + sampler_seed = getattr(cfg.dataset, "random_seed", None) + if sampler_seed is None: + sampler_seed = getattr(getattr(cfg, "rng", None), "seed", None) # Build dataloders. train_dataloader = build_pretraining_data_loader( @@ -325,6 +333,7 @@ def worker_init_fn(_): data_parallel_size=dp_size, global_batch_size=cfg.train.global_batch_size, drop_last=drop_last, + seed=sampler_seed, ) eval_gbs = ( cfg.validation.eval_global_batch_size @@ -351,6 +360,7 @@ def worker_init_fn(_): data_parallel_rank=eval_dp_rank, data_parallel_size=eval_dp_size, global_batch_size=eval_gbs, + seed=sampler_seed, ) elif cfg.validation.eval_iters > 0: val_dataloader_type = "cyclic" if isinstance(cfg.dataset, GPTDatasetConfig) else cfg.dataset.dataloader_type @@ -368,6 +378,7 @@ def worker_init_fn(_): data_parallel_rank=eval_dp_rank, data_parallel_size=eval_dp_size, global_batch_size=eval_gbs, + seed=sampler_seed, ) if cfg.validation.eval_iters > 0: @@ -385,6 +396,7 @@ def worker_init_fn(_): data_parallel_rank=eval_dp_rank, data_parallel_size=eval_dp_size, global_batch_size=eval_gbs, + seed=sampler_seed, ) # Flags to know if we need to do training/validation/testing. diff --git a/src/megatron/bridge/data/samplers.py b/src/megatron/bridge/data/samplers.py index 1c7670515d..4610d2f634 100644 --- a/src/megatron/bridge/data/samplers.py +++ b/src/megatron/bridge/data/samplers.py @@ -25,6 +25,7 @@ def build_pretraining_data_loader( data_parallel_size: int = 1, drop_last: Optional[bool] = True, global_batch_size: Optional[int] = None, + seed: int | None = None, ) -> Optional[DataLoader]: """Build a dataloader for pretraining. @@ -49,6 +50,9 @@ def build_pretraining_data_loader( drop_last: Whether to drop last incomplete batch. global_batch_size: Total batch size across all data parallel ranks. Required for 'batch' dataloader_type. + seed: Explicit shuffle seed for the global-batch sampler. Supplying the + dataset seed keeps pipeline stages on the same sample order even + though their model RNG seeds differ. Returns: A PyTorch DataLoader instance, or the dataset itself if dataloader_type is @@ -97,6 +101,7 @@ def build_pretraining_data_loader( data_parallel_size=data_parallel_size, drop_last=drop_last, pad_samples_to_global_batch_size=not drop_last, + seed=seed, ) elif dataloader_type == "external": # External dataloaders are passed through. User is expected to provide a diff --git a/src/megatron/bridge/models/conversion/auto_bridge.py b/src/megatron/bridge/models/conversion/auto_bridge.py index 736409ddd6..0c5e40e79b 100644 --- a/src/megatron/bridge/models/conversion/auto_bridge.py +++ b/src/megatron/bridge/models/conversion/auto_bridge.py @@ -279,6 +279,9 @@ def __init__(self, hf_pretrained: PreTrainedCausalLM | PreTrainedMaskedLM | Pret # Data type for exporting weights self.export_weight_dtype: Literal["bf16", "fp16", "fp8"] = "bf16" self.hf_model_id: Optional[str] = None + init_kwargs = getattr(hf_pretrained, "init_kwargs", {}) + revision = init_kwargs.get("revision") if isinstance(init_kwargs, dict) else None + self.hf_model_revision: str | None = revision if isinstance(revision, str) else None trust_remote_code = getattr(hf_pretrained, "trust_remote_code", False) self.trust_remote_code = trust_remote_code if isinstance(trust_remote_code, bool) else False @@ -1681,6 +1684,8 @@ def to_megatron_provider(self, load_weights: bool = True, hf_path: str | Path | if hf_identifier: setattr(provider, "hf_model_id", hf_identifier) + if hf_path is None and self.hf_model_revision: + setattr(provider, "hf_model_revision", self.hf_model_revision) return provider diff --git a/src/megatron/bridge/models/deepseek/deepseek_v3_bridge.py b/src/megatron/bridge/models/deepseek/deepseek_v3_bridge.py index a62435e2b0..015930a6e6 100644 --- a/src/megatron/bridge/models/deepseek/deepseek_v3_bridge.py +++ b/src/megatron/bridge/models/deepseek/deepseek_v3_bridge.py @@ -22,7 +22,6 @@ from megatron.bridge.models.conversion import quantization_utils from megatron.bridge.models.conversion.mapping_registry import MegatronMappingRegistry from megatron.bridge.models.conversion.model_bridge import MegatronModelBridge, WeightConversionTask -from megatron.bridge.models.conversion.transformers_compat import rope_theta_from_hf from megatron.bridge.models.deepseek.common import get_common_mapping_list from megatron.bridge.models.hf_pretrained.causal_lm import PreTrainedCausalLM from megatron.bridge.models.mla_provider import MLAModelProvider @@ -72,7 +71,6 @@ def provider_bridge(self, hf_pretrained: PreTrainedCausalLM) -> MLAModelProvider provider.moe_router_enable_expert_bias = True provider.moe_router_dtype = "fp32" provider.moe_permute_fusion = True - provider.moe_aux_loss_coeff = 0.0001 provider.apply_rope_fusion = False provider.gradient_accumulation_fusion = True @@ -87,7 +85,6 @@ def provider_bridge(self, hf_pretrained: PreTrainedCausalLM) -> MLAModelProvider provider.attention_softmax_in_fp32 = False provider.make_vocab_size_divisible_by = 1280 - provider.seq_length = 4096 provider.moe_layer_freq = [0] * hf_config.first_k_dense_replace + [1] * ( hf_config.num_hidden_layers - hf_config.first_k_dense_replace @@ -172,7 +169,7 @@ def maybe_modify_converted_hf_weight( converted_weights_dict: Dict[str, torch.Tensor], hf_state_dict: Mapping[str, torch.Tensor], ) -> Dict[str, torch.Tensor]: - """Add rotary embedding inverse frequency parameter if needed.""" + """Preserve a source rotary embedding inverse frequency tensor if present.""" global_name = task.global_param_name if not global_name.startswith("decoder.layers.") or not global_name.endswith(".input_layernorm.weight"): return converted_weights_dict @@ -181,36 +178,17 @@ def maybe_modify_converted_hf_weight( if len(parts) < 4 or not parts[2].isdigit(): return converted_weights_dict - inv_freq_prefix = "model.layers." - inv_freq_suffix = ".self_attn.rotary_emb.inv_freq" layer_idx = int(parts[2]) - inv_freq_key = f"{inv_freq_prefix}{layer_idx}{inv_freq_suffix}" + inv_freq_key = f"model.layers.{layer_idx}.self_attn.rotary_emb.inv_freq" if inv_freq_key in converted_weights_dict: return converted_weights_dict - has_inv_freq = getattr(self, "_deepseek_has_inv_freq", None) - if has_inv_freq is None: - has_inv_freq = False - for key in hf_state_dict.keys(): - if key.startswith(inv_freq_prefix) and key.endswith(inv_freq_suffix): - has_inv_freq = True - break - self._deepseek_has_inv_freq = has_inv_freq - if not has_inv_freq: + source_inv_freq = hf_state_dict.get(inv_freq_key) + if source_inv_freq is not None: + if converted_weights_dict: + reference_tensor = next(iter(converted_weights_dict.values())) + source_inv_freq = source_inv_freq.to(device=reference_tensor.device) + converted_weights_dict[inv_freq_key] = source_inv_freq return converted_weights_dict - inv_freq = getattr(self, "_deepseek_inv_freq", None) - if inv_freq is None: - rotary_dim = self.hf_config.qk_rope_head_dim - rotary_base = rope_theta_from_hf(self.hf_config) - inv_freq = 1.0 / (rotary_base ** (torch.arange(0, rotary_dim, 2, dtype=torch.float32) / rotary_dim)) - self._deepseek_inv_freq = inv_freq - - if converted_weights_dict: - reference_tensor = next(iter(converted_weights_dict.values())) - if inv_freq.device != reference_tensor.device: - inv_freq = inv_freq.to(device=reference_tensor.device) - self._deepseek_inv_freq = inv_freq - - converted_weights_dict[inv_freq_key] = inv_freq return converted_weights_dict diff --git a/src/megatron/bridge/models/gpt_provider.py b/src/megatron/bridge/models/gpt_provider.py index e86a6e6d96..6bb4a9145c 100644 --- a/src/megatron/bridge/models/gpt_provider.py +++ b/src/megatron/bridge/models/gpt_provider.py @@ -172,6 +172,9 @@ class GPTModelProvider(TransformerConfig, ModelProviderMixin[MCoreGPTModel]): hf_model_id: str | None = None """Optional HuggingFace model identifier associated with this provider.""" + hf_model_revision: str | None = None + """Optional immutable HuggingFace revision used to construct this provider.""" + # This represents the unpadded vocab size # The padded vocab size is automatically calculated in the provide() method. vocab_size: Optional[int] = None diff --git a/src/megatron/bridge/models/hf_pretrained/base.py b/src/megatron/bridge/models/hf_pretrained/base.py index 84e75c51c2..8a32767f2b 100644 --- a/src/megatron/bridge/models/hf_pretrained/base.py +++ b/src/megatron/bridge/models/hf_pretrained/base.py @@ -211,7 +211,28 @@ def save_artifacts( for name in self.OPTIONAL_ARTIFACTS: artifact = getattr(self, name, None) if artifact is not None and hasattr(artifact, "save_pretrained"): - artifact.save_pretrained(save_path) + try: + artifact.save_pretrained(save_path) + except ValueError: + if name != "generation_config": + raise + + source_path = original_source_path or getattr(self, "model_name_or_path", None) + if source_path is None: + raise + + copied_files = self._copy_custom_modeling_files( + source_path=source_path, + target_path=save_path, + file_patterns=["generation_config.json"], + ) + if "generation_config.json" not in copied_files: + raise + + logger.warning( + "Generation config validation failed during export; preserving the source " + "generation_config.json unchanged" + ) # Download/copy additional files if specified if additional_files: diff --git a/src/megatron/bridge/models/hybrid/hybrid_provider.py b/src/megatron/bridge/models/hybrid/hybrid_provider.py index b4d01afabf..9f1e69679b 100644 --- a/src/megatron/bridge/models/hybrid/hybrid_provider.py +++ b/src/megatron/bridge/models/hybrid/hybrid_provider.py @@ -12,6 +12,7 @@ # See the License for the specific language governing permissions and # limitations under the License. +import copy import inspect import logging import warnings @@ -42,6 +43,19 @@ logger = logging.getLogger(__name__) +DEFAULT_MAMBA_CHUNK_SIZE = 128 + + +def _configure_mamba_chunk_size(stack_spec: ModuleSpec, chunk_size: int) -> ModuleSpec: + """Return a stack spec whose Mamba mixer uses the requested scan chunk size.""" + if chunk_size == DEFAULT_MAMBA_CHUNK_SIZE: + return stack_spec + + configured_spec = copy.deepcopy(stack_spec) + mixer_spec = configured_spec.submodules.mamba_layer.submodules.mixer + mixer_spec.params = {**mixer_spec.params, "chunk_size": chunk_size} + return configured_spec + def modelopt_hybrid_stack_spec(config: "HybridModelProvider | None" = None) -> ModuleSpec: """Hybrid stack specification for quantization with ModelOpt. @@ -107,6 +121,7 @@ class HybridModelProvider(TransformerConfig, ModelProviderMixin[MCoreHybridModel bf16: bool = True num_layers: int | None = None mamba_num_groups: int = 8 + mamba_chunk_size: int = DEFAULT_MAMBA_CHUNK_SIZE num_attention_heads: int = 1 hybrid_attention_ratio: float = 0.0 hybrid_mlp_ratio: float = 0.0 @@ -144,6 +159,11 @@ class HybridModelProvider(TransformerConfig, ModelProviderMixin[MCoreHybridModel vocab_size: int | None = None should_pad_vocab: bool = False hf_model_id: str | None = None + """Optional HuggingFace model identifier associated with this provider.""" + + hf_model_revision: str | None = None + """Optional immutable HuggingFace revision used to construct this provider.""" + _pg_collection: ProcessGroupCollection | None = None # MTP @@ -151,8 +171,6 @@ class HybridModelProvider(TransformerConfig, ModelProviderMixin[MCoreHybridModel mtp_hybrid_override_pattern: str | None = None keep_mtp_spec_in_bf16: bool = False - """Optional HuggingFace model identifier associated with this provider.""" - # If True, restore modelopt_state that contains quantization, sparsity, and speculative decoding state. restore_modelopt_state: bool = False @@ -162,6 +180,9 @@ def finalize(self) -> None: Calculates the number of layers from ``hybrid_layer_pattern`` and executes the deferred MCore post-init logic. """ + if self.mamba_chunk_size < 1: + raise ValueError("mamba_chunk_size must be at least 1.") + # Check if hybrid_override_pattern is specified and throw deprecation warning. used_hybrid_override_pattern = False if self.hybrid_override_pattern is not None: @@ -240,12 +261,15 @@ def _resolve_hybrid_stack_spec(self) -> ModuleSpec: """Resolve the configured Hybrid stack spec.""" hybrid_stack_spec = self.hybrid_stack_spec if hybrid_stack_spec is None: - return get_default_hybrid_stack_spec(self) - if not isinstance(hybrid_stack_spec, ModuleSpec): + resolved_spec = get_default_hybrid_stack_spec(self) + elif not isinstance(hybrid_stack_spec, ModuleSpec): if len(inspect.signature(hybrid_stack_spec).parameters) > 0: - return hybrid_stack_spec(self) - return hybrid_stack_spec() - return hybrid_stack_spec + resolved_spec = hybrid_stack_spec(self) + else: + resolved_spec = hybrid_stack_spec() + else: + resolved_spec = hybrid_stack_spec + return _configure_mamba_chunk_size(resolved_spec, self.mamba_chunk_size) def provide(self, pre_process=None, post_process=None, vp_stage=None) -> MCoreHybridModel: """Configure and instantiate a Megatron Core Hybrid model based on this configuration. diff --git a/src/megatron/bridge/models/nemotronh/nemotron_h_bridge.py b/src/megatron/bridge/models/nemotronh/nemotron_h_bridge.py index 309bb2508d..6e2e2d36a7 100644 --- a/src/megatron/bridge/models/nemotronh/nemotron_h_bridge.py +++ b/src/megatron/bridge/models/nemotronh/nemotron_h_bridge.py @@ -238,6 +238,7 @@ class NemotronHBridge(MegatronModelBridge): # Mamba-specific fields ("mamba_head_dim", "mamba_head_dim"), ("mamba_num_heads", "mamba_num_heads"), + ("chunk_size", "mamba_chunk_size"), ("n_groups", "mamba_num_groups"), ("ssm_state_size", "mamba_state_dim"), ("hybrid_override_pattern", "hybrid_layer_pattern"), diff --git a/src/megatron/bridge/models/qwen/qwen3_moe_bridge.py b/src/megatron/bridge/models/qwen/qwen3_moe_bridge.py index a6476f4f49..23666f66fa 100755 --- a/src/megatron/bridge/models/qwen/qwen3_moe_bridge.py +++ b/src/megatron/bridge/models/qwen/qwen3_moe_bridge.py @@ -1,4 +1,4 @@ -# Copyright (c) 2025, NVIDIA CORPORATION. All rights reserved. +# Copyright (c) 2025-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. diff --git a/src/megatron/bridge/perf_recipes/qwen/common.py b/src/megatron/bridge/perf_recipes/qwen/common.py index 50779ec84c..d56a891899 100644 --- a/src/megatron/bridge/perf_recipes/qwen/common.py +++ b/src/megatron/bridge/perf_recipes/qwen/common.py @@ -19,10 +19,10 @@ _enable_overlap_param_gather_with_optimizer_step, _perf_precision, ) -from megatron.bridge.recipes.qwen.qwen3_moe import ( - qwen3_30b_a3b_pretrain_config, - qwen3_235b_a22b_pretrain_config, +from megatron.bridge.recipes.qwen.h100.qwen3_moe import ( + qwen3_30b_a3b_pretrain_8gpu_h100_bf16_config as qwen3_30b_a3b_pretrain_config, ) +from megatron.bridge.recipes.qwen.qwen3_moe import qwen3_235b_a22b_pretrain_config from megatron.bridge.recipes.qwen.qwen3_next import qwen3_next_80b_a3b_pretrain_config from megatron.bridge.training.comm_overlap import CommOverlapConfig from megatron.bridge.training.config import ConfigContainer diff --git a/src/megatron/bridge/perf_recipes/qwen/h100/qwen3_moe.py b/src/megatron/bridge/perf_recipes/qwen/h100/qwen3_moe.py index 55740443a3..9205bbca14 100644 --- a/src/megatron/bridge/perf_recipes/qwen/h100/qwen3_moe.py +++ b/src/megatron/bridge/perf_recipes/qwen/h100/qwen3_moe.py @@ -25,39 +25,17 @@ qwen3_235b_a22b_pretrain_config, qwen3_next_80b_a3b_pretrain_config, ) +from megatron.bridge.recipes.qwen.h100.qwen3_moe import ( + qwen3_30b_a3b_pretrain_16gpu_h100_bf16_config as _qwen3_30b_a3b_default_pretrain_config, +) def qwen3_30b_a3b_pretrain_16gpu_h100_bf16_config() -> ConfigContainer: """Qwen3 30B-A3B pretrain: 16× H100, BF16, EP=16.""" - cfg = qwen3_30b_a3b_pretrain_config() + cfg = _qwen3_30b_a3b_default_pretrain_config() cfg.mixed_precision = _perf_precision("bf16") - cfg.model.bias_activation_fusion = True - cfg.model.recompute_granularity = None - cfg.model.recompute_method = None - cfg.model.recompute_num_layers = None - cfg.model.moe_router_fusion = True - cfg.model.seq_length = 4096 - cfg.dataset.seq_length = 4096 cfg.model.moe_router_force_load_balancing = True - cfg.model.tensor_model_parallel_size = 1 - cfg.model.pipeline_model_parallel_size = 1 - cfg.model.context_parallel_size = 1 - cfg.model.virtual_pipeline_model_parallel_size = None - cfg.model.expert_model_parallel_size = 16 - cfg.model.sequence_parallel = False - cfg.train.global_batch_size = 1024 - cfg.train.micro_batch_size = 1 - - cfg.model.moe_flex_dispatcher_backend = "hybridep" - cfg.model.moe_token_dispatcher_type = "flex" - cfg.model.moe_a2a_overlap = False - - cfg.model.cuda_graph_impl = "transformer_engine" - cfg.model.cuda_graph_scope = ["moe_router", "moe_preprocess"] - - cfg.comm_overlap = CommOverlapConfig(tp_comm_overlap=True) - _benchmark_common(cfg) # Keep process settings next to the recipe so users can see the exact benchmark environment. cfg.env_vars = { diff --git a/src/megatron/bridge/recipes/moonlight/__init__.py b/src/megatron/bridge/recipes/moonlight/__init__.py index d3d06919dd..c86af2903c 100644 --- a/src/megatron/bridge/recipes/moonlight/__init__.py +++ b/src/megatron/bridge/recipes/moonlight/__init__.py @@ -1,4 +1,4 @@ -# Copyright (c) 2025, NVIDIA CORPORATION. All rights reserved. +# 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. @@ -15,6 +15,7 @@ from megatron.bridge.recipes.moonlight.moonlight_16b import ( moonlight_16b_peft_config, moonlight_16b_pretrain_config, + moonlight_16b_sft_8k_config, moonlight_16b_sft_config, ) @@ -22,5 +23,6 @@ __all__ = [ "moonlight_16b_pretrain_config", "moonlight_16b_sft_config", + "moonlight_16b_sft_8k_config", "moonlight_16b_peft_config", ] diff --git a/src/megatron/bridge/recipes/moonlight/h100/__init__.py b/src/megatron/bridge/recipes/moonlight/h100/__init__.py index 5435278c1f..adae78bda5 100644 --- a/src/megatron/bridge/recipes/moonlight/h100/__init__.py +++ b/src/megatron/bridge/recipes/moonlight/h100/__init__.py @@ -17,6 +17,10 @@ __all__ = [ "moonlight_16b_peft_2gpu_h100_bf16_config", + "moonlight_16b_peft_4gpu_h100_bf16_config", + "moonlight_16b_pretrain_16gpu_h100_bf16_config", "moonlight_16b_pretrain_8gpu_h100_bf16_config", + "moonlight_16b_sft_8gpu_h100_bf16_8k_config", "moonlight_16b_sft_8gpu_h100_bf16_config", + "moonlight_16b_sft_8gpu_h100_bf16_tp1_config", ] diff --git a/src/megatron/bridge/recipes/moonlight/h100/moonlight_16b.py b/src/megatron/bridge/recipes/moonlight/h100/moonlight_16b.py index d041cdd3c4..3c4125145b 100644 --- a/src/megatron/bridge/recipes/moonlight/h100/moonlight_16b.py +++ b/src/megatron/bridge/recipes/moonlight/h100/moonlight_16b.py @@ -14,7 +14,6 @@ import torch -import torch.nn.functional as F from megatron.bridge import AutoBridge from megatron.bridge.models.mla_provider import MLAModelProvider @@ -27,6 +26,110 @@ from megatron.bridge.training.mixed_precision import MixedPrecisionConfig +_MOONLIGHT_16B_MODEL_ID = "moonshotai/Moonlight-16B-A3B" +_MOONLIGHT_16B_MODEL_REVISION = "476b36a473d4467f94469414bef6cee75c9c8172" # pragma: allowlist secret +_MOONLIGHT_16B_FINETUNING_UNPADDED_VOCAB_SIZE = 163842 + + +def _moonlight_16b_model_provider() -> MLAModelProvider: + """Build the Moonlight architecture from its Hugging Face configuration.""" + return AutoBridge.from_hf_pretrained( + _MOONLIGHT_16B_MODEL_ID, revision=_MOONLIGHT_16B_MODEL_REVISION + ).to_megatron_provider(load_weights=False) + + +def _moonlight_16b_finetuning_model_provider( + *, + tensor_parallel_size: int, + pipeline_parallel_size: int, + context_parallel_size: int, + expert_parallel_size: int, + sequence_parallel: bool, +) -> MLAModelProvider: + """Build a checkpoint-compatible Moonlight provider for SFT or PEFT.""" + model = _moonlight_16b_model_provider() + + # Preserve all 163842 finetuning-token rows, then explicitly pad to the TP + # multiple because setup treats a preset model vocab as already padded. + model.vocab_size = ( + (_MOONLIGHT_16B_FINETUNING_UNPADDED_VOCAB_SIZE + tensor_parallel_size - 1) // tensor_parallel_size + ) * tensor_parallel_size + + model.tensor_model_parallel_size = tensor_parallel_size + model.pipeline_model_parallel_size = pipeline_parallel_size + model.pipeline_dtype = torch.bfloat16 + model.virtual_pipeline_model_parallel_size = None + model.context_parallel_size = context_parallel_size + model.expert_model_parallel_size = expert_parallel_size + model.expert_tensor_parallel_size = 1 + model.sequence_parallel = sequence_parallel + + model.recompute_granularity = "selective" + model.recompute_modules = None + model.recompute_method = None + model.recompute_num_layers = None + + # Preserve the finetuning recipe's source-model auxiliary-loss value. + model.moe_aux_loss_coeff = 0.001 + return model + + +def _apply_moonlight_16b_finetuning_convergence_contract(cfg: ConfigContainer) -> None: + """Apply the shared bounded SFT/PEFT convergence contract.""" + cfg.train.train_iters = 100 + cfg.train.global_batch_size = 32 + cfg.train.micro_batch_size = 1 + cfg.dataset.seed = 1234 + cfg.rng.seed = 5678 + cfg.train.manual_gc = True + cfg.train.manual_gc_interval = 100 + cfg.train.manual_gc_eval = 100 + + cfg.validation.eval_interval = 50 + cfg.validation.eval_iters = 10 + + cfg.optimizer.adam_beta1 = 0.9 + cfg.optimizer.adam_beta2 = 0.95 + cfg.optimizer.adam_eps = 1e-8 + cfg.optimizer.weight_decay = 0.1 + cfg.optimizer.clip_grad = 1.0 + cfg.optimizer.use_distributed_optimizer = True + cfg.optimizer.use_precision_aware_optimizer = False + cfg.optimizer.main_params_dtype = torch.float32 + cfg.optimizer.main_grads_dtype = torch.float32 + cfg.optimizer.exp_avg_dtype = torch.float32 + cfg.optimizer.exp_avg_sq_dtype = torch.float32 + + cfg.scheduler.lr_warmup_iters = 10 + cfg.scheduler.lr_decay_iters = cfg.train.train_iters + cfg.scheduler.lr_warmup_init = 0.0 + cfg.scheduler.start_weight_decay = 0.033 + cfg.scheduler.end_weight_decay = 0.033 + cfg.scheduler.weight_decay_incr_style = "constant" + cfg.scheduler.lr_decay_style = "cosine" + + cfg.mixed_precision = MixedPrecisionConfig( + bf16=True, + params_dtype=torch.bfloat16, + pipeline_dtype=torch.bfloat16, + autocast_enabled=False, + grad_reduce_in_fp32=True, + ) + + cfg.checkpoint.save_interval = cfg.train.train_iters + cfg.checkpoint.load = None + cfg.checkpoint.ckpt_format = "torch_dist" + cfg.checkpoint.fully_parallel_save = True + cfg.checkpoint.async_save = False + + cfg.ddp.check_for_nan_in_grad = True + cfg.ddp.grad_reduce_in_fp32 = True + cfg.ddp.overlap_grad_reduce = True + cfg.ddp.overlap_param_gather = True + cfg.ddp.average_in_collective = True + cfg.ddp.use_distributed_optimizer = True + + def _get_moonlight_pipeline_layout(pp_size: int, vp_size: int): """Get pipeline layout for Moonlight-16B based on PP and VP size.""" map_pp_vp_to_layout = { @@ -57,10 +160,7 @@ def moonlight_16b_pretrain_8gpu_h100_bf16_config() -> ConfigContainer: cfg = _pretrain_common() # Model config via AutoBridge (dispatches to DeepSeekV3Bridge) - cfg.model = AutoBridge.from_hf_pretrained("moonshotai/Moonlight-16B-A3B").to_megatron_provider(load_weights=False) - # TEMPFIX(yuya): Moonlight has no Q LoRA compression (HF q_lora_rank=null), - # but CONFIG_MAPPING skips None so MLATransformerConfig default (512) would be used - cfg.model.q_lora_rank = None + cfg.model = _moonlight_16b_model_provider() # Tokenizer - uses NullTokenizer with model vocab_size cfg.tokenizer.tokenizer_type = "NullTokenizer" @@ -186,77 +286,98 @@ def moonlight_16b_pretrain_8gpu_h100_bf16_config() -> ConfigContainer: return cfg +def moonlight_16b_pretrain_16gpu_h100_bf16_config() -> ConfigContainer: + """Return the bounded-convergence pre-training config for Moonlight-16B. + + Recommended parallelism: TP=1, PP=1, CP=1, EP=16 on 16 H100 GPUs. + The 100-step schedule uses the same batch, accumulation, precision, and + optimizer fingerprint as ``qwen3_30b_a3b_convergence_v1`` while retaining + Moonlight's model-native routing configuration. + + Returns: + ConfigContainer with the Moonlight-16B pre-training contract. + """ + cfg = moonlight_16b_pretrain_8gpu_h100_bf16_config() + + cfg.model.tensor_model_parallel_size = 1 + cfg.model.pipeline_model_parallel_size = 1 + cfg.model.pipeline_model_parallel_layout = _get_moonlight_pipeline_layout(1, 1) + cfg.model.virtual_pipeline_model_parallel_size = None + cfg.model.context_parallel_size = 1 + cfg.model.expert_model_parallel_size = 16 + cfg.model.expert_tensor_parallel_size = 1 + cfg.model.sequence_parallel = False + + cfg.train.train_iters = 100 + cfg.train.global_batch_size = 1024 + cfg.train.micro_batch_size = 1 + cfg.dataset.random_seed = 1234 + cfg.rng.seed = 1234 + cfg.scheduler.lr_warmup_iters = 40 + cfg.scheduler.lr_decay_iters = cfg.train.train_iters + cfg.scheduler.lr_warmup_init = 0.0 + cfg.checkpoint.save_interval = 50 + cfg.checkpoint.load = None + + cfg.optimizer.adam_beta1 = 0.9 + cfg.optimizer.adam_beta2 = 0.95 + cfg.optimizer.adam_eps = 1e-8 + cfg.optimizer.weight_decay = 0.1 + cfg.optimizer.lr = 3e-4 + cfg.optimizer.min_lr = 3e-5 + cfg.optimizer.clip_grad = 1.0 + cfg.optimizer.use_distributed_optimizer = True + cfg.optimizer.use_precision_aware_optimizer = True + cfg.optimizer.main_params_dtype = torch.float32 + cfg.optimizer.main_grads_dtype = torch.float32 + cfg.optimizer.exp_avg_dtype = torch.float32 + cfg.optimizer.exp_avg_sq_dtype = torch.float32 + + cfg.scheduler.start_weight_decay = 0.033 + cfg.scheduler.end_weight_decay = 0.033 + cfg.scheduler.weight_decay_incr_style = "constant" + cfg.scheduler.lr_decay_style = "cosine" + + cfg.mixed_precision = MixedPrecisionConfig( + bf16=True, + params_dtype=torch.bfloat16, + pipeline_dtype=torch.bfloat16, + autocast_enabled=False, + grad_reduce_in_fp32=False, + ) + cfg.ddp.grad_reduce_in_fp32 = False + + # Keep the complete process environment visible on the recipe. + cfg.env_vars = { + **COMMON_RECIPE_ENV_VARS, + } + return cfg + + # ============================================================================= # SFT Config # ============================================================================= def moonlight_16b_sft_8gpu_h100_bf16_config() -> ConfigContainer: - """Return a full SFT config for Moonlight-16B. + """Return the legacy TP4/PP2 full SFT config for Moonlight-16B. - Default parallelism: TP=2, PP=1, EP=8, SP=True + Recommended parallelism: TP=4, PP=2, CP=1, EP=4 on 8 H100 GPUs. + Sequence parallelism resolves offline packing to ``pad_seq_to_mult=4``, so + this remains a support recipe rather than the shared pad-1 convergence + cohort. Returns: ConfigContainer with all settings pre-configured for Moonlight-16B SFT. """ cfg = _sft_common() - # Model config - uses MLAModelProvider with Moonlight-16B architecture - cfg.model = MLAModelProvider( - # Architecture - num_layers=27, - hidden_size=2048, - ffn_hidden_size=11264, - num_attention_heads=16, - kv_channels=16, - q_lora_rank=None, - kv_lora_rank=512, - num_moe_experts=64, - moe_ffn_hidden_size=1408, - moe_shared_expert_intermediate_size=2816, - moe_layer_freq=[0] * 1 + [1] * 26, - moe_router_topk=6, - moe_router_num_groups=1, - moe_router_group_topk=1, - moe_router_topk_scaling_factor=2.446, - moe_aux_loss_coeff=0.001, - make_vocab_size_divisible_by=1280, - moe_router_score_function="sigmoid", - moe_router_enable_expert_bias=True, - rotary_scaling_factor=1.0, - mscale=1.0, - mscale_all_dim=1.0, - rotary_base=50000, - layernorm_epsilon=1e-5, - init_method_std=0.02, - moe_router_bias_update_rate=1e-3, - rotary_percent=1.0, - vocab_size=163842, - # Common defaults - normalization="RMSNorm", - activation_func=F.silu, - gated_linear_unit=True, - position_embedding_type="rope", - add_bias_linear=False, - share_embeddings_and_output_weights=False, - qk_layernorm=True, - bf16=True, - params_dtype=torch.bfloat16, - moe_grouped_gemm=True, - moe_token_dispatcher_type="alltoall", - # Parallelism - tensor_model_parallel_size=2, - pipeline_model_parallel_size=1, - pipeline_dtype=torch.bfloat16, - virtual_pipeline_model_parallel_size=None, + cfg.model = _moonlight_16b_finetuning_model_provider( + tensor_parallel_size=4, + pipeline_parallel_size=2, context_parallel_size=1, - expert_model_parallel_size=8, + expert_parallel_size=4, sequence_parallel=True, - expert_tensor_parallel_size=1, - recompute_granularity="selective", - recompute_modules=None, - recompute_method=None, - recompute_num_layers=None, ) # Pipeline split settings @@ -266,14 +387,15 @@ def moonlight_16b_sft_8gpu_h100_bf16_config() -> ConfigContainer: cfg.model.num_layers_in_last_pipeline_stage = None # Set pipeline layout - cfg.model.pipeline_model_parallel_layout = _get_moonlight_pipeline_layout(1, 1) + cfg.model.pipeline_model_parallel_layout = _get_moonlight_pipeline_layout(2, 1) # Sequence length - seq_length = 4096 + seq_length = 2048 cfg.model.seq_length = seq_length # Dataset config - enable_offline_packing=True by default cfg.dataset.seq_length = seq_length cfg.dataset.offline_packing_specs.packed_sequence_size = seq_length + cfg.dataset.offline_packing_specs.pad_seq_to_mult = 4 # MoE performance optimizations cfg.model.moe_permute_fusion = True @@ -331,33 +453,130 @@ def moonlight_16b_sft_8gpu_h100_bf16_config() -> ConfigContainer: # MoE Force Load Balancing cfg.model.moe_router_force_load_balancing = False - # Training config overrides - cfg.validation.eval_interval = 50 - cfg.train.manual_gc = True - cfg.train.manual_gc_interval = 5 - cfg.train.manual_gc_eval = 5 + _apply_moonlight_16b_finetuning_convergence_contract(cfg) # Adjust pad_seq_to_mult for context parallelism if cfg.model.context_parallel_size > 1: cfg.dataset.offline_packing_specs.pad_seq_to_mult = cfg.model.context_parallel_size * 2 - # Optimizer overrides - Moonlight uses specific optimizer settings + # Tokenizer - HuggingFace tokenizer with trust_remote_code + cfg.tokenizer.tokenizer_model = _MOONLIGHT_16B_MODEL_ID + cfg.tokenizer.hf_tokenizer_kwargs = { + "revision": _MOONLIGHT_16B_MODEL_REVISION, + "trust_remote_code": True, + } + # Uncomment below if using a pretrained checkpoint and provide path to the directory containing pretrained model for finetuning + # cfg.checkpoint.pretrained_checkpoint = "/path/to/checkpoint" + + # Communication overlap settings + cfg.comm_overlap = CommOverlapConfig(tp_comm_overlap=False) + cfg.comm_overlap.delay_wgrad_compute = False + cfg.comm_overlap.overlap_moe_expert_parallel_comm = False + + # Keep the complete process environment visible on the recipe. + cfg.env_vars = { + **COMMON_RECIPE_ENV_VARS, + } + return cfg + + +def moonlight_16b_sft_8gpu_h100_bf16_tp1_config() -> ConfigContainer: + """Return the bounded-convergence full SFT config for Moonlight-16B. + + Recommended parallelism is TP=1, PP=1, CP=1, EP=8 on 8 H100 GPUs. This + topology keeps sequence parallelism disabled so the shared SFT packing + contract remains ``pad_seq_to_mult=1``. With GBS=32 and MBS=1 it uses + dense DP=8 and four gradient-accumulation steps. + + Returns: + ConfigContainer with the Moonlight-16B full-SFT convergence contract. + """ + cfg = moonlight_16b_sft_8gpu_h100_bf16_config() + + cfg.model.tensor_model_parallel_size = 1 + cfg.model.pipeline_model_parallel_size = 1 + cfg.model.pipeline_model_parallel_layout = _get_moonlight_pipeline_layout(1, 1) + cfg.model.virtual_pipeline_model_parallel_size = None + cfg.model.context_parallel_size = 1 + cfg.model.expert_model_parallel_size = 8 + cfg.model.expert_tensor_parallel_size = 1 + cfg.model.sequence_parallel = False + cfg.model.vocab_size = _MOONLIGHT_16B_FINETUNING_UNPADDED_VOCAB_SIZE + + cfg.dataset.offline_packing_specs.pad_seq_to_mult = 1 + + # Plain all-to-all is the correctness-first single-node EP8 path. Clear + # inactive flex-dispatcher settings so the execution fingerprint is exact. + cfg.model.moe_token_dispatcher_type = "alltoall" + cfg.model.moe_flex_dispatcher_backend = None + cfg.model.moe_hybridep_num_sms = None + cfg.model.moe_flex_dispatcher_num_sms = None + cfg.model.moe_a2a_overlap = False + cfg.model.moe_shared_expert_overlap = False + cfg.comm_overlap = None + + # Keep the complete process environment visible on the recipe. + cfg.env_vars = { + **COMMON_RECIPE_ENV_VARS, + } + return cfg + + +def moonlight_16b_sft_8gpu_h100_bf16_8k_config() -> ConfigContainer: + """Return the separate 8K-context SFT config for Moonlight-16B. + + This recipe preserves the previously verified long-context execution and + precision cohort instead of inheriting the 2K bounded-convergence SFT + contract. Recommended parallelism is TP=2, PP=1, CP=2, EP=8 on 8 H100 + GPUs. + + Returns: + ConfigContainer with the Moonlight-16B 8K-context SFT contract. + """ + cfg = moonlight_16b_sft_8gpu_h100_bf16_config() + + cfg.model.tensor_model_parallel_size = 2 + cfg.model.pipeline_model_parallel_size = 1 + cfg.model.pipeline_model_parallel_layout = _get_moonlight_pipeline_layout(1, 1) + cfg.model.virtual_pipeline_model_parallel_size = None + cfg.model.context_parallel_size = 2 + cfg.model.expert_model_parallel_size = 8 + cfg.model.expert_tensor_parallel_size = 1 + cfg.model.sequence_parallel = True + cfg.model.vocab_size = _MOONLIGHT_16B_FINETUNING_UNPADDED_VOCAB_SIZE + + seq_length = 8192 + cfg.model.seq_length = seq_length + cfg.dataset.seq_length = seq_length + cfg.dataset.offline_packing_specs.packed_sequence_size = seq_length + cfg.dataset.offline_packing_specs.pad_seq_to_mult = cfg.model.context_parallel_size * 2 + + cfg.train.train_iters = 20 + cfg.train.global_batch_size = 128 + cfg.train.micro_batch_size = 1 + cfg.train.manual_gc_interval = 5 + cfg.train.manual_gc_eval = 5 + cfg.validation.eval_interval = 0 + cfg.validation.eval_iters = 0 + + cfg.optimizer.lr = 1e-6 + cfg.optimizer.min_lr = 0.0 + cfg.optimizer.adam_beta1 = 0.9 + cfg.optimizer.adam_beta2 = 0.98 cfg.optimizer.adam_eps = 1e-5 cfg.optimizer.weight_decay = 0.1 - - # Precision-aware optimizer settings cfg.optimizer.use_precision_aware_optimizer = True cfg.optimizer.main_params_dtype = torch.float32 cfg.optimizer.main_grads_dtype = torch.bfloat16 cfg.optimizer.exp_avg_dtype = torch.bfloat16 cfg.optimizer.exp_avg_sq_dtype = torch.bfloat16 - # Scheduler overrides - cfg.scheduler.lr_warmup_iters = 50 - cfg.scheduler.lr_decay_iters = 1000 - cfg.scheduler.min_lr = 0.0 + cfg.scheduler.lr_warmup_iters = 2 + cfg.scheduler.lr_decay_iters = cfg.train.train_iters + cfg.scheduler.lr_warmup_init = 0.0 + cfg.checkpoint.save_interval = 50 + cfg.checkpoint.load = None - # Mixed precision config - Moonlight uses MixedPrecisionConfig (not "bf16_mixed" string) cfg.mixed_precision = MixedPrecisionConfig( bf16=True, params_dtype=torch.bfloat16, @@ -365,31 +584,10 @@ def moonlight_16b_sft_8gpu_h100_bf16_config() -> ConfigContainer: autocast_enabled=False, grad_reduce_in_fp32=False, ) - - # Tokenizer - HuggingFace tokenizer with trust_remote_code - cfg.tokenizer.tokenizer_model = "moonshotai/Moonlight-16B-A3B" - cfg.tokenizer.hf_tokenizer_kwargs = {"trust_remote_code": True} - - # Checkpoint config overrides - cfg.checkpoint.save_interval = 50 - cfg.checkpoint.ckpt_format = "torch_dist" - cfg.checkpoint.fully_parallel_save = True - cfg.checkpoint.async_save = False - # Uncomment below if using a pretrained checkpoint and provide path to the directory containing pretrained model for finetuning - # cfg.checkpoint.pretrained_checkpoint = "/path/to/checkpoint" - - # DDP config - Moonlight uses grad_reduce_in_fp32=False - cfg.ddp.check_for_nan_in_grad = True + cfg.model.cross_entropy_loss_fusion = False + cfg.model.calculate_per_token_loss = True cfg.ddp.grad_reduce_in_fp32 = False - cfg.ddp.overlap_grad_reduce = True - cfg.ddp.overlap_param_gather = True - cfg.ddp.average_in_collective = True - cfg.ddp.use_distributed_optimizer = True - - # Communication overlap settings - cfg.comm_overlap = CommOverlapConfig(tp_comm_overlap=False) - cfg.comm_overlap.delay_wgrad_compute = False - cfg.comm_overlap.overlap_moe_expert_parallel_comm = False + cfg.ddp.average_in_collective = False # Keep the complete process environment visible on the recipe. cfg.env_vars = { @@ -418,62 +616,12 @@ def moonlight_16b_peft_2gpu_h100_bf16_config( """ cfg = _peft_common() - # Model config - PEFT uses TP=1, EP=2 - cfg.model = MLAModelProvider( - # Architecture - num_layers=27, - hidden_size=2048, - ffn_hidden_size=11264, - num_attention_heads=16, - kv_channels=16, - q_lora_rank=None, - kv_lora_rank=512, - num_moe_experts=64, - moe_ffn_hidden_size=1408, - moe_shared_expert_intermediate_size=2816, - moe_layer_freq=[0] * 1 + [1] * 26, - moe_router_topk=6, - moe_router_num_groups=1, - moe_router_group_topk=1, - moe_router_topk_scaling_factor=2.446, - moe_aux_loss_coeff=0.001, - make_vocab_size_divisible_by=1280, - moe_router_score_function="sigmoid", - moe_router_enable_expert_bias=True, - rotary_scaling_factor=1.0, - mscale=1.0, - mscale_all_dim=1.0, - rotary_base=50000, - layernorm_epsilon=1e-5, - init_method_std=0.02, - moe_router_bias_update_rate=1e-3, - rotary_percent=1.0, - vocab_size=163842, - # Common defaults - normalization="RMSNorm", - activation_func=F.silu, - gated_linear_unit=True, - position_embedding_type="rope", - add_bias_linear=False, - share_embeddings_and_output_weights=False, - qk_layernorm=True, - bf16=True, - params_dtype=torch.bfloat16, - moe_grouped_gemm=True, - moe_token_dispatcher_type="alltoall", - # Parallelism - tensor_model_parallel_size=1, - pipeline_model_parallel_size=1, - pipeline_dtype=torch.bfloat16, - virtual_pipeline_model_parallel_size=None, + cfg.model = _moonlight_16b_finetuning_model_provider( + tensor_parallel_size=1, + pipeline_parallel_size=1, context_parallel_size=1, - expert_model_parallel_size=2, + expert_parallel_size=2, sequence_parallel=False, - expert_tensor_parallel_size=1, - recompute_granularity="selective", - recompute_modules=None, - recompute_method=None, - recompute_num_layers=None, ) # Pipeline split settings @@ -588,8 +736,11 @@ def moonlight_16b_peft_2gpu_h100_bf16_config( ) # Tokenizer - HuggingFace tokenizer with trust_remote_code - cfg.tokenizer.tokenizer_model = "moonshotai/Moonlight-16B-A3B" - cfg.tokenizer.hf_tokenizer_kwargs = {"trust_remote_code": True} + cfg.tokenizer.tokenizer_model = _MOONLIGHT_16B_MODEL_ID + cfg.tokenizer.hf_tokenizer_kwargs = { + "revision": _MOONLIGHT_16B_MODEL_REVISION, + "trust_remote_code": True, + } # Checkpoint config overrides cfg.checkpoint.save_interval = 50 @@ -619,8 +770,66 @@ def moonlight_16b_peft_2gpu_h100_bf16_config( return cfg +def moonlight_16b_peft_4gpu_h100_bf16_config( + peft_scheme: str | PEFT = "lora", +) -> ConfigContainer: + """Return the bounded-convergence PEFT config for Moonlight-16B. + + Recommended parallelism is TP=4, PP=1, CP=1, EP=4 on 4 H100 GPUs. + The default LoRA keeps the base model frozen and targets only the attention + QKV and output projections with rank 8, alpha 16, and zero dropout. + + Args: + peft_scheme: PEFT scheme - "lora", "dora", or a custom PEFT instance. + + Returns: + ConfigContainer with the Moonlight-16B PEFT convergence contract. + """ + cfg = moonlight_16b_peft_2gpu_h100_bf16_config(peft_scheme=peft_scheme) + + cfg.model.tensor_model_parallel_size = 4 + cfg.model.pipeline_model_parallel_size = 1 + cfg.model.pipeline_model_parallel_layout = _get_moonlight_pipeline_layout(1, 1) + cfg.model.virtual_pipeline_model_parallel_size = None + cfg.model.context_parallel_size = 1 + cfg.model.expert_model_parallel_size = 4 + cfg.model.expert_tensor_parallel_size = 1 + cfg.model.sequence_parallel = True + cfg.model.vocab_size = ( + (_MOONLIGHT_16B_FINETUNING_UNPADDED_VOCAB_SIZE + cfg.model.tensor_model_parallel_size - 1) + // cfg.model.tensor_model_parallel_size + ) * cfg.model.tensor_model_parallel_size + + seq_length = 2048 + cfg.model.seq_length = seq_length + cfg.dataset.seq_length = seq_length + cfg.dataset.offline_packing_specs.packed_sequence_size = seq_length + cfg.dataset.offline_packing_specs.pad_seq_to_mult = 4 + + if isinstance(peft_scheme, str) and peft_scheme.lower() in ["lora", "dora"]: + cfg.peft = default_peft_config( + peft_scheme, + dim=8, + alpha=16, + dropout=0.0, + target_modules=["linear_q_proj", "linear_kv_down_proj", "linear_kv_up_proj", "linear_proj"], + ) + + _apply_moonlight_16b_finetuning_convergence_contract(cfg) + + # Keep the complete process environment visible on the recipe. + cfg.env_vars = { + **COMMON_RECIPE_ENV_VARS, + } + return cfg + + __all__ = [ "moonlight_16b_peft_2gpu_h100_bf16_config", + "moonlight_16b_peft_4gpu_h100_bf16_config", + "moonlight_16b_pretrain_16gpu_h100_bf16_config", "moonlight_16b_pretrain_8gpu_h100_bf16_config", + "moonlight_16b_sft_8gpu_h100_bf16_8k_config", "moonlight_16b_sft_8gpu_h100_bf16_config", + "moonlight_16b_sft_8gpu_h100_bf16_tp1_config", ] diff --git a/src/megatron/bridge/recipes/moonlight/moonlight_16b.py b/src/megatron/bridge/recipes/moonlight/moonlight_16b.py index 0ad8c48047..422bcf49f1 100644 --- a/src/megatron/bridge/recipes/moonlight/moonlight_16b.py +++ b/src/megatron/bridge/recipes/moonlight/moonlight_16b.py @@ -20,18 +20,22 @@ _get_moonlight_pipeline_layout, ) from megatron.bridge.recipes.moonlight.h100.moonlight_16b import ( - moonlight_16b_peft_2gpu_h100_bf16_config as moonlight_16b_peft_config, + moonlight_16b_peft_4gpu_h100_bf16_config as moonlight_16b_peft_config, ) from megatron.bridge.recipes.moonlight.h100.moonlight_16b import ( - moonlight_16b_pretrain_8gpu_h100_bf16_config as moonlight_16b_pretrain_config, + moonlight_16b_pretrain_16gpu_h100_bf16_config as moonlight_16b_pretrain_config, ) from megatron.bridge.recipes.moonlight.h100.moonlight_16b import ( - moonlight_16b_sft_8gpu_h100_bf16_config as moonlight_16b_sft_config, + moonlight_16b_sft_8gpu_h100_bf16_8k_config as moonlight_16b_sft_8k_config, +) +from megatron.bridge.recipes.moonlight.h100.moonlight_16b import ( + moonlight_16b_sft_8gpu_h100_bf16_tp1_config as moonlight_16b_sft_config, ) __all__ = [ "moonlight_16b_pretrain_config", "moonlight_16b_sft_config", + "moonlight_16b_sft_8k_config", "moonlight_16b_peft_config", ] diff --git a/src/megatron/bridge/recipes/nemotronh/__init__.py b/src/megatron/bridge/recipes/nemotronh/__init__.py index 91bfe4feeb..4a3ab0a63c 100644 --- a/src/megatron/bridge/recipes/nemotronh/__init__.py +++ b/src/megatron/bridge/recipes/nemotronh/__init__.py @@ -19,6 +19,12 @@ nemotron_3_nano_pretrain_config, nemotron_3_nano_sft_config, ) +from megatron.bridge.recipes.nemotronh.nemotron_3_nano_4b import ( + nemotron_3_nano_4b_peft_config, + nemotron_3_nano_4b_pretrain_config, + nemotron_3_nano_4b_sft_32k_config, + nemotron_3_nano_4b_sft_config, +) # Nemotron 3 Super models from megatron.bridge.recipes.nemotronh.nemotron_3_super import ( @@ -82,6 +88,11 @@ "nemotron_3_nano_pretrain_config", "nemotron_3_nano_sft_config", "nemotron_3_nano_peft_config", + # Nemotron 3 Nano 4B model + "nemotron_3_nano_4b_pretrain_config", + "nemotron_3_nano_4b_sft_config", + "nemotron_3_nano_4b_sft_32k_config", + "nemotron_3_nano_4b_peft_config", # Nemotron 3 Super models "nemotron_3_super_pretrain_config", "nemotron_3_super_sft_config", diff --git a/src/megatron/bridge/recipes/nemotronh/h100/__init__.py b/src/megatron/bridge/recipes/nemotronh/h100/__init__.py index a152f9675e..462b6666c8 100644 --- a/src/megatron/bridge/recipes/nemotronh/h100/__init__.py +++ b/src/megatron/bridge/recipes/nemotronh/h100/__init__.py @@ -13,6 +13,7 @@ # limitations under the License. from megatron.bridge.recipes.nemotronh.h100.nemotron_3_nano import * # noqa: F403 +from megatron.bridge.recipes.nemotronh.h100.nemotron_3_nano_4b import * # noqa: F403 from megatron.bridge.recipes.nemotronh.h100.nemotron_3_super import * # noqa: F403 from megatron.bridge.recipes.nemotronh.h100.nemotron_3_ultra import * # noqa: F403 from megatron.bridge.recipes.nemotronh.h100.nemotron_nano_v2 import * # noqa: F403 @@ -20,6 +21,10 @@ __all__ = [ + "nemotron_3_nano_4b_peft_8gpu_h100_bf16_config", + "nemotron_3_nano_4b_pretrain_8gpu_h100_bf16_config", + "nemotron_3_nano_4b_sft_8gpu_h100_bf16_32k_config", + "nemotron_3_nano_4b_sft_8gpu_h100_bf16_config", "nemotron_3_nano_peft_8gpu_h100_bf16_config", "nemotron_3_nano_pretrain_8gpu_h100_bf16_config", "nemotron_3_nano_sft_8gpu_h100_bf16_config", diff --git a/src/megatron/bridge/recipes/nemotronh/h100/nemotron_3_nano_4b.py b/src/megatron/bridge/recipes/nemotronh/h100/nemotron_3_nano_4b.py new file mode 100644 index 0000000000..1fe38ca8a0 --- /dev/null +++ b/src/megatron/bridge/recipes/nemotronh/h100/nemotron_3_nano_4b.py @@ -0,0 +1,288 @@ +# 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. + +"""Recipes for the dense NVIDIA Nemotron 3 Nano 4B model.""" + +from __future__ import annotations + +import torch +from megatron.core.activations import squared_relu + +from megatron.bridge.models.hybrid.hybrid_provider import HybridModelProvider +from megatron.bridge.peft.base import PEFT +from megatron.bridge.recipes.common import _peft_common, _pretrain_common, _sft_common +from megatron.bridge.recipes.utils.dataset_utils import default_peft_config +from megatron.bridge.recipes.utils.environment_utils import COMMON_RECIPE_ENV_VARS +from megatron.bridge.training.config import ConfigContainer +from megatron.bridge.training.mixed_precision import bf16_mixed + + +_HF_MODEL_ID = "nvidia/NVIDIA-Nemotron-3-Nano-4B-BF16" +_HF_MODEL_REVISION = "dfaf35de3e30f1867dd8dbc38a7fc9fb52d3914f" # pragma: allowlist secret + + +def _model_config(*, seq_length: int, context_parallel_size: int = 1) -> HybridModelProvider: + """Build the exact dense Nemotron 3 Nano 4B architecture.""" + return HybridModelProvider( + hybrid_layer_pattern="M-M-M-MM-M-M*-M-M*-M-M-M*-M-M-MM*-MMM-M-M-", + num_layers=42, + hidden_size=3136, + ffn_hidden_size=12544, + num_attention_heads=40, + num_query_groups=8, + kv_channels=128, + mamba_num_heads=96, + mamba_head_dim=80, + mamba_state_dim=128, + mamba_num_groups=8, + mamba_chunk_size=256, + seq_length=seq_length, + vocab_size=131072, + should_pad_vocab=False, + make_vocab_size_divisible_by=128, + share_embeddings_and_output_weights=False, + position_embedding_type="none", + activation_func=squared_relu, + gated_linear_unit=False, + add_bias_linear=False, + normalization="RMSNorm", + layernorm_epsilon=1.0e-5, + init_method_std=0.02, + hidden_dropout=0.0, + attention_dropout=0.0, + masked_softmax_fusion=True, + apply_query_key_layer_scaling=False, + persist_layer_norm=True, + attention_softmax_in_fp32=False, + first_last_layers_bf16=True, + is_hybrid_model=True, + use_mamba_mem_eff_path=True, + transformer_impl="transformer_engine", + attention_backend=None, + cross_entropy_loss_fusion=True, + cross_entropy_fusion_impl="native", + tensor_model_parallel_size=1, + pipeline_model_parallel_size=1, + pipeline_dtype=torch.bfloat16, + virtual_pipeline_model_parallel_size=None, + context_parallel_size=context_parallel_size, + sequence_parallel=False, + ) + + +def _configure_tokenizer(cfg: ConfigContainer) -> None: + """Pin the model-native tokenizer revision.""" + cfg.tokenizer.tokenizer_type = "HuggingFaceTokenizer" + cfg.tokenizer.tokenizer_model = _HF_MODEL_ID + cfg.tokenizer.hf_tokenizer_kwargs = {"revision": _HF_MODEL_REVISION, "use_fast": True} + + +def _configure_kernels(cfg: ConfigContainer) -> None: + """Set the common BF16 execution and memory policy.""" + cfg.model.pipeline_model_parallel_layout = None + cfg.model.cuda_graph_impl = "none" + cfg.model.cuda_graph_scope = "full" + cfg.model.cuda_graph_warmup_steps = 3 + cfg.model.recompute_granularity = None + cfg.model.recompute_modules = None + cfg.model.fine_grained_activation_offloading = False + cfg.model.offload_modules = None + cfg.train.manual_gc = False + cfg.train.manual_gc_interval = 0 + cfg.logger.log_interval = 1 + cfg.logger.log_throughput = True + + +def _configure_optimizer(cfg: ConfigContainer, *, lr: float, min_lr: float, grad_reduce_in_fp32: bool) -> None: + """Apply the bounded model-card convergence cohort optimizer contract.""" + cfg.optimizer.lr = lr + cfg.optimizer.min_lr = min_lr + cfg.optimizer.adam_beta1 = 0.9 + cfg.optimizer.adam_beta2 = 0.95 + cfg.optimizer.adam_eps = 1.0e-8 + cfg.optimizer.weight_decay = 0.1 + cfg.optimizer.clip_grad = 1.0 + cfg.scheduler.lr_warmup_iters = 10 + cfg.scheduler.lr_decay_iters = 100 + cfg.scheduler.lr_warmup_init = 0.0 + cfg.scheduler.start_weight_decay = 0.033 + cfg.scheduler.end_weight_decay = 0.033 + cfg.scheduler.weight_decay_incr_style = "constant" + cfg.scheduler.lr_decay_style = "cosine" + + cfg.mixed_precision = bf16_mixed() + cfg.mixed_precision.grad_reduce_in_fp32 = grad_reduce_in_fp32 + cfg.optimizer.use_precision_aware_optimizer = False + cfg.optimizer.main_grads_dtype = torch.float32 + cfg.optimizer.main_params_dtype = torch.float32 + cfg.optimizer.exp_avg_dtype = torch.float32 + cfg.optimizer.exp_avg_sq_dtype = torch.float32 + cfg.optimizer.use_distributed_optimizer = True + cfg.ddp.grad_reduce_in_fp32 = grad_reduce_in_fp32 + cfg.ddp.check_for_nan_in_grad = True + cfg.ddp.use_distributed_optimizer = True + + +def nemotron_3_nano_4b_pretrain_8gpu_h100_bf16_config() -> ConfigContainer: + """Return the bounded random-init pretraining config for eight H100 GPUs. + + Recommended parallelism is TP=1, PP=1, CP=1, DP=8. The recipe is random + initialization by default; loading the released HF checkpoint is not part of + this recipe's verification contract. + """ + cfg = _pretrain_common() + cfg.model = _model_config(seq_length=4096) + _configure_tokenizer(cfg) + _configure_kernels(cfg) + + cfg.dataset.seq_length = 4096 + cfg.dataset.random_seed = 1234 + cfg.dataset.num_dataset_builder_threads = 1 + cfg.train.train_iters = 100 + cfg.train.global_batch_size = 1024 + cfg.train.micro_batch_size = 1 + cfg.rng.seed = 1234 + cfg.validation.eval_interval = 0 + cfg.validation.eval_iters = 0 + + _configure_optimizer(cfg, lr=3.0e-4, min_lr=3.0e-5, grad_reduce_in_fp32=False) + cfg.scheduler.lr_warmup_iters = 40 + cfg.checkpoint.save_interval = 50 + cfg.checkpoint.load = None + cfg.checkpoint.dist_ckpt_strictness = "log_all" + cfg.optimizer.use_precision_aware_optimizer = True + cfg.ddp.overlap_grad_reduce = True + cfg.ddp.overlap_param_gather = True + + # Keep the complete process environment visible on the recipe. + cfg.env_vars = { + **COMMON_RECIPE_ENV_VARS, + } + return cfg + + +def nemotron_3_nano_4b_sft_8gpu_h100_bf16_config() -> ConfigContainer: + """Return the 2K packed full-SFT config for eight H100 GPUs. + + Recommended parallelism is TP=1, PP=1, CP=1, DP=8. + """ + cfg = _sft_common() + cfg.model = _model_config(seq_length=2048) + _configure_tokenizer(cfg) + _configure_kernels(cfg) + + cfg.dataset.seq_length = 2048 + cfg.dataset.seed = 1234 + cfg.dataset.offline_packing_specs.packed_sequence_size = 2048 + cfg.dataset.offline_packing_specs.pad_seq_to_mult = 1 + cfg.train.train_iters = 100 + cfg.train.global_batch_size = 32 + cfg.train.micro_batch_size = 1 + cfg.rng.seed = 5678 + cfg.validation.eval_interval = 0 + cfg.validation.eval_iters = 0 + + _configure_optimizer(cfg, lr=5.0e-6, min_lr=0.0, grad_reduce_in_fp32=True) + cfg.checkpoint.save_interval = 100 + cfg.checkpoint.load = None + cfg.checkpoint.dist_ckpt_strictness = "log_all" + cfg.ddp.overlap_grad_reduce = False + cfg.ddp.overlap_param_gather = False + + # Keep the complete process environment visible on the recipe. + cfg.env_vars = { + **COMMON_RECIPE_ENV_VARS, + } + return cfg + + +def nemotron_3_nano_4b_sft_8gpu_h100_bf16_32k_config() -> ConfigContainer: + """Return the packed 32K full-SFT config with CP=2 for eight H100 GPUs. + + Recommended parallelism is TP=1, PP=1, CP=2, DP=4. + """ + cfg = nemotron_3_nano_4b_sft_8gpu_h100_bf16_config() + cfg.model.context_parallel_size = 2 + cfg.model.cp_comm_type = "a2a" + cfg.model.seq_length = 32768 + cfg.model.cross_entropy_loss_fusion = False + cfg.model.calculate_per_token_loss = True + cfg.dataset.seq_length = 32768 + cfg.dataset.offline_packing_specs.packed_sequence_size = 32768 + cfg.dataset.offline_packing_specs.pad_seq_to_mult = 4 + cfg.train.global_batch_size = 8 + cfg.ddp.average_in_collective = False + + # Keep the complete process environment visible on the recipe. + cfg.env_vars = { + **COMMON_RECIPE_ENV_VARS, + } + return cfg + + +def nemotron_3_nano_4b_peft_8gpu_h100_bf16_config( + peft_scheme: str | PEFT = "lora", +) -> ConfigContainer: + """Return the packed attention-LoRA config for eight H100 GPUs. + + Args: + peft_scheme: PEFT scheme (``"lora"``, ``"dora"``), or a custom PEFT + instance. + + Recommended parallelism is TP=1, PP=1, CP=1, DP=8. + """ + cfg = _peft_common() + cfg.model = _model_config(seq_length=2048) + _configure_tokenizer(cfg) + _configure_kernels(cfg) + + peft_cfg = default_peft_config(peft_scheme) + if isinstance(peft_scheme, str) and peft_scheme.lower() in {"lora", "dora"}: + peft_cfg.dim = 8 + peft_cfg.alpha = 16 + peft_cfg.dropout = 0.0 + peft_cfg.target_modules = ["linear_qkv", "linear_proj"] + cfg.peft = peft_cfg + + cfg.dataset.seq_length = 2048 + cfg.dataset.seed = 1234 + cfg.dataset.offline_packing_specs.packed_sequence_size = 2048 + cfg.dataset.offline_packing_specs.pad_seq_to_mult = 4 + cfg.train.train_iters = 100 + cfg.train.global_batch_size = 32 + cfg.train.micro_batch_size = 1 + cfg.rng.seed = 5678 + cfg.validation.eval_interval = 0 + cfg.validation.eval_iters = 0 + + _configure_optimizer(cfg, lr=1.0e-4, min_lr=0.0, grad_reduce_in_fp32=True) + cfg.checkpoint.save_interval = 100 + cfg.checkpoint.load = None + cfg.checkpoint.dist_ckpt_strictness = "log_all" + cfg.ddp.overlap_grad_reduce = False + cfg.ddp.overlap_param_gather = False + + # Keep the complete process environment visible on the recipe. + cfg.env_vars = { + **COMMON_RECIPE_ENV_VARS, + } + return cfg + + +__all__ = [ + "nemotron_3_nano_4b_peft_8gpu_h100_bf16_config", + "nemotron_3_nano_4b_pretrain_8gpu_h100_bf16_config", + "nemotron_3_nano_4b_sft_8gpu_h100_bf16_32k_config", + "nemotron_3_nano_4b_sft_8gpu_h100_bf16_config", +] diff --git a/src/megatron/bridge/recipes/nemotronh/nemotron_3_nano_4b.py b/src/megatron/bridge/recipes/nemotronh/nemotron_3_nano_4b.py new file mode 100644 index 0000000000..1b47eefbb5 --- /dev/null +++ b/src/megatron/bridge/recipes/nemotronh/nemotron_3_nano_4b.py @@ -0,0 +1,38 @@ +# 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. +# ruff: noqa: F401 +"""Compatibility aliases for Nemotron 3 Nano 4B recipes.""" + +from __future__ import annotations + +from megatron.bridge.recipes.nemotronh.h100.nemotron_3_nano_4b import ( + nemotron_3_nano_4b_peft_8gpu_h100_bf16_config as nemotron_3_nano_4b_peft_config, +) +from megatron.bridge.recipes.nemotronh.h100.nemotron_3_nano_4b import ( + nemotron_3_nano_4b_pretrain_8gpu_h100_bf16_config as nemotron_3_nano_4b_pretrain_config, +) +from megatron.bridge.recipes.nemotronh.h100.nemotron_3_nano_4b import ( + nemotron_3_nano_4b_sft_8gpu_h100_bf16_32k_config as nemotron_3_nano_4b_sft_32k_config, +) +from megatron.bridge.recipes.nemotronh.h100.nemotron_3_nano_4b import ( + nemotron_3_nano_4b_sft_8gpu_h100_bf16_config as nemotron_3_nano_4b_sft_config, +) + + +__all__ = [ + "nemotron_3_nano_4b_peft_config", + "nemotron_3_nano_4b_pretrain_config", + "nemotron_3_nano_4b_sft_32k_config", + "nemotron_3_nano_4b_sft_config", +] diff --git a/src/megatron/bridge/recipes/qwen/__init__.py b/src/megatron/bridge/recipes/qwen/__init__.py index ea92b68d5f..b228b95ab4 100644 --- a/src/megatron/bridge/recipes/qwen/__init__.py +++ b/src/megatron/bridge/recipes/qwen/__init__.py @@ -56,6 +56,7 @@ qwen3_4b_sft_config, qwen3_8b_peft_config, qwen3_8b_pretrain_config, + qwen3_8b_sft_32k_config, qwen3_8b_sft_config, qwen3_14b_peft_config, qwen3_14b_pretrain_config, @@ -134,6 +135,7 @@ "qwen3_1p7b_sft_config", "qwen3_4b_sft_config", "qwen3_8b_sft_config", + "qwen3_8b_sft_32k_config", "qwen3_14b_sft_config", "qwen3_32b_sft_config", "qwen3_600m_peft_config", diff --git a/src/megatron/bridge/recipes/qwen/h100/__init__.py b/src/megatron/bridge/recipes/qwen/h100/__init__.py index dce7993044..3dbdc53391 100644 --- a/src/megatron/bridge/recipes/qwen/h100/__init__.py +++ b/src/megatron/bridge/recipes/qwen/h100/__init__.py @@ -59,7 +59,9 @@ "qwen3_235b_a22b_pretrain_256gpu_h100_bf16_config", # pragma: allowlist secret "qwen3_235b_a22b_sft_64gpu_h100_bf16_config", "qwen3_30b_a3b_peft_4gpu_h100_bf16_config", + "qwen3_30b_a3b_pretrain_16gpu_h100_bf16_config", "qwen3_30b_a3b_pretrain_8gpu_h100_bf16_config", + "qwen3_30b_a3b_sft_16gpu_h100_bf16_config", "qwen3_30b_a3b_sft_8gpu_h100_bf16_config", "qwen3_32b_peft_1gpu_h100_bf16_config", "qwen3_32b_pretrain_16gpu_h100_bf16_config", @@ -73,8 +75,10 @@ "qwen3_600m_sft_8gpu_h100_bf16_128k_config", "qwen3_600m_sft_8gpu_h100_bf16_yarn_128k_config", "qwen3_8b_peft_1gpu_h100_bf16_config", + "qwen3_8b_pretrain_16gpu_h100_bf16_config", "qwen3_8b_pretrain_4gpu_h100_bf16_config", "qwen3_8b_sft_4gpu_h100_bf16_config", + "qwen3_8b_sft_8gpu_h100_bf16_32k_config", "qwen3_next_80b_a3b_peft_1gpu_h100_bf16_config", "qwen3_next_80b_a3b_pretrain_32gpu_h100_bf16_config", "qwen3_next_80b_a3b_sft_16gpu_h100_bf16_config", diff --git a/src/megatron/bridge/recipes/qwen/h100/qwen3.py b/src/megatron/bridge/recipes/qwen/h100/qwen3.py index 19c3ea4923..eda55c855f 100644 --- a/src/megatron/bridge/recipes/qwen/h100/qwen3.py +++ b/src/megatron/bridge/recipes/qwen/h100/qwen3.py @@ -22,6 +22,11 @@ from megatron.bridge.recipes.utils.dataset_utils import default_peft_config from megatron.bridge.recipes.utils.environment_utils import COMMON_RECIPE_ENV_VARS from megatron.bridge.training.config import ConfigContainer +from megatron.bridge.training.mixed_precision import bf16_mixed + + +_QWEN3_8B_MODEL_ID = "Qwen/Qwen3-8B" +_QWEN3_8B_MODEL_REVISION = "b968826d9c46dd6066d109eabc6255188de91218" # pragma: allowlist secret def qwen3_600m_pretrain_1gpu_h100_bf16_config() -> ConfigContainer: @@ -287,10 +292,13 @@ def qwen3_8b_pretrain_4gpu_h100_bf16_config() -> ConfigContainer: cfg = _pretrain_common() # Model config - cfg.model = AutoBridge.from_hf_pretrained("Qwen/Qwen3-8B").to_megatron_provider(load_weights=False) + cfg.model = AutoBridge.from_hf_pretrained( + _QWEN3_8B_MODEL_ID, revision=_QWEN3_8B_MODEL_REVISION + ).to_megatron_provider(load_weights=False) # Tokenizer - cfg.tokenizer.tokenizer_model = "Qwen/Qwen3-8B" + cfg.tokenizer.tokenizer_model = _QWEN3_8B_MODEL_ID + cfg.tokenizer.hf_tokenizer_kwargs = {"revision": _QWEN3_8B_MODEL_REVISION} # Dataset config - mock data by default cfg.dataset.blend = None # Pass the path to the dataset here if not using mock data, along with weight. Ex: (["path/to/data1"], 0.2), [("path/to/data2", 0.8)] @@ -364,6 +372,61 @@ def qwen3_8b_pretrain_4gpu_h100_bf16_config() -> ConfigContainer: return cfg +def qwen3_8b_pretrain_16gpu_h100_bf16_config() -> ConfigContainer: + """Return a convergence-cohort pre-training config for Qwen3 8B on 16 H100 GPUs. + + Recommended parallelism: TP=1, PP=1, CP=1, DP=16. + """ + cfg = qwen3_8b_pretrain_4gpu_h100_bf16_config() + + # Preserve the cohort's per-rank micro batch while using DP for the full global batch. + cfg.model.tensor_model_parallel_size = 1 + cfg.model.pipeline_model_parallel_size = 1 + cfg.model.context_parallel_size = 1 + cfg.model.virtual_pipeline_model_parallel_size = None + cfg.model.sequence_parallel = False + cfg.train.train_iters = 100 + cfg.train.global_batch_size = 1024 + cfg.train.micro_batch_size = 1 + cfg.dataset.random_seed = 1234 + cfg.rng.seed = 1234 + + cfg.optimizer.lr = 3.0e-4 + cfg.optimizer.min_lr = 3.0e-5 + cfg.optimizer.adam_beta1 = 0.9 + cfg.optimizer.adam_beta2 = 0.95 + cfg.optimizer.adam_eps = 1.0e-8 + cfg.optimizer.weight_decay = 0.1 + cfg.optimizer.clip_grad = 1.0 + cfg.scheduler.lr_warmup_iters = 40 + cfg.scheduler.lr_decay_iters = 100 + cfg.scheduler.lr_warmup_init = 0.0 + cfg.scheduler.start_weight_decay = 0.033 + cfg.scheduler.end_weight_decay = 0.033 + cfg.scheduler.weight_decay_incr_style = "constant" + cfg.scheduler.lr_decay_style = "cosine" + + # qwen3_30b_a3b_convergence_v1 precision and optimizer-state contract. + cfg.mixed_precision = bf16_mixed() + cfg.mixed_precision.grad_reduce_in_fp32 = False + cfg.ddp.grad_reduce_in_fp32 = False + cfg.optimizer.use_precision_aware_optimizer = True + cfg.optimizer.main_grads_dtype = torch.float32 + cfg.optimizer.main_params_dtype = torch.float32 + cfg.optimizer.exp_avg_dtype = torch.float32 + cfg.optimizer.exp_avg_sq_dtype = torch.float32 + cfg.ddp.use_distributed_optimizer = True + cfg.optimizer.use_distributed_optimizer = True + cfg.checkpoint.save_interval = 50 + cfg.checkpoint.load = None + + # Keep the complete process environment visible on the recipe. + cfg.env_vars = { + **COMMON_RECIPE_ENV_VARS, + } + return cfg + + def qwen3_14b_pretrain_8gpu_h100_bf16_config() -> ConfigContainer: """Return a pre-training config for Qwen3 14B. @@ -933,15 +996,18 @@ def qwen3_4b_sft_2gpu_h100_bf16_config() -> ConfigContainer: def qwen3_8b_sft_4gpu_h100_bf16_config() -> ConfigContainer: """Return a full SFT config for Qwen3 8B. - Recommended parallelism: TP=4, PP=1 (1 node, 8 GPUs) + Recommended parallelism: TP=4, PP=1 (4 GPUs total). """ cfg = _sft_common() # Model config - cfg.model = AutoBridge.from_hf_pretrained("Qwen/Qwen3-8B").to_megatron_provider(load_weights=False) + cfg.model = AutoBridge.from_hf_pretrained( + _QWEN3_8B_MODEL_ID, revision=_QWEN3_8B_MODEL_REVISION + ).to_megatron_provider(load_weights=False) # Tokenizer - cfg.tokenizer.tokenizer_model = "Qwen/Qwen3-8B" + cfg.tokenizer.tokenizer_model = _QWEN3_8B_MODEL_ID + cfg.tokenizer.hf_tokenizer_kwargs = {"revision": _QWEN3_8B_MODEL_REVISION} # Parallelism settings cfg.model.tensor_model_parallel_size = 4 @@ -955,11 +1021,13 @@ def qwen3_8b_sft_4gpu_h100_bf16_config() -> ConfigContainer: # Sequence length (2048 for packed sequences) cfg.model.seq_length = 2048 - # Global batch size is 8 for packed sequences, 128 otherwise - cfg.train.global_batch_size = 8 - # Set pad_seq_to_mult for context parallelism - if cfg.model.context_parallel_size > 1: - cfg.dataset.offline_packing_specs.pad_seq_to_mult = cfg.model.context_parallel_size * 2 + # qwen3_30b_a3b_convergence_v1 batch contract + cfg.train.train_iters = 100 + cfg.train.global_batch_size = 32 + cfg.train.micro_batch_size = 1 + cfg.dataset.seed = 1234 + cfg.rng.seed = 5678 + cfg.dataset.offline_packing_specs.pad_seq_to_mult = 1 # Training config cfg.validation.eval_interval = 30 @@ -993,15 +1061,32 @@ def qwen3_8b_sft_4gpu_h100_bf16_config() -> ConfigContainer: # cfg.mixed_precision.fp8_param_gather = False # default # cfg.mixed_precision.reuse_grad_buf_for_mxfp8_param_ag = False # default - # Optimizer precision settings + # qwen3_30b_a3b_convergence_v1 optimizer and precision contract + cfg.optimizer.adam_beta1 = 0.9 + cfg.optimizer.adam_beta2 = 0.95 + cfg.optimizer.adam_eps = 1.0e-8 + cfg.optimizer.lr = 5.0e-6 + cfg.optimizer.min_lr = 0.0 + cfg.optimizer.weight_decay = 0.1 + cfg.optimizer.clip_grad = 1.0 + cfg.scheduler.lr_warmup_iters = 10 + cfg.scheduler.lr_decay_iters = 100 + cfg.scheduler.lr_warmup_init = 0.0 + cfg.scheduler.start_weight_decay = 0.033 + cfg.scheduler.end_weight_decay = 0.033 + cfg.scheduler.weight_decay_incr_style = "constant" + cfg.scheduler.lr_decay_style = "cosine" cfg.optimizer.use_precision_aware_optimizer = False cfg.optimizer.main_grads_dtype = torch.float32 cfg.optimizer.main_params_dtype = torch.float32 cfg.optimizer.exp_avg_dtype = torch.float32 cfg.optimizer.exp_avg_sq_dtype = torch.float32 + cfg.mixed_precision = bf16_mixed() + cfg.mixed_precision.grad_reduce_in_fp32 = True # Checkpoint config - cfg.checkpoint.save_interval = 50 + cfg.checkpoint.save_interval = 100 + cfg.checkpoint.load = None # cfg.checkpoint.save and cfg.checkpoint.load are set in _sft_common. To override: # cfg.checkpoint.save = "path/to/save" # cfg.checkpoint.load = "path/to/load" @@ -1009,11 +1094,45 @@ def qwen3_8b_sft_4gpu_h100_bf16_config() -> ConfigContainer: # cfg.checkpoint.pretrained_checkpoint = "/path/to/checkpoint" # DDP config - cfg.ddp.grad_reduce_in_fp32 = False + cfg.ddp.grad_reduce_in_fp32 = True cfg.ddp.overlap_grad_reduce = False cfg.ddp.overlap_param_gather = False cfg.ddp.check_for_nan_in_grad = True - cfg.ddp.use_distributed_optimizer = False + cfg.ddp.use_distributed_optimizer = True + cfg.optimizer.use_distributed_optimizer = True + + # Keep the complete process environment visible on the recipe. + cfg.env_vars = { + **COMMON_RECIPE_ENV_VARS, + } + return cfg + + +def qwen3_8b_sft_8gpu_h100_bf16_32k_config() -> ConfigContainer: + """Return the separate 32K-context full-SFT config for Qwen3 8B. + + Recommended parallelism: TP=4, PP=1, CP=2 (8 H100 GPUs total). The + long-context workload retains its verified GBS=8 execution cohort instead + of inheriting the bounded 2K SFT batch contract. + """ + cfg = qwen3_8b_sft_4gpu_h100_bf16_config() + + cfg.model.context_parallel_size = 2 + cfg.model.sequence_parallel = True + cfg.model.cp_comm_type = "a2a" + cfg.model.seq_length = 32768 + cfg.dataset.seq_length = 32768 + cfg.dataset.offline_packing_specs.packed_sequence_size = 32768 + cfg.dataset.offline_packing_specs.pad_seq_to_mult = 8 + + cfg.train.global_batch_size = 8 + cfg.train.micro_batch_size = 1 + + # CP SFT requires per-token loss accounting because some CP ranks may have + # no supervised tokens after label masking. + cfg.model.cross_entropy_loss_fusion = False + cfg.model.calculate_per_token_loss = True + cfg.ddp.average_in_collective = False # Keep the complete process environment visible on the recipe. cfg.env_vars = { @@ -1514,15 +1633,18 @@ def qwen3_8b_peft_1gpu_h100_bf16_config(peft_scheme: str | PEFT = "lora") -> Con Args: peft_scheme: PEFT scheme - 'lora', 'dora', or a PEFT instance. Default: 'lora' - Recommended parallelism: TP=1, PP=1 (1 node, 8 GPUs) + Recommended parallelism: TP=1, PP=1 (1 GPU total). """ cfg = _peft_common() # Model config - cfg.model = AutoBridge.from_hf_pretrained("Qwen/Qwen3-8B").to_megatron_provider(load_weights=False) + cfg.model = AutoBridge.from_hf_pretrained( + _QWEN3_8B_MODEL_ID, revision=_QWEN3_8B_MODEL_REVISION + ).to_megatron_provider(load_weights=False) # Tokenizer - cfg.tokenizer.tokenizer_model = "Qwen/Qwen3-8B" + cfg.tokenizer.tokenizer_model = _QWEN3_8B_MODEL_ID + cfg.tokenizer.hf_tokenizer_kwargs = {"revision": _QWEN3_8B_MODEL_REVISION} # Parallelism settings - TP=1 for PEFT (only training adapters) cfg.model.tensor_model_parallel_size = 1 @@ -1536,14 +1658,22 @@ def qwen3_8b_peft_1gpu_h100_bf16_config(peft_scheme: str | PEFT = "lora") -> Con # Sequence length (2048 for packed sequences) cfg.model.seq_length = 2048 - # PEFT config - use user-provided scheme or default to LoRA - cfg.peft = default_peft_config(peft_scheme) - - # Global batch size is 8 for packed sequences, 128 otherwise - cfg.train.global_batch_size = 8 - # Set pad_seq_to_mult for context parallelism - if cfg.model.context_parallel_size > 1: - cfg.dataset.offline_packing_specs.pad_seq_to_mult = cfg.model.context_parallel_size * 2 + # PEFT config - use the cohort LoRA/DoRA shape while preserving custom PEFT instances. + peft_cfg = default_peft_config(peft_scheme) + if isinstance(peft_scheme, str) and peft_scheme.lower() in ["lora", "dora"]: + peft_cfg.dim = 8 + peft_cfg.alpha = 16 + peft_cfg.dropout = 0.0 + peft_cfg.target_modules = ["linear_qkv", "linear_proj"] + cfg.peft = peft_cfg + + # qwen3_30b_a3b_convergence_v1 batch contract + cfg.train.train_iters = 100 + cfg.train.global_batch_size = 32 + cfg.train.micro_batch_size = 1 + cfg.dataset.seed = 1234 + cfg.rng.seed = 5678 + cfg.dataset.offline_packing_specs.pad_seq_to_mult = 4 # Training config cfg.validation.eval_interval = 30 @@ -1577,15 +1707,32 @@ def qwen3_8b_peft_1gpu_h100_bf16_config(peft_scheme: str | PEFT = "lora") -> Con # cfg.mixed_precision.fp8_param_gather = False # default # cfg.mixed_precision.reuse_grad_buf_for_mxfp8_param_ag = False # default - # Optimizer precision settings + # qwen3_30b_a3b_convergence_v1 optimizer and precision contract + cfg.optimizer.adam_beta1 = 0.9 + cfg.optimizer.adam_beta2 = 0.95 + cfg.optimizer.adam_eps = 1.0e-8 + cfg.optimizer.lr = 1.0e-4 + cfg.optimizer.min_lr = 0.0 + cfg.optimizer.weight_decay = 0.1 + cfg.optimizer.clip_grad = 1.0 + cfg.scheduler.lr_warmup_iters = 10 + cfg.scheduler.lr_decay_iters = 100 + cfg.scheduler.lr_warmup_init = 0.0 + cfg.scheduler.start_weight_decay = 0.033 + cfg.scheduler.end_weight_decay = 0.033 + cfg.scheduler.weight_decay_incr_style = "constant" + cfg.scheduler.lr_decay_style = "cosine" cfg.optimizer.use_precision_aware_optimizer = False cfg.optimizer.main_grads_dtype = torch.float32 cfg.optimizer.main_params_dtype = torch.float32 cfg.optimizer.exp_avg_dtype = torch.float32 cfg.optimizer.exp_avg_sq_dtype = torch.float32 + cfg.mixed_precision = bf16_mixed() + cfg.mixed_precision.grad_reduce_in_fp32 = True # Checkpoint config - cfg.checkpoint.save_interval = 50 + cfg.checkpoint.save_interval = 100 + cfg.checkpoint.load = None # cfg.checkpoint.save and cfg.checkpoint.load are set in _peft_common. To override: # cfg.checkpoint.save = "path/to/save" # cfg.checkpoint.load = "path/to/load" @@ -1593,11 +1740,12 @@ def qwen3_8b_peft_1gpu_h100_bf16_config(peft_scheme: str | PEFT = "lora") -> Con # cfg.checkpoint.pretrained_checkpoint = "/path/to/checkpoint" # DDP config - cfg.ddp.grad_reduce_in_fp32 = False + cfg.ddp.grad_reduce_in_fp32 = True cfg.ddp.overlap_grad_reduce = False cfg.ddp.overlap_param_gather = False cfg.ddp.check_for_nan_in_grad = True - cfg.ddp.use_distributed_optimizer = False + cfg.ddp.use_distributed_optimizer = True + cfg.optimizer.use_distributed_optimizer = True # Keep the complete process environment visible on the recipe. cfg.env_vars = { @@ -1824,6 +1972,8 @@ def qwen3_32b_peft_1gpu_h100_bf16_config(peft_scheme: str | PEFT = "lora") -> Co "qwen3_600m_sft_8gpu_h100_bf16_128k_config", "qwen3_600m_sft_8gpu_h100_bf16_yarn_128k_config", "qwen3_8b_peft_1gpu_h100_bf16_config", + "qwen3_8b_pretrain_16gpu_h100_bf16_config", "qwen3_8b_pretrain_4gpu_h100_bf16_config", "qwen3_8b_sft_4gpu_h100_bf16_config", + "qwen3_8b_sft_8gpu_h100_bf16_32k_config", ] diff --git a/src/megatron/bridge/recipes/qwen/h100/qwen3_moe.py b/src/megatron/bridge/recipes/qwen/h100/qwen3_moe.py index 16b8a88d1f..99b1ab4f9c 100644 --- a/src/megatron/bridge/recipes/qwen/h100/qwen3_moe.py +++ b/src/megatron/bridge/recipes/qwen/h100/qwen3_moe.py @@ -19,11 +19,16 @@ from megatron.bridge.recipes.common import _peft_common, _pretrain_common, _sft_common from megatron.bridge.recipes.utils.dataset_utils import default_peft_config from megatron.bridge.recipes.utils.environment_utils import COMMON_RECIPE_ENV_VARS +from megatron.bridge.training.comm_overlap import CommOverlapConfig from megatron.bridge.training.config import ConfigContainer from megatron.bridge.training.flex_dispatcher_backend import apply_flex_dispatcher_backend from megatron.bridge.training.mixed_precision import bf16_mixed +_QWEN3_30B_A3B_MODEL_ID = "Qwen/Qwen3-30B-A3B" +_QWEN3_30B_A3B_MODEL_REVISION = "ad44e777bcd18fa416d9da3bd8f70d33ebb85d39" # pragma: allowlist secret + + def qwen3_30b_a3b_pretrain_8gpu_h100_bf16_config() -> ConfigContainer: """Return a pre-training config for Qwen3-30B-A3B MoE. @@ -32,10 +37,13 @@ def qwen3_30b_a3b_pretrain_8gpu_h100_bf16_config() -> ConfigContainer: cfg = _pretrain_common() # Model config - cfg.model = AutoBridge.from_hf_pretrained("Qwen/Qwen3-30B-A3B").to_megatron_provider(load_weights=False) + cfg.model = AutoBridge.from_hf_pretrained( + _QWEN3_30B_A3B_MODEL_ID, revision=_QWEN3_30B_A3B_MODEL_REVISION + ).to_megatron_provider(load_weights=False) # Tokenizer (--tokenizer-model) - cfg.tokenizer.tokenizer_model = "Qwen/Qwen3-30B-A3B" + cfg.tokenizer.tokenizer_model = _QWEN3_30B_A3B_MODEL_ID + cfg.tokenizer.hf_tokenizer_kwargs = {"revision": _QWEN3_30B_A3B_MODEL_REVISION} # Dataset config - mock data by default cfg.dataset.blend = None # Pass the path to the dataset here if not using mock data, along with weight. Ex: (["path/to/data1"], 0.2), [("path/to/data2", 0.8)] @@ -135,6 +143,88 @@ def qwen3_30b_a3b_pretrain_8gpu_h100_bf16_config() -> ConfigContainer: return cfg +def qwen3_30b_a3b_pretrain_16gpu_h100_bf16_config() -> ConfigContainer: + """Return a pre-training config for Qwen3-30B-A3B on 16 H100 GPUs. + + Recommended parallelism: TP=1, PP=1, EP=16 with HybridEP. + """ + cfg = qwen3_30b_a3b_pretrain_8gpu_h100_bf16_config() + + # Precision and fused kernels + cfg.mixed_precision = bf16_mixed() + cfg.mixed_precision.grad_reduce_in_fp32 = False + cfg.ddp.grad_reduce_in_fp32 = False + cfg.optimizer.use_precision_aware_optimizer = True + cfg.model.bias_activation_fusion = True + cfg.model.apply_rope_fusion = True + cfg.model.moe_router_fusion = True + + # The 16-GPU topology fits without activation recomputation. + cfg.model.recompute_granularity = None + cfg.model.recompute_method = None + cfg.model.recompute_num_layers = None + + cfg.model.seq_length = 4096 + cfg.dataset.seq_length = 4096 + cfg.model.tensor_model_parallel_size = 1 + cfg.model.pipeline_model_parallel_size = 1 + cfg.model.context_parallel_size = 1 + cfg.model.virtual_pipeline_model_parallel_size = None + cfg.model.expert_model_parallel_size = 16 + cfg.model.expert_tensor_parallel_size = 1 + cfg.model.sequence_parallel = False + cfg.train.train_iters = 100 + cfg.train.global_batch_size = 1024 + cfg.train.micro_batch_size = 1 + cfg.dataset.random_seed = 1234 + cfg.rng.seed = 1234 + + cfg.optimizer.lr = 3e-4 + cfg.optimizer.min_lr = 3e-5 + cfg.optimizer.adam_beta1 = 0.9 + cfg.optimizer.adam_beta2 = 0.95 + cfg.optimizer.adam_eps = 1e-8 + cfg.optimizer.weight_decay = 0.1 + cfg.optimizer.clip_grad = 1.0 + cfg.optimizer.use_distributed_optimizer = True + cfg.optimizer.main_grads_dtype = torch.float32 + cfg.optimizer.main_params_dtype = torch.float32 + cfg.optimizer.exp_avg_dtype = torch.float32 + cfg.optimizer.exp_avg_sq_dtype = torch.float32 + cfg.scheduler.lr_warmup_iters = 40 + cfg.scheduler.lr_decay_iters = 100 + cfg.scheduler.lr_warmup_init = 0.0 + cfg.scheduler.start_weight_decay = 0.033 + cfg.scheduler.end_weight_decay = 0.033 + cfg.scheduler.weight_decay_incr_style = "constant" + cfg.scheduler.lr_decay_style = "cosine" + cfg.checkpoint.save_interval = 50 + cfg.checkpoint.load = None + + # HybridEP dispatcher + cfg.model.moe_flex_dispatcher_backend = "hybridep" + cfg.model.moe_token_dispatcher_type = "flex" + cfg.model.moe_a2a_overlap = False + cfg.model.moe_shared_expert_overlap = False + cfg.model.moe_hybridep_num_sms = 32 + cfg.model.moe_router_force_load_balancing = False + + # Transformer Engine scoped CUDA graphs + cfg.model.cuda_graph_impl = "transformer_engine" + cfg.model.cuda_graph_scope = ["moe_router", "moe_preprocess"] + cfg.rng.te_rng_tracker = True + cfg.model.use_te_rng_tracker = True + + cfg.comm_overlap = CommOverlapConfig(tp_comm_overlap=True) + + apply_flex_dispatcher_backend(cfg.model, cfg.model.moe_flex_dispatcher_backend) + # Keep the complete process environment visible on the recipe. + cfg.env_vars = { + **COMMON_RECIPE_ENV_VARS, + } + return cfg + + def qwen3_235b_a22b_pretrain_256gpu_h100_bf16_config() -> ConfigContainer: """Return a pre-training config for Qwen3-235B-A22B MoE. @@ -265,10 +355,13 @@ def qwen3_30b_a3b_sft_8gpu_h100_bf16_config() -> ConfigContainer: cfg = _sft_common() # Model config - cfg.model = AutoBridge.from_hf_pretrained("Qwen/Qwen3-30B-A3B").to_megatron_provider(load_weights=False) + cfg.model = AutoBridge.from_hf_pretrained( + _QWEN3_30B_A3B_MODEL_ID, revision=_QWEN3_30B_A3B_MODEL_REVISION + ).to_megatron_provider(load_weights=False) # Tokenizer - cfg.tokenizer.tokenizer_model = "Qwen/Qwen3-30B-A3B" + cfg.tokenizer.tokenizer_model = _QWEN3_30B_A3B_MODEL_ID + cfg.tokenizer.hf_tokenizer_kwargs = {"revision": _QWEN3_30B_A3B_MODEL_REVISION} # Parallelism settings (MoE-specific: includes expert parallelism) cfg.model.tensor_model_parallel_size = 4 @@ -284,11 +377,12 @@ def qwen3_30b_a3b_sft_8gpu_h100_bf16_config() -> ConfigContainer: # Sequence length (2048 for packed sequences) cfg.model.seq_length = 2048 - # Global batch size is 32 for MoE packed sequences + # qwen3_30b_a3b_convergence_v1 batch and RNG contract cfg.train.global_batch_size = 32 - # Set pad_seq_to_mult for context parallelism - if cfg.model.context_parallel_size > 1: - cfg.dataset.offline_packing_specs.pad_seq_to_mult = cfg.model.context_parallel_size * 2 + cfg.train.micro_batch_size = 1 + cfg.dataset.seed = 1234 + cfg.rng.seed = 5678 + cfg.dataset.offline_packing_specs.pad_seq_to_mult = 4 # MoE Token Dispatcher settings, may be overridden by apply_flex_dispatcher_backend at the end cfg.model.moe_token_dispatcher_type = "alltoall" @@ -314,7 +408,18 @@ def qwen3_30b_a3b_sft_8gpu_h100_bf16_config() -> ConfigContainer: # Optimizer and scheduler overrides for MoE cfg.scheduler.lr_warmup_iters = 10 cfg.scheduler.lr_decay_iters = 100 # Same as train_iters + cfg.scheduler.lr_warmup_init = 0.0 + cfg.scheduler.start_weight_decay = 0.033 + cfg.scheduler.end_weight_decay = 0.033 + cfg.scheduler.weight_decay_incr_style = "constant" + cfg.scheduler.lr_decay_style = "cosine" + cfg.optimizer.lr = 5e-6 + cfg.optimizer.min_lr = 0.0 + cfg.optimizer.adam_beta1 = 0.9 cfg.optimizer.adam_beta2 = 0.95 + cfg.optimizer.adam_eps = 1e-8 + cfg.optimizer.weight_decay = 0.1 + cfg.optimizer.clip_grad = 1.0 # TE (Transformer Engine) cfg.model.transformer_impl = "transformer_engine" @@ -353,6 +458,7 @@ def qwen3_30b_a3b_sft_8gpu_h100_bf16_config() -> ConfigContainer: cfg.optimizer.main_params_dtype = torch.float32 cfg.optimizer.exp_avg_dtype = torch.float32 cfg.optimizer.exp_avg_sq_dtype = torch.float32 + cfg.optimizer.use_distributed_optimizer = True # Communication overlap # cfg.comm_overlap = CommOverlapConfig(tp_comm_overlap=False) # Uncomment to enable @@ -363,6 +469,7 @@ def qwen3_30b_a3b_sft_8gpu_h100_bf16_config() -> ConfigContainer: # Checkpoint config cfg.checkpoint.save_interval = 100 + cfg.checkpoint.load = None # cfg.checkpoint.save and cfg.checkpoint.load are set in _sft_common. To override: # cfg.checkpoint.save = "path/to/save" # cfg.checkpoint.load = "path/to/load" @@ -389,6 +496,50 @@ def qwen3_30b_a3b_sft_8gpu_h100_bf16_config() -> ConfigContainer: return cfg +def qwen3_30b_a3b_sft_16gpu_h100_bf16_config() -> ConfigContainer: + """Return a full SFT config for Qwen3-30B-A3B MoE on 16 H100 GPUs. + + Recommended parallelism: TP=1, PP=1, EP=16 (2 nodes, 16 GPUs). + This topology keeps the global batch size at 32 while using DP=16 and + two gradient-accumulation steps for the default micro batch size of 1. + """ + cfg = qwen3_30b_a3b_sft_8gpu_h100_bf16_config() + + # Parallelism settings + cfg.model.tensor_model_parallel_size = 1 + cfg.model.pipeline_model_parallel_size = 1 + cfg.model.context_parallel_size = 1 + cfg.model.expert_model_parallel_size = 16 + cfg.model.expert_tensor_parallel_size = 1 + cfg.model.sequence_parallel = False + + # Keep GBS fixed while increasing data parallelism. Offline packing + # currently requires MBS=1, which gives two accumulation steps at DP=16. + cfg.train.global_batch_size = 32 + cfg.train.micro_batch_size = 1 + cfg.dataset.offline_packing_specs.pad_seq_to_mult = 1 + + # The plain all-to-all dispatcher and fused kernels are faster for this + # short-sequence SFT workload than DeepEP, EP overlap, or CUDA graphs. + cfg.model.moe_token_dispatcher_type = "alltoall" + cfg.model.moe_flex_dispatcher_backend = None + cfg.model.moe_hybridep_num_sms = None + cfg.model.moe_flex_dispatcher_num_sms = None + cfg.model.moe_a2a_overlap = False + cfg.model.moe_shared_expert_overlap = False + cfg.model.bias_activation_fusion = True + cfg.model.apply_rope_fusion = True + cfg.model.moe_router_fusion = True + cfg.model.cuda_graph_impl = "none" + cfg.comm_overlap = None + + # Keep the complete process environment visible on the recipe. + cfg.env_vars = { + **COMMON_RECIPE_ENV_VARS, + } + return cfg + + def qwen3_235b_a22b_sft_64gpu_h100_bf16_config() -> ConfigContainer: """Return a full SFT config for Qwen3-235B-A22B MoE. @@ -537,16 +688,19 @@ def qwen3_30b_a3b_peft_4gpu_h100_bf16_config(peft_scheme: str | PEFT = "lora") - Args: peft_scheme: PEFT scheme - 'lora', 'dora', or a PEFT instance. Default: 'lora' - Recommended parallelism: TP=4, PP=1, EP=4 (1 node, 8 GPUs with SP=True) + Recommended parallelism: TP=4, PP=1, EP=4 (1 node, 4 GPUs with SP=True) LoRA/DoRA uses dim=8, alpha=16, target_modules=['linear_qkv', 'linear_proj'] """ cfg = _peft_common() # Model config - cfg.model = AutoBridge.from_hf_pretrained("Qwen/Qwen3-30B-A3B").to_megatron_provider(load_weights=False) + cfg.model = AutoBridge.from_hf_pretrained( + _QWEN3_30B_A3B_MODEL_ID, revision=_QWEN3_30B_A3B_MODEL_REVISION + ).to_megatron_provider(load_weights=False) # Tokenizer - cfg.tokenizer.tokenizer_model = "Qwen/Qwen3-30B-A3B" + cfg.tokenizer.tokenizer_model = _QWEN3_30B_A3B_MODEL_ID + cfg.tokenizer.hf_tokenizer_kwargs = {"revision": _QWEN3_30B_A3B_MODEL_REVISION} # Parallelism settings (MoE-specific: includes expert parallelism) # PEFT uses PP=1 (less parallelism needed since only adapters are trained) @@ -569,14 +723,16 @@ def qwen3_30b_a3b_peft_4gpu_h100_bf16_config(peft_scheme: str | PEFT = "lora") - if isinstance(peft_scheme, str) and peft_scheme.lower() in ["lora", "dora"]: peft_cfg.dim = 8 peft_cfg.alpha = 16 + peft_cfg.dropout = 0.0 peft_cfg.target_modules = ["linear_qkv", "linear_proj"] cfg.peft = peft_cfg - # Global batch size is 32 for MoE packed sequences + # qwen3_30b_a3b_convergence_v1 batch and RNG contract cfg.train.global_batch_size = 32 - # Set pad_seq_to_mult for context parallelism - if cfg.model.context_parallel_size > 1: - cfg.dataset.offline_packing_specs.pad_seq_to_mult = cfg.model.context_parallel_size * 2 + cfg.train.micro_batch_size = 1 + cfg.dataset.seed = 1234 + cfg.rng.seed = 5678 + cfg.dataset.offline_packing_specs.pad_seq_to_mult = 4 # MoE Token Dispatcher settings, moe_token_dispatcher_type may be overridden by apply_flex_dispatcher_backend at the end cfg.model.moe_token_dispatcher_type = "alltoall" @@ -602,7 +758,18 @@ def qwen3_30b_a3b_peft_4gpu_h100_bf16_config(peft_scheme: str | PEFT = "lora") - # Optimizer and scheduler overrides for MoE cfg.scheduler.lr_warmup_iters = 10 cfg.scheduler.lr_decay_iters = 100 # Same as train_iters + cfg.scheduler.lr_warmup_init = 0.0 + cfg.scheduler.start_weight_decay = 0.033 + cfg.scheduler.end_weight_decay = 0.033 + cfg.scheduler.weight_decay_incr_style = "constant" + cfg.scheduler.lr_decay_style = "cosine" + cfg.optimizer.lr = 1e-4 + cfg.optimizer.min_lr = 0.0 + cfg.optimizer.adam_beta1 = 0.9 cfg.optimizer.adam_beta2 = 0.95 + cfg.optimizer.adam_eps = 1e-8 + cfg.optimizer.weight_decay = 0.1 + cfg.optimizer.clip_grad = 1.0 # TE (Transformer Engine) cfg.model.transformer_impl = "transformer_engine" @@ -641,6 +808,7 @@ def qwen3_30b_a3b_peft_4gpu_h100_bf16_config(peft_scheme: str | PEFT = "lora") - cfg.optimizer.main_params_dtype = torch.float32 cfg.optimizer.exp_avg_dtype = torch.float32 cfg.optimizer.exp_avg_sq_dtype = torch.float32 + cfg.optimizer.use_distributed_optimizer = True # Communication overlap # cfg.comm_overlap = CommOverlapConfig(tp_comm_overlap=False) # Uncomment to enable @@ -651,6 +819,7 @@ def qwen3_30b_a3b_peft_4gpu_h100_bf16_config(peft_scheme: str | PEFT = "lora") - # Checkpoint config cfg.checkpoint.save_interval = 100 + cfg.checkpoint.load = None # cfg.checkpoint.save and cfg.checkpoint.load are set in _peft_common. To override: # cfg.checkpoint.save = "path/to/save" # cfg.checkpoint.load = "path/to/load" @@ -833,6 +1002,8 @@ def qwen3_235b_a22b_peft_16gpu_h100_bf16_config(peft_scheme: str | PEFT = "lora" "qwen3_235b_a22b_pretrain_256gpu_h100_bf16_config", # pragma: allowlist secret "qwen3_235b_a22b_sft_64gpu_h100_bf16_config", "qwen3_30b_a3b_peft_4gpu_h100_bf16_config", + "qwen3_30b_a3b_pretrain_16gpu_h100_bf16_config", "qwen3_30b_a3b_pretrain_8gpu_h100_bf16_config", + "qwen3_30b_a3b_sft_16gpu_h100_bf16_config", "qwen3_30b_a3b_sft_8gpu_h100_bf16_config", ] diff --git a/src/megatron/bridge/recipes/qwen/qwen3.py b/src/megatron/bridge/recipes/qwen/qwen3.py index c8899937f7..92009bba11 100644 --- a/src/megatron/bridge/recipes/qwen/qwen3.py +++ b/src/megatron/bridge/recipes/qwen/qwen3.py @@ -38,11 +38,14 @@ qwen3_8b_peft_1gpu_h100_bf16_config as qwen3_8b_peft_config, ) from megatron.bridge.recipes.qwen.h100.qwen3 import ( - qwen3_8b_pretrain_4gpu_h100_bf16_config as qwen3_8b_pretrain_config, + qwen3_8b_pretrain_16gpu_h100_bf16_config as qwen3_8b_pretrain_config, ) from megatron.bridge.recipes.qwen.h100.qwen3 import ( qwen3_8b_sft_4gpu_h100_bf16_config as qwen3_8b_sft_config, ) +from megatron.bridge.recipes.qwen.h100.qwen3 import ( + qwen3_8b_sft_8gpu_h100_bf16_32k_config as qwen3_8b_sft_32k_config, +) from megatron.bridge.recipes.qwen.h100.qwen3 import ( qwen3_14b_peft_1gpu_h100_bf16_config as qwen3_14b_peft_config, ) @@ -98,5 +101,6 @@ "qwen3_600m_sft_yarn_128k_config", "qwen3_8b_peft_config", "qwen3_8b_pretrain_config", + "qwen3_8b_sft_32k_config", "qwen3_8b_sft_config", ] diff --git a/src/megatron/bridge/recipes/qwen/qwen3_moe.py b/src/megatron/bridge/recipes/qwen/qwen3_moe.py index da33ee04bb..b31472e3d0 100644 --- a/src/megatron/bridge/recipes/qwen/qwen3_moe.py +++ b/src/megatron/bridge/recipes/qwen/qwen3_moe.py @@ -12,7 +12,7 @@ # See the License for the specific language governing permissions and # limitations under the License. # ruff: noqa: F401 -"""Compatibility aliases for legacy recipe names.""" +"""Compatibility aliases for generic Qwen3 MoE recipe names.""" from __future__ import annotations @@ -20,10 +20,10 @@ qwen3_30b_a3b_peft_4gpu_h100_bf16_config as qwen3_30b_a3b_peft_config, ) from megatron.bridge.recipes.qwen.h100.qwen3_moe import ( - qwen3_30b_a3b_pretrain_8gpu_h100_bf16_config as qwen3_30b_a3b_pretrain_config, + qwen3_30b_a3b_pretrain_16gpu_h100_bf16_config as qwen3_30b_a3b_pretrain_config, ) from megatron.bridge.recipes.qwen.h100.qwen3_moe import ( - qwen3_30b_a3b_sft_8gpu_h100_bf16_config as qwen3_30b_a3b_sft_config, + qwen3_30b_a3b_sft_16gpu_h100_bf16_config as qwen3_30b_a3b_sft_config, ) from megatron.bridge.recipes.qwen.h100.qwen3_moe import ( qwen3_235b_a22b_peft_16gpu_h100_bf16_config as qwen3_235b_a22b_peft_config, diff --git a/src/megatron/bridge/training/tokenizers/config.py b/src/megatron/bridge/training/tokenizers/config.py index c22c1f8363..5cbeafe435 100644 --- a/src/megatron/bridge/training/tokenizers/config.py +++ b/src/megatron/bridge/training/tokenizers/config.py @@ -40,6 +40,7 @@ class TokenizerConfig(MTrainTokenizerConfig): - use_fast (bool): Whether to use fast tokenizer implementation - trust_remote_code (bool): Whether to trust remote code when loading tokenizer - include_special_tokens (bool): Whether to include special tokens when converting text to ids + - revision (str): Hugging Face Hub revision used to resolve an immutable tokenizer snapshot Example: hf_tokenizer_kwargs = { diff --git a/src/megatron/bridge/training/tokenizers/tokenizer.py b/src/megatron/bridge/training/tokenizers/tokenizer.py index cd4a1d06a9..c4c0cc80e2 100644 --- a/src/megatron/bridge/training/tokenizers/tokenizer.py +++ b/src/megatron/bridge/training/tokenizers/tokenizer.py @@ -2,12 +2,91 @@ """Megatron tokenizers.""" +from copy import copy +from pathlib import Path + from megatron.core.tokenizers import MegatronTokenizer from megatron.core.tokenizers.utils.build_tokenizer import build_tokenizer as build_mcore_tokenizer from megatron.bridge.training.tokenizers.config import TokenizerConfig +_HF_TOKENIZER_SNAPSHOT_ALLOW_PATTERNS = ( + "config.json", + "tokenizer*", + "special_tokens_map.json", + "added_tokens.json", + "vocab*", + "merges*", + "*.model", + "*.spm", + "*.tiktoken", + "*.py", + "*.jinja", + "chat_template*", + "chat_templates/*", +) +_HF_MODEL_WEIGHT_IGNORE_PATTERNS = ( + "*.safetensors", + "*.bin", + "*.pt", + "*.pth", + "*.ckpt", + "*.h5", + "*.msgpack", + "*.gguf", + "*.onnx", + "*.npz", +) + + +def _is_local_tokenizer_path(tokenizer_model: object) -> bool: + """Return whether a tokenizer reference is explicitly a local path.""" + if isinstance(tokenizer_model, Path): + return True + if not isinstance(tokenizer_model, str): + return False + + path = Path(tokenizer_model).expanduser() + return path.exists() or path.is_absolute() or tokenizer_model.startswith(("./", "../", "~")) + + +def _resolve_hf_tokenizer_revision(config: TokenizerConfig) -> TokenizerConfig: + """Resolve an immutable Hugging Face tokenizer revision without mutating persisted config.""" + if config.tokenizer_type != "HuggingFaceTokenizer": + return config + + hf_tokenizer_kwargs = config.hf_tokenizer_kwargs or {} + revision = hf_tokenizer_kwargs.get("revision") + if revision is None: + return config + if not isinstance(revision, str) or not revision.strip(): + raise ValueError("hf_tokenizer_kwargs.revision must be a non-empty string") + + tokenizer_model = config.tokenizer_model + if not isinstance(tokenizer_model, str) or _is_local_tokenizer_path(tokenizer_model): + return config + + from huggingface_hub import snapshot_download + + snapshot_path = snapshot_download( + repo_id=tokenizer_model, + revision=revision, + allow_patterns=list(_HF_TOKENIZER_SNAPSHOT_ALLOW_PATTERNS), + ignore_patterns=list(_HF_MODEL_WEIGHT_IGNORE_PATTERNS), + ) + + resolved_config = copy(config) + resolved_config.tokenizer_model = snapshot_path + if "use_fast" in hf_tokenizer_kwargs: + resolved_config.tokenizer_hf_no_use_fast = not hf_tokenizer_kwargs["use_fast"] + if "include_special_tokens" in hf_tokenizer_kwargs: + resolved_config.tokenizer_hf_no_include_special_tokens = not hf_tokenizer_kwargs["include_special_tokens"] + if "trust_remote_code" in hf_tokenizer_kwargs: + resolved_config.trust_remote_code = hf_tokenizer_kwargs["trust_remote_code"] + return resolved_config + + def build_tokenizer(config: TokenizerConfig, **kwargs) -> MegatronTokenizer: """Initialize tokenizer from megatron.core.tokenizers based on the provided configuration. @@ -26,4 +105,4 @@ def build_tokenizer(config: TokenizerConfig, **kwargs) -> MegatronTokenizer: "Please, use `megatron.core.tokenizers.utils.build_tokenizer` instead." ) - return build_mcore_tokenizer(config, **kwargs) + return build_mcore_tokenizer(_resolve_hf_tokenizer_revision(config), **kwargs) diff --git a/tests/functional_tests/test_groups/data/test_loaders.py b/tests/functional_tests/test_groups/data/test_loaders.py index 492799e8d8..6249cf329a 100644 --- a/tests/functional_tests/test_groups/data/test_loaders.py +++ b/tests/functional_tests/test_groups/data/test_loaders.py @@ -304,10 +304,28 @@ def test_build_train_valid_test_data_loaders_allows_none_train_dataset( @mock.patch("torch.distributed.get_rank") @mock.patch("megatron.bridge.data.loaders.build_pretraining_data_loader", return_value=object()) @mock.patch("megatron.bridge.data.loaders.build_train_valid_test_datasets") + @pytest.mark.parametrize("dataset_kind", ["gpt", "sft"]) def test_build_train_valid_test_data_loaders_uses_eval_dp_group( - self, mock_build_datasets, mock_build_loader, mock_get_rank, mock_get_world_size, _mock_broadcast + self, + mock_build_datasets, + mock_build_loader, + mock_get_rank, + mock_get_world_size, + _mock_broadcast, + dataset_kind, ): cfg = create_simple_test_config() + if dataset_kind == "sft": + cfg.dataset = GPTSFTDatasetConfig( + dataset_root="/tmp/dataset", + seq_length=512, + seed=4321, + num_workers=0, + persistent_workers=False, + ) + expected_seed = cfg.dataset.seed + else: + expected_seed = cfg.dataset.random_seed train_ds = mock.MagicMock() train_ds.__len__.return_value = cfg.train.global_batch_size valid_ds = mock.MagicMock() @@ -334,6 +352,9 @@ def test_build_train_valid_test_data_loaders_uses_eval_dp_group( assert valid_call.kwargs["data_parallel_size"] == 1 assert test_call.kwargs["data_parallel_rank"] == 0 assert test_call.kwargs["data_parallel_size"] == 1 + assert train_call.kwargs["seed"] == expected_seed + assert valid_call.kwargs["seed"] == expected_seed + assert test_call.kwargs["seed"] == expected_seed @mock.patch("torch.distributed.broadcast") @mock.patch("torch.distributed.get_world_size", return_value=1) diff --git a/tests/functional_tests/test_groups/data/test_samplers.py b/tests/functional_tests/test_groups/data/test_samplers.py index 4c1355f347..aa13e148cb 100644 --- a/tests/functional_tests/test_groups/data/test_samplers.py +++ b/tests/functional_tests/test_groups/data/test_samplers.py @@ -845,6 +845,26 @@ def mock_dataloader(): class TestBatchDataloaderIntegration: """Integration tests for batch dataloader type.""" + def test_build_batch_dataloader_explicit_seed_is_independent_of_model_rng(self): + """Pipeline-stage model seeds must not change the fine-tuning sample order.""" + import torch + + def first_global_batch(model_seed: int) -> list[int]: + torch.manual_seed(model_seed) + dataloader = build_pretraining_data_loader( + dataset=list(range(64)), + consumed_samples=0, + dataloader_type="batch", + micro_batch_size=1, + num_workers=0, + data_sharding=False, + global_batch_size=8, + seed=1234, + ) + return next(iter(dataloader.batch_sampler)) + + assert first_global_batch(5678) == first_global_batch(5778) + def test_build_batch_dataloader_basic(self): """Test building a dataloader with dataloader_type='batch'.""" from unittest import mock as _mock diff --git a/tests/unit_tests/conversion/launcher/test_arguments.py b/tests/unit_tests/conversion/launcher/test_arguments.py index d8befff614..10acd45731 100644 --- a/tests/unit_tests/conversion/launcher/test_arguments.py +++ b/tests/unit_tests/conversion/launcher/test_arguments.py @@ -146,6 +146,8 @@ def test_roundtrip_alias_and_worker_args(): "8", "--hf-model-id", "hf/model", + "--hf-revision", + "0123456789abcdef", "--tp", "2", "--pp", @@ -176,12 +178,33 @@ def test_roundtrip_alias_and_worker_args(): "4", "--etp", "2", + "--hf-revision", + "0123456789abcdef", "--trust-remote-code", "--distributed-timeout-minutes", "30", ] +def test_import_worker_args_forward_hf_revision(): + module = _load_arguments_module() + args = module.build_parser(include_execution=True).parse_args( + [ + "import", + "--hf-model", + "hf/model", + "--hf-revision", + "0123456789abcdef", + "--megatron-path", + "/checkpoint", + ] + ) + + worker_args = module.conversion_worker_args(args) + + assert worker_args[worker_args.index("--hf-revision") + 1] == "0123456789abcdef" + + def test_roundtrip_worker_parser_accepts_serialized_args(): module = _load_arguments_module() diff --git a/tests/unit_tests/conversion/launcher/test_run_conversion.py b/tests/unit_tests/conversion/launcher/test_run_conversion.py index e97fbe21a0..a8508a0780 100644 --- a/tests/unit_tests/conversion/launcher/test_run_conversion.py +++ b/tests/unit_tests/conversion/launcher/test_run_conversion.py @@ -80,6 +80,33 @@ def test_cpu_import_dispatches_to_cpu_backend(): ] +def test_cpu_import_resolves_hf_revision_before_dispatch(monkeypatch): + module, cpu_backend, _ = _load_run_conversion_module() + calls = [] + cpu_backend.import_checkpoint = lambda **kwargs: calls.append(kwargs) + monkeypatch.setattr( + module, + "resolve_hf_model_revision", + lambda model, revision: f"/snapshots/{model.replace('/', '--')}/{revision}", + ) + + module.main( + [ + "import", + "--device", + "cpu", + "--hf-model", + "hf/model", + "--hf-revision", + "0123456789abcdef", + "--megatron-path", + "/checkpoint", + ] + ) + + assert calls[0]["hf_model"] == "/snapshots/hf--model/0123456789abcdef" # pragma: allowlist secret + + def test_gpu_export_enables_distributed_save_by_default(): module, _, gpu_backend = _load_run_conversion_module() calls = [] diff --git a/tests/unit_tests/conversion/launcher/test_setup_conversion.py b/tests/unit_tests/conversion/launcher/test_setup_conversion.py index 3d730c10c8..e9e92636c5 100644 --- a/tests/unit_tests/conversion/launcher/test_setup_conversion.py +++ b/tests/unit_tests/conversion/launcher/test_setup_conversion.py @@ -358,12 +358,15 @@ def __init__(self, **kwargs): "partition", "--container-image", "image.sqsh", + "--experiment-name", + "mb4909-nano4b-conversion", ) module._validate_args(args) executor = module._build_executor(args, ["HF_TOKEN"], ["/host:/container"]) assert executor.kwargs["ntasks_per_node"] == 1 + assert executor.kwargs["job_name_prefix"] == "mb4909-nano4b-conversion" assert "cpus_per_task" not in executor.kwargs assert "gpus_per_node" not in executor.kwargs assert executor.kwargs["container_env"] == ["HF_TOKEN", "PYTHONPATH"] diff --git a/tests/unit_tests/conversion/launcher/test_utils.py b/tests/unit_tests/conversion/launcher/test_utils.py index 96e6d7da8f..6ee3797de1 100644 --- a/tests/unit_tests/conversion/launcher/test_utils.py +++ b/tests/unit_tests/conversion/launcher/test_utils.py @@ -26,6 +26,25 @@ conversion_utils = importlib.util.module_from_spec(spec) spec.loader.exec_module(conversion_utils) prepare_output_directory = conversion_utils.prepare_output_directory +resolve_hf_model_revision = conversion_utils.resolve_hf_model_revision + + +def test_resolve_hf_model_revision_downloads_exact_snapshot(monkeypatch): + calls = [] + monkeypatch.setattr( + "huggingface_hub.snapshot_download", + lambda **kwargs: calls.append(kwargs) or "/cache/models--hf--model/snapshots/0123456789abcdef", + ) + + resolved = resolve_hf_model_revision("hf/model", "0123456789abcdef") + + assert resolved == "/cache/models--hf--model/snapshots/0123456789abcdef" + assert calls == [{"repo_id": "hf/model", "revision": "0123456789abcdef"}] + + +def test_resolve_hf_model_revision_rejects_local_path(tmp_path): + with pytest.raises(ValueError, match="only to Hugging Face Hub model IDs"): + resolve_hf_model_revision(str(tmp_path), "0123456789abcdef") @pytest.mark.parametrize("relationship", ["equal", "destination-parent", "destination-child"]) diff --git a/tests/unit_tests/data/packing/test_parquet.py b/tests/unit_tests/data/packing/test_parquet.py index 590a54525e..ec2111e882 100644 --- a/tests/unit_tests/data/packing/test_parquet.py +++ b/tests/unit_tests/data/packing/test_parquet.py @@ -1,4 +1,4 @@ -# Copyright (c) 2025, NVIDIA CORPORATION. All rights reserved. +# Copyright (c) 2025-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. diff --git a/tests/unit_tests/data/test_conversation_processing.py b/tests/unit_tests/data/test_conversation_processing.py index 06019b4027..6b37c202dd 100644 --- a/tests/unit_tests/data/test_conversation_processing.py +++ b/tests/unit_tests/data/test_conversation_processing.py @@ -243,6 +243,51 @@ def __init__(self): self.tokenizer = _ChatMLBoundaryTokenizer() +class _MoonlightBoundaryTokenizer(_Tokenizer): + chat_template = ( + "{%- for message in messages -%}" + "{%- if message['role'] == 'system' -%}<|im_system|>{%- endif -%}" + "{%- if message['role'] == 'user' -%}<|im_user|>{%- endif -%}" + "{%- if message['role'] == 'assistant' -%}<|im_assistant|>{%- endif -%}" + "{{ message['role'] }}<|im_middle|>{{ message['content'] }}<|im_end|>" + "{%- endfor -%}" + ) + + _role_markers = { + "system": [405, 415, 401], + "user": [400, 414, 401], + "assistant": [402, 412, 401], + } + + def __call__(self, text, add_special_tokens=False): + mapping = { + "<|im_system|>system<|im_middle|>": self._role_markers["system"], + "<|im_user|>user<|im_middle|>": self._role_markers["user"], + "<|im_assistant|>assistant<|im_middle|>": self._role_markers["assistant"], + "<|im_end|>": [403], + "question": [16], + "answer": [3, 4], + } + return {"input_ids": mapping.get(text, [42])} + + def apply_chat_template(self, conversation, tokenize=True, add_generation_prompt=False, return_dict=False): + assert tokenize is True + assert add_generation_prompt is False + assert return_dict is True + input_ids = [] + for turn in conversation: + input_ids.extend(self._role_markers[turn["role"]]) + input_ids.extend(self(turn["content"])["input_ids"]) + input_ids.append(403) + return {"input_ids": input_ids} + + +class _MoonlightBoundaryProcessor(_Processor): + def __init__(self): + super().__init__() + self.tokenizer = _MoonlightBoundaryTokenizer() + + class _ProcessorTemplateBoundaryProcessor(_ChatMLBoundaryProcessor): chat_template = "<|turn>model\n{{ content }}" @@ -493,6 +538,44 @@ def test_infer_assistant_mask_boundary_config_from_chatml_template(): assert all(token_variants == [[103]] for token_variants in boundary_config.role_end_token_variants.values()) +def test_infer_assistant_mask_boundary_config_from_moonlight_template(): + processor = _MoonlightBoundaryProcessor() + assert "<|im_assistant|>assistant<|im_middle|>" not in processor.tokenizer.chat_template + boundary_config = infer_assistant_mask_boundary_config(processor) + + assert boundary_config is not None + assert boundary_config.role_start_tokens == { + "assistant": [402, 412, 401], + "system": [405, 415, 401], + "user": [400, 414, 401], + } + assert all(token_ids == [403] for token_ids in boundary_config.role_end_tokens.values()) + + tokenized = tokenize_chat_example( + { + "messages": [ + {"role": "user", "content": "question"}, + {"role": "assistant", "content": "answer"}, + ] + }, + processor, + ) + assert tokenized.input_ids.tolist() == [400, 414, 401, 16, 403, 402, 412, 401, 3, 4, 403] + assert tokenized.assistant_mask.tolist() == [ + False, + False, + False, + False, + False, + False, + False, + False, + True, + True, + True, + ] + + def test_infer_assistant_mask_boundary_config_handles_jinja_separated_chatml_newline(): boundary_config = infer_assistant_mask_boundary_config(_JinjaSeparatedChatMLBoundaryProcessor()) diff --git a/tests/unit_tests/examples/test_hf_to_megatron_generate_text.py b/tests/unit_tests/examples/test_hf_to_megatron_generate_text.py index 26dd29defd..e45d349003 100644 --- a/tests/unit_tests/examples/test_hf_to_megatron_generate_text.py +++ b/tests/unit_tests/examples/test_hf_to_megatron_generate_text.py @@ -14,7 +14,7 @@ import runpy from pathlib import Path -from unittest.mock import MagicMock +from unittest.mock import MagicMock, patch import pytest import torch @@ -22,8 +22,14 @@ _SCRIPT = Path(__file__).parents[3] / "examples" / "conversion" / "hf_to_megatron_generate_text.py" _SCRIPT_GLOBALS = runpy.run_path(_SCRIPT) +_build_parser = _SCRIPT_GLOBALS["build_parser"] _decode_completion = _SCRIPT_GLOBALS["_decode_completion"] +_hf_revision_kwargs = _SCRIPT_GLOBALS["_hf_revision_kwargs"] +_maybe_gather_tensor_parallel_logits = _SCRIPT_GLOBALS["_maybe_gather_tensor_parallel_logits"] +_run_megatron_forward = _SCRIPT_GLOBALS["_run_megatron_forward"] +_text_forward_step = _SCRIPT_GLOBALS["text_forward_step"] _tokenize_prompt = _SCRIPT_GLOBALS["_tokenize_prompt"] +_InferenceMode = _SCRIPT_GLOBALS["InferenceMode"] @pytest.mark.unit @@ -66,3 +72,75 @@ def test_decode_completion_excludes_prompt_and_special_tokens() -> None: assert text == "The sky appears blue." tokenizer.decode.assert_called_once_with([20, 21], skip_special_tokens=True) + + +@pytest.mark.unit +def test_hf_revision_is_parsed_and_forwarded() -> None: + revision = "0123456789abcdef0123456789abcdef01234567" # pragma: allowlist secret + + args = _build_parser().parse_args(["--hf_model_path", "org/model", "--hf-revision", revision]) + + assert args.hf_revision == revision + assert _hf_revision_kwargs(args.hf_revision) == {"revision": revision} + assert _hf_revision_kwargs(None) == {} + + +@pytest.mark.unit +def test_text_forward_step_passes_static_context_and_gather_request() -> None: + inference_context = object() + batch = { + "tokens": torch.tensor([[1, 2, 3]]), + "position_ids": torch.arange(3).unsqueeze(0), + "inference_context": inference_context, + } + model = MagicMock(return_value=torch.randn(1, 3, 16)) + + _text_forward_step(iter([batch]), model) + + assert model.call_args.kwargs["inference_context"] is inference_context + assert model.call_args.kwargs["runtime_gather_output"] is True + + +@pytest.mark.unit +def test_generation_forward_activates_inference_mode() -> None: + forward = MagicMock(return_value="output") + context = MagicMock() + + with patch.object(_InferenceMode, "active", return_value=context) as active: + result = _run_megatron_forward(forward, forward_only=True) + + assert result == "output" + active.assert_called_once_with() + context.__enter__.assert_called_once_with() + context.__exit__.assert_called_once() + forward.assert_called_once_with(forward_only=True) + + +@pytest.mark.unit +def test_generation_skips_tp_gather_for_complete_vocabulary() -> None: + full_logits = torch.randn(1, 1, 128) + + with patch.object(torch.distributed, "all_gather") as all_gather: + result = _maybe_gather_tensor_parallel_logits(full_logits, 128, 2, object()) + + assert result is full_logits + all_gather.assert_not_called() + + +@pytest.mark.unit +def test_generation_gathers_tp_vocabulary_shards() -> None: + local_logits = torch.arange(64, dtype=torch.float32).reshape(1, 1, 64) + tp_group = object() + + def mock_all_gather(outputs, tensor, group): + assert group is tp_group + outputs[0].copy_(tensor) + outputs[1].copy_(tensor + 64) + + with patch.object(torch.distributed, "all_gather", side_effect=mock_all_gather) as all_gather: + result = _maybe_gather_tensor_parallel_logits(local_logits, 128, 2, tp_group) + + assert result.shape == (1, 1, 128) + assert torch.equal(result[..., :64], local_logits) + assert torch.equal(result[..., 64:], local_logits + 64) + all_gather.assert_called_once() diff --git a/tests/unit_tests/models/deepseek/test_deepseek_bridges.py b/tests/unit_tests/models/deepseek/test_deepseek_bridges.py index 04dc364caa..8f2785ba28 100644 --- a/tests/unit_tests/models/deepseek/test_deepseek_bridges.py +++ b/tests/unit_tests/models/deepseek/test_deepseek_bridges.py @@ -296,6 +296,19 @@ def test_provider_bridge_maps_config(self, mock_pretrained_v3): assert provider.bf16 is True assert provider.params_dtype == torch.bfloat16 + def test_provider_bridge_preserves_model_specific_context_and_aux_loss(self, mock_pretrained_v3): + mock_pretrained_v3.config.max_position_embeddings = 8192 + mock_pretrained_v3.config.aux_loss_alpha = 0.001 + bridge = DeepSeekV3Bridge() + + provider = bridge.provider_bridge(mock_pretrained_v3) + exported = bridge.megatron_to_hf_config(provider) + + assert provider.seq_length == 8192 + assert provider.moe_aux_loss_coeff == 0.001 + assert exported["max_position_embeddings"] == 8192 + assert exported["aux_loss_alpha"] == 0.001 + def test_hf_config_to_provider_kwargs_preserves_none_q_lora_rank(self, mock_pretrained_v3): mock_pretrained_v3.config.q_lora_rank = None bridge = DeepSeekV3Bridge() @@ -323,7 +336,7 @@ def test_megatron_to_hf_config_preserves_none_q_lora_rank(self, mock_pretrained_ assert "q_lora_rank" in hf_config assert hf_config["q_lora_rank"] is None - def test_export_injects_inv_freq_for_layer(self, mock_pretrained_v3): + def test_export_does_not_synthesize_inv_freq_from_another_layer(self, mock_pretrained_v3): bridge = DeepSeekV3Bridge() bridge.hf_config = mock_pretrained_v3.config mock_pretrained_v3.state = {"model.layers.1.self_attn.rotary_emb.inv_freq": torch.randn(1)} @@ -336,16 +349,26 @@ def test_export_injects_inv_freq_for_layer(self, mock_pretrained_v3): result = bridge.maybe_modify_converted_hf_weight(task, dict(converted), mock_pretrained_v3.state) inv_key = "model.layers.0.self_attn.rotary_emb.inv_freq" - expected = 1.0 / ( - mock_pretrained_v3.config.rope_theta - ** ( - torch.arange(0, mock_pretrained_v3.config.qk_rope_head_dim, 2, dtype=torch.float32) - / mock_pretrained_v3.config.qk_rope_head_dim - ) + assert inv_key not in result + + def test_export_preserves_source_inv_freq_for_layer(self, mock_pretrained_v3): + bridge = DeepSeekV3Bridge() + bridge.hf_config = mock_pretrained_v3.config + inv_key = "model.layers.0.self_attn.rotary_emb.inv_freq" + source_inv_freq = torch.linspace(1.0, 0.0, 64, dtype=torch.bfloat16) + mock_pretrained_v3.state = {inv_key: source_inv_freq} + task = WeightConversionTask( + param_name="decoder.layers.0.input_layernorm.weight", + global_param_name="decoder.layers.0.input_layernorm.weight", + mapping=Mock(), ) + converted = {"model.layers.0.input_layernorm.weight": torch.randn(1)} + + result = bridge.maybe_modify_converted_hf_weight(task, dict(converted), mock_pretrained_v3.state) - assert inv_key in result - assert torch.allclose(result[inv_key], expected) + assert result[inv_key].shape == source_inv_freq.shape + assert result[inv_key].dtype == source_inv_freq.dtype + assert torch.equal(result[inv_key], source_inv_freq) def test_export_skips_inv_freq_for_non_layernorm(self, mock_pretrained_v3): bridge = DeepSeekV3Bridge() diff --git a/tests/unit_tests/models/hf_pretrained/test_base.py b/tests/unit_tests/models/hf_pretrained/test_base.py index 2f7388b8a4..6578d7bc3f 100644 --- a/tests/unit_tests/models/hf_pretrained/test_base.py +++ b/tests/unit_tests/models/hf_pretrained/test_base.py @@ -20,6 +20,7 @@ from pathlib import Path from unittest.mock import Mock, patch +import pytest from transformers.configuration_utils import PretrainedConfig @@ -274,6 +275,42 @@ def test_save_artifacts_without_model_name_or_path(): print("✅ test_save_artifacts_without_model_name_or_path passed") +def test_save_artifacts_preserves_source_generation_config_after_validation_failure(): + """Test invalid loaded generation configs are copied unchanged from the source.""" + with tempfile.TemporaryDirectory() as tmp_dir: + tmp_path = Path(tmp_dir) + source_dir = tmp_path / "source" + source_dir.mkdir() + generation_config = '{"do_sample": false, "top_p": 0.95}\n' + (source_dir / "generation_config.json").write_text(generation_config) + + target_dir = tmp_path / "target" + base = MockPreTrainedBase(model_name_or_path=str(source_dir)) + base._config = Mock() + base._generation_config = Mock() + base._generation_config.save_pretrained.side_effect = ValueError("invalid generation config") + + base.save_artifacts(target_dir) + + assert (target_dir / "generation_config.json").read_text() == generation_config + + +def test_save_artifacts_reraises_generation_config_error_without_source_file(): + """Test generation config validation errors remain fatal without a source artifact.""" + with tempfile.TemporaryDirectory() as tmp_dir: + tmp_path = Path(tmp_dir) + source_dir = tmp_path / "source" + source_dir.mkdir() + + base = MockPreTrainedBase(model_name_or_path=str(source_dir)) + base._config = Mock() + base._generation_config = Mock() + base._generation_config.save_pretrained.side_effect = ValueError("invalid generation config") + + with pytest.raises(ValueError, match="invalid generation config"): + base.save_artifacts(tmp_path / "target") + + def test_copy_handles_permission_errors(): """Test that copy failures are handled gracefully.""" with tempfile.TemporaryDirectory() as tmp_dir: diff --git a/tests/unit_tests/models/hybrid/test_hybrid_provider.py b/tests/unit_tests/models/hybrid/test_hybrid_provider.py index f8d34ef508..4742ef7d6c 100644 --- a/tests/unit_tests/models/hybrid/test_hybrid_provider.py +++ b/tests/unit_tests/models/hybrid/test_hybrid_provider.py @@ -40,6 +40,7 @@ def test_hybrid_provider_initialization(self): assert provider.fp16 is False assert provider.bf16 is True assert provider.mamba_num_groups == 8 + assert provider.mamba_chunk_size == 128 assert provider.hybrid_layer_pattern is None assert provider.hybrid_stack_spec is None assert provider.seq_length == 8192 @@ -111,6 +112,26 @@ def test_provide_method_with_vocab_padding(self): mock_calc_vocab.assert_called_once_with(50000, 128, 8) assert mock_model.call_args.kwargs["vocab_size"] == 50176 + def test_nondefault_mamba_chunk_size_is_applied_without_mutating_default_spec(self): + provider = HybridModelProvider( + num_layers=2, + hidden_size=128, + num_attention_heads=1, + vocab_size=1000, + tensor_model_parallel_size=1, + mamba_chunk_size=256, + ) + provider._pg_collection = type("PG", (), {"pp": object()})() + + with patch("megatron.bridge.models.hybrid.hybrid_provider.MCoreHybridModel") as mock_model: + provider.provide(pre_process=True, post_process=True) + + configured_spec = mock_model.call_args.kwargs["hybrid_stack_spec"] + configured_mixer = configured_spec.submodules.mamba_layer.submodules.mixer + default_mixer = hybrid_provider.default_hybrid_stack_spec.submodules.mamba_layer.submodules.mixer + assert configured_mixer.params["chunk_size"] == 256 + assert "chunk_size" not in default_mixer.params + @patch("megatron.bridge.models.hybrid.hybrid_provider.is_pp_first_stage", return_value=True) @patch("megatron.bridge.models.hybrid.hybrid_provider.is_pp_last_stage", return_value=True) def test_provide_method_respects_explicit_pipeline_stages(self, *_): diff --git a/tests/unit_tests/models/nemotronh/test_nemotron_h_bridge.py b/tests/unit_tests/models/nemotronh/test_nemotron_h_bridge.py index 37c550e3c2..b668f315ae 100644 --- a/tests/unit_tests/models/nemotronh/test_nemotron_h_bridge.py +++ b/tests/unit_tests/models/nemotronh/test_nemotron_h_bridge.py @@ -169,6 +169,7 @@ def test_provider_bridge_mamba_config(self, mock_pretrained_nemotronh, mock_nemo assert result.mamba_head_dim == mock_nemotronh_config.mamba_head_dim assert result.mamba_num_heads == mock_nemotronh_config.mamba_num_heads assert result.mamba_num_groups == mock_nemotronh_config.n_groups + assert result.mamba_chunk_size == mock_nemotronh_config.chunk_size assert result.hybrid_layer_pattern == mock_nemotronh_config.hybrid_override_pattern def test_provider_bridge_mlp_config(self, mock_pretrained_nemotronh, mock_nemotronh_config): diff --git a/tests/unit_tests/models/qwen/test_qwen3_moe_bridge.py b/tests/unit_tests/models/qwen/test_qwen3_moe_bridge.py index e9f23d16e6..b1be9898eb 100644 --- a/tests/unit_tests/models/qwen/test_qwen3_moe_bridge.py +++ b/tests/unit_tests/models/qwen/test_qwen3_moe_bridge.py @@ -1,4 +1,4 @@ -# Copyright (c) 2025, NVIDIA CORPORATION. All rights reserved. +# Copyright (c) 2025-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. diff --git a/tests/unit_tests/models/test_auto_bridge.py b/tests/unit_tests/models/test_auto_bridge.py index 6158352d8c..2e1f957321 100644 --- a/tests/unit_tests/models/test_auto_bridge.py +++ b/tests/unit_tests/models/test_auto_bridge.py @@ -662,6 +662,23 @@ def test_to_megatron_provider_sets_hf_model_id_from_pretrained(self): assert provider.hf_model_id == "hf/model-id" + def test_to_megatron_provider_persists_hf_model_revision(self): + """to_megatron_provider records the immutable revision used for config loading.""" + revision = "b968826d9c46dd6066d109eabc6255188de91218" # pragma: allowlist secret + mock_hf_model = PreTrainedCausalLM(model_name_or_path="hf/model-id", revision=revision) + mock_model_bridge = Mock() + mock_provider = Mock(spec=GPTModelProvider) + mock_provider.hf_model_id = None + mock_provider.hf_model_revision = None + mock_model_bridge.provider_bridge.return_value = mock_provider + + with patch.object(AutoBridge, "_model_bridge", mock_model_bridge): + bridge = AutoBridge(mock_hf_model) + provider = bridge.to_megatron_provider(load_weights=False) + + assert provider.hf_model_id == "hf/model-id" + assert provider.hf_model_revision == revision + def test_get_hf_model_id_from_checkpoint_delegates(self): """AutoBridge helper delegates to checkpoint utilities.""" with patch( diff --git a/tests/unit_tests/recipes/nemotronh/test_nemotron_3_nano_4b.py b/tests/unit_tests/recipes/nemotronh/test_nemotron_3_nano_4b.py new file mode 100644 index 0000000000..2679ff1db1 --- /dev/null +++ b/tests/unit_tests/recipes/nemotronh/test_nemotron_3_nano_4b.py @@ -0,0 +1,113 @@ +"""Unit tests for the dense Nemotron 3 Nano 4B recipe configurations.""" + +import pytest + +from megatron.bridge.models.hybrid.hybrid_provider import HybridModelProvider +from megatron.bridge.peft.lora import LoRA +from megatron.bridge.recipes.nemotronh import ( + nemotron_3_nano_4b_peft_config, + nemotron_3_nano_4b_pretrain_config, + nemotron_3_nano_4b_sft_32k_config, + nemotron_3_nano_4b_sft_config, +) +from megatron.bridge.training.config import ConfigContainer + + +EXPECTED_PATTERN = "M-M-M-MM-M-M*-M-M*-M-M-M*-M-M-MM*-MMM-M-M-" + + +def _assert_exact_architecture(config: ConfigContainer) -> None: + assert isinstance(config, ConfigContainer) + assert isinstance(config.model, HybridModelProvider) + assert config.model.hybrid_layer_pattern == EXPECTED_PATTERN + assert config.model.num_layers == 42 + assert config.model.hidden_size == 3136 + assert config.model.ffn_hidden_size == 12544 + assert config.model.num_attention_heads == 40 + assert config.model.num_query_groups == 8 + assert config.model.kv_channels == 128 + assert config.model.mamba_num_heads == 96 + assert config.model.mamba_head_dim == 80 + assert config.model.mamba_state_dim == 128 + assert config.model.mamba_num_groups == 8 + assert config.model.mamba_chunk_size == 256 + assert config.model.vocab_size == 131072 + assert config.model.share_embeddings_and_output_weights is False + assert config.model.position_embedding_type == "none" + assert config.tokenizer.tokenizer_model == "nvidia/NVIDIA-Nemotron-3-Nano-4B-BF16" + assert config.tokenizer.hf_tokenizer_kwargs["revision"] == ( + "dfaf35de3e30f1867dd8dbc38a7fc9fb52d3914f" # pragma: allowlist secret + ) + + +@pytest.mark.unit +def test_pretrain_config_matches_exact_architecture_and_bounded_contract() -> None: + config = nemotron_3_nano_4b_pretrain_config() + + _assert_exact_architecture(config) + assert config.model.tensor_model_parallel_size == 1 + assert config.model.context_parallel_size == 1 + assert config.dataset.seq_length == 4096 + assert config.train.train_iters == 100 + assert config.train.global_batch_size == 1024 + assert config.train.micro_batch_size == 1 + assert config.scheduler.lr_warmup_iters == 40 + assert config.scheduler.lr_decay_iters == 100 + assert config.checkpoint.save_interval == 50 + assert config.checkpoint.load is None + + +@pytest.mark.unit +def test_sft_config_matches_packed_cohort_contract() -> None: + config = nemotron_3_nano_4b_sft_config() + + _assert_exact_architecture(config) + assert config.dataset.seq_length == 2048 + assert config.dataset.offline_packing_specs.packed_sequence_size == 2048 + assert config.dataset.offline_packing_specs.pad_seq_to_mult == 1 + assert config.train.train_iters == 100 + assert config.train.global_batch_size == 32 + assert config.train.micro_batch_size == 1 + assert config.optimizer.lr == 5.0e-6 + assert config.checkpoint.save_interval == 100 + + +@pytest.mark.unit +def test_long_context_sft_config_combines_packing_and_context_parallelism() -> None: + config = nemotron_3_nano_4b_sft_32k_config() + + _assert_exact_architecture(config) + assert config.model.context_parallel_size == 2 + assert config.model.cp_comm_type == "a2a" + assert config.model.calculate_per_token_loss is True + assert config.model.cross_entropy_loss_fusion is False + assert config.dataset.seq_length == 32768 + assert config.dataset.offline_packing_specs.packed_sequence_size == 32768 + assert config.dataset.offline_packing_specs.pad_seq_to_mult == 4 + assert config.train.global_batch_size == 8 + assert config.ddp.average_in_collective is False + + +@pytest.mark.unit +def test_peft_config_uses_attention_only_cohort_lora() -> None: + config = nemotron_3_nano_4b_peft_config() + + _assert_exact_architecture(config) + assert isinstance(config.peft, LoRA) + assert config.peft.dim == 8 + assert config.peft.alpha == 16 + assert config.peft.dropout == 0.0 + assert config.peft.target_modules == ["linear_qkv", "linear_proj"] + assert config.dataset.offline_packing_specs.pad_seq_to_mult == 4 + assert config.train.train_iters == 100 + assert config.train.global_batch_size == 32 + assert config.optimizer.lr == 1.0e-4 + + +@pytest.mark.unit +def test_peft_config_preserves_custom_peft_instance() -> None: + custom = LoRA(target_modules=["linear_qkv"], dim=4, alpha=8) + + config = nemotron_3_nano_4b_peft_config(custom) + + assert config.peft is custom diff --git a/tests/unit_tests/recipes/test_moonlight_recipes.py b/tests/unit_tests/recipes/test_moonlight_recipes.py index 56a850b408..fcbf97a159 100644 --- a/tests/unit_tests/recipes/test_moonlight_recipes.py +++ b/tests/unit_tests/recipes/test_moonlight_recipes.py @@ -1,4 +1,4 @@ -# Copyright (c) 2025, NVIDIA CORPORATION. All rights reserved. +# 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. @@ -15,7 +15,7 @@ # # Test purpose: # - Parametrize over all exported Moonlight recipe functions in `megatron.bridge.recipes.moonlight`. -# - For each recipe, monkeypatch `MLAModelProvider` with a lightweight fake to avoid I/O. +# - For each recipe, monkeypatch `AutoBridge` with a lightweight fake to avoid I/O. # - Build a config with small, safe overrides and assert it forms a valid `ConfigContainer`. # - Verify tokenizer selection and sanity-check parallelism fields. # @@ -24,6 +24,7 @@ from typing import Callable import pytest +import torch from tests.unit_tests.recipes.recipe_test_utils import patch_recipe_module_global @@ -38,7 +39,7 @@ # Moonlight SFT-specific tests _MOONLIGHT_SFT_FUNCS = [ getattr(_moonlight_module, name) - for name in ["moonlight_16b_sft_config"] + for name in ["moonlight_16b_sft_config", "moonlight_16b_sft_8k_config"] if callable(getattr(_moonlight_module, name, None)) ] @@ -64,7 +65,7 @@ def _apply_test_overrides(cfg, name: str): cfg.train.train_iters = 10 cfg.train.micro_batch_size = 1 cfg.dataset.seq_length = 64 - cfg.scheduler.min_lr = 1e-5 + cfg.optimizer.min_lr = 1e-5 cfg.scheduler.lr_warmup_iters = 2 cfg.optimizer.lr = 1e-4 cfg.logger.name = f"unit_{name}" @@ -84,7 +85,21 @@ def __init__(self, *args, **kwargs): setattr(self, key, value) # Set required attributes - self.vocab_size = 151936 # Default vocab size for Moonlight + self.vocab_size = 163840 + self.kv_channels = 128 + self.multi_latent_attention = True + self.q_lora_rank = None + self.num_moe_experts = 64 + self.moe_router_topk = 6 + self.moe_router_num_groups = 1 + self.moe_router_group_topk = 1 + self.moe_router_topk_scaling_factor = 2.446 + self.moe_aux_loss_coeff = 0.001 + self.moe_router_pre_softmax = True + self.moe_router_load_balancing_type = "seq_aux_loss" + self.moe_router_score_function = "sigmoid" + self.moe_router_enable_expert_bias = True + self.moe_router_bias_update_rate = 1e-3 self.account_for_embedding_in_pipeline_split = False self.account_for_loss_in_pipeline_split = False self.num_layers_in_first_pipeline_stage = None @@ -110,6 +125,22 @@ def finalize(self): return None +class _FakeBridge: + """Return the checkpoint-compatible fake provider without model I/O.""" + + @classmethod + def from_hf_pretrained(cls, model_id: str, **kwargs) -> "_FakeBridge": + assert model_id == "moonshotai/Moonlight-16B-A3B" + assert kwargs == { + "revision": "476b36a473d4467f94469414bef6cee75c9c8172" # pragma: allowlist secret + } + return cls() + + def to_megatron_provider(self, *, load_weights: bool) -> _FakeMoonlightModelProvider: + assert load_weights is False + return _FakeMoonlightModelProvider() + + def _assert_basic_config(cfg): from megatron.bridge.training.config import ConfigContainer @@ -134,8 +165,8 @@ def test_each_moonlight_recipe_builds_config(recipe_func: Callable, monkeypatch: module_name = recipe_func.__module__ mod = importlib.import_module(module_name) - # Monkeypatch the MLAModelProvider class - patch_recipe_module_global(monkeypatch, mod, "MLAModelProvider", _FakeMoonlightModelProvider) + # Keep every recipe construction offline. + patch_recipe_module_global(monkeypatch, mod, "AutoBridge", _FakeBridge) func_name = recipe_func.__name__ is_peft = "peft" in func_name.lower() @@ -170,7 +201,7 @@ def test_moonlight_sft_config_builds(recipe_func: Callable, monkeypatch: pytest. """Test that each Moonlight SFT recipe builds a valid config.""" module_name = recipe_func.__module__ mod = importlib.import_module(module_name) - patch_recipe_module_global(monkeypatch, mod, "MLAModelProvider", _FakeMoonlightModelProvider) + patch_recipe_module_global(monkeypatch, mod, "AutoBridge", _FakeBridge) cfg = recipe_func() _apply_test_overrides(cfg, recipe_func.__name__) @@ -195,7 +226,7 @@ def test_moonlight_peft_config_builds(recipe_func: Callable, monkeypatch: pytest """Test that each Moonlight PEFT recipe builds a valid config.""" module_name = recipe_func.__module__ mod = importlib.import_module(module_name) - patch_recipe_module_global(monkeypatch, mod, "MLAModelProvider", _FakeMoonlightModelProvider) + patch_recipe_module_global(monkeypatch, mod, "AutoBridge", _FakeBridge) cfg = recipe_func(peft_scheme="lora") _apply_test_overrides(cfg, recipe_func.__name__) @@ -221,7 +252,7 @@ def test_moonlight_peft_schemes(recipe_func: Callable, peft_scheme: str, monkeyp """Test that PEFT configurations are correctly applied with different schemes.""" module_name = recipe_func.__module__ mod = importlib.import_module(module_name) - patch_recipe_module_global(monkeypatch, mod, "MLAModelProvider", _FakeMoonlightModelProvider) + patch_recipe_module_global(monkeypatch, mod, "AutoBridge", _FakeBridge) cfg = recipe_func(peft_scheme=peft_scheme) _apply_test_overrides(cfg, recipe_func.__name__) @@ -232,109 +263,297 @@ def test_moonlight_peft_schemes(recipe_func: Callable, peft_scheme: str, monkeyp assert cfg.peft is not None -def test_moonlight_16b_peft_lora_defaults(monkeypatch: pytest.MonkeyPatch): - """Test that Moonlight-16B LoRA has correct default parallelism.""" - from megatron.bridge.recipes.moonlight import moonlight_16b_peft_config - - mod = importlib.import_module("megatron.bridge.recipes.moonlight.moonlight_16b") - patch_recipe_module_global(monkeypatch, mod, "MLAModelProvider", _FakeMoonlightModelProvider) - - cfg = moonlight_16b_peft_config(peft_scheme="lora") - _apply_test_overrides(cfg, "moonlight_16b_peft_config") - - _assert_basic_config(cfg) - - # For LoRA, Moonlight-16B should use TP=1, PP=1, EP=2 +def _assert_moonlight_router_identity(cfg, *, aux_loss_coeff: float): + """Assert that cohort alignment did not replace Moonlight routing semantics.""" + assert cfg.model.num_moe_experts == 64 + assert cfg.model.moe_router_topk == 6 + assert cfg.model.moe_router_num_groups == 1 + assert cfg.model.moe_router_group_topk == 1 + assert cfg.model.moe_router_topk_scaling_factor == 2.446 + assert cfg.model.moe_aux_loss_coeff == aux_loss_coeff + assert cfg.model.moe_router_pre_softmax is True + assert cfg.model.moe_router_load_balancing_type == "seq_aux_loss" + assert cfg.model.moe_router_score_function == "sigmoid" + assert cfg.model.moe_router_enable_expert_bias is True + assert cfg.model.moe_router_bias_update_rate == 1e-3 + assert cfg.model.moe_router_force_load_balancing is False + + +def _assert_finetuning_optimizer_contract(cfg): + """Assert the shared SFT/PEFT optimizer and arithmetic contract.""" + assert cfg.optimizer.optimizer == "adam" + assert cfg.optimizer.adam_beta1 == 0.9 + assert cfg.optimizer.adam_beta2 == 0.95 + assert cfg.optimizer.adam_eps == 1e-8 + assert cfg.optimizer.weight_decay == 0.1 + assert cfg.optimizer.clip_grad == 1.0 + assert cfg.optimizer.use_distributed_optimizer is True + assert cfg.optimizer.use_precision_aware_optimizer is False + assert cfg.optimizer.main_params_dtype == torch.float32 + assert cfg.optimizer.main_grads_dtype == torch.float32 + assert cfg.optimizer.exp_avg_dtype == torch.float32 + assert cfg.optimizer.exp_avg_sq_dtype == torch.float32 + assert cfg.scheduler.start_weight_decay == 0.033 + assert cfg.scheduler.end_weight_decay == 0.033 + assert cfg.scheduler.weight_decay_incr_style == "constant" + assert cfg.scheduler.lr_decay_style == "cosine" + assert cfg.scheduler.lr_warmup_init == 0.0 + assert cfg.mixed_precision.bf16 is True + assert cfg.mixed_precision.params_dtype == torch.bfloat16 + assert cfg.mixed_precision.pipeline_dtype == torch.bfloat16 + assert cfg.mixed_precision.grad_reduce_in_fp32 is True + assert cfg.ddp.grad_reduce_in_fp32 is True + + +def test_moonlight_16b_pretrain_convergence_contract(monkeypatch: pytest.MonkeyPatch): + """Test the exact 16-GPU pretrain convergence contract.""" + from megatron.bridge.recipes.moonlight import moonlight_16b_pretrain_config + + patch_recipe_module_global(monkeypatch, moonlight_16b_pretrain_config, "AutoBridge", _FakeBridge) + cfg = moonlight_16b_pretrain_config() + + assert moonlight_16b_pretrain_config.__name__ == "moonlight_16b_pretrain_16gpu_h100_bf16_config" assert cfg.model.tensor_model_parallel_size == 1 assert cfg.model.pipeline_model_parallel_size == 1 - assert cfg.model.expert_model_parallel_size == 2 + assert cfg.model.pipeline_model_parallel_layout is None + assert cfg.model.context_parallel_size == 1 + assert cfg.model.expert_model_parallel_size == 16 + assert cfg.model.expert_tensor_parallel_size == 1 assert cfg.model.sequence_parallel is False + assert cfg.model.seq_length == 4096 + assert cfg.train.train_iters == 100 + assert cfg.train.global_batch_size == 1024 + assert cfg.train.micro_batch_size == 1 + assert cfg.train.global_batch_size // 16 == 64 + assert cfg.dataset.random_seed == 1234 + assert cfg.rng.seed == 1234 + assert cfg.optimizer.optimizer == "adam" + assert cfg.optimizer.lr == 3e-4 + assert cfg.optimizer.min_lr == 3e-5 + assert cfg.optimizer.adam_beta1 == 0.9 + assert cfg.optimizer.adam_beta2 == 0.95 + assert cfg.optimizer.adam_eps == 1e-8 + assert cfg.optimizer.weight_decay == 0.1 + assert cfg.optimizer.clip_grad == 1.0 + assert cfg.optimizer.use_distributed_optimizer is True + assert cfg.optimizer.use_precision_aware_optimizer is True + assert cfg.optimizer.main_params_dtype == torch.float32 + assert cfg.optimizer.main_grads_dtype == torch.float32 + assert cfg.optimizer.exp_avg_dtype == torch.float32 + assert cfg.optimizer.exp_avg_sq_dtype == torch.float32 + assert cfg.scheduler.lr_warmup_iters == 40 + assert cfg.scheduler.lr_decay_iters == 100 + assert cfg.scheduler.lr_decay_style == "cosine" + assert cfg.scheduler.lr_warmup_init == 0.0 + assert cfg.scheduler.start_weight_decay == 0.033 + assert cfg.scheduler.end_weight_decay == 0.033 + assert cfg.scheduler.weight_decay_incr_style == "constant" + assert cfg.checkpoint.save_interval == 50 + assert cfg.checkpoint.load is None + assert cfg.mixed_precision.bf16 is True + assert cfg.mixed_precision.params_dtype == torch.bfloat16 + assert cfg.mixed_precision.grad_reduce_in_fp32 is False + assert cfg.ddp.grad_reduce_in_fp32 is False + _assert_moonlight_router_identity(cfg, aux_loss_coeff=0.001) + + +def test_moonlight_16b_sft_convergence_contract(monkeypatch: pytest.MonkeyPatch): + """Test the exact 8-GPU full-SFT convergence contract.""" + from megatron.bridge.recipes.moonlight import moonlight_16b_sft_config - # Check manual GC is enabled - assert cfg.train.manual_gc is True - assert cfg.train.manual_gc_interval == 5 - - -def test_moonlight_16b_peft_dora_defaults(monkeypatch: pytest.MonkeyPatch): - """Test that Moonlight-16B DoRA has correct default parallelism.""" - from megatron.bridge.recipes.moonlight import moonlight_16b_peft_config - - mod = importlib.import_module("megatron.bridge.recipes.moonlight.moonlight_16b") - patch_recipe_module_global(monkeypatch, mod, "MLAModelProvider", _FakeMoonlightModelProvider) - - cfg = moonlight_16b_peft_config(peft_scheme="dora") - _apply_test_overrides(cfg, "moonlight_16b_peft_config") - - _assert_basic_config(cfg) + patch_recipe_module_global(monkeypatch, moonlight_16b_sft_config, "AutoBridge", _FakeBridge) + cfg = moonlight_16b_sft_config() - # For DoRA, Moonlight-16B should use TP=1, PP=1, EP=2 (same as LoRA) + assert moonlight_16b_sft_config.__name__ == "moonlight_16b_sft_8gpu_h100_bf16_tp1_config" assert cfg.model.tensor_model_parallel_size == 1 assert cfg.model.pipeline_model_parallel_size == 1 - assert cfg.model.expert_model_parallel_size == 2 + assert cfg.model.pipeline_model_parallel_layout is None + assert cfg.model.context_parallel_size == 1 + assert cfg.model.expert_model_parallel_size == 8 + assert cfg.model.expert_tensor_parallel_size == 1 assert cfg.model.sequence_parallel is False + assert cfg.model.vocab_size == 163842 + assert cfg.get_data_parallel_size(8) == 8 + assert cfg.train.global_batch_size // (cfg.train.micro_batch_size * cfg.get_data_parallel_size(8)) == 4 + assert cfg.model.num_moe_experts // cfg.model.expert_model_parallel_size == 8 + assert cfg.model.seq_length == 2048 + assert cfg.dataset.seq_length == 2048 + assert cfg.dataset.offline_packing_specs.packed_sequence_size == 2048 + assert cfg.dataset.offline_packing_specs.pad_seq_to_mult == 1 + assert cfg.train.train_iters == 100 + assert cfg.train.global_batch_size == 32 + assert cfg.train.micro_batch_size == 1 + assert cfg.dataset.seed == 1234 + assert cfg.rng.seed == 5678 + assert cfg.optimizer.lr == 5e-6 + assert cfg.optimizer.min_lr == 0.0 + assert cfg.scheduler.lr_warmup_iters == 10 + assert cfg.scheduler.lr_decay_iters == 100 + assert cfg.checkpoint.save_interval == 100 + assert cfg.checkpoint.load is None + assert cfg.tokenizer.tokenizer_type == "HuggingFaceTokenizer" + assert cfg.tokenizer.tokenizer_model == "moonshotai/Moonlight-16B-A3B" + assert cfg.tokenizer.hf_tokenizer_kwargs == { + "revision": "476b36a473d4467f94469414bef6cee75c9c8172", # pragma: allowlist secret + "trust_remote_code": True, + } + from megatron.bridge.training.utils.omegaconf_utils import process_config_with_overrides + + process_config_with_overrides( + cfg.tokenizer, + cli_overrides=[ + '++hf_tokenizer_kwargs.revision="476b36a473d4467f94469414bef6cee75c9c8172"' # pragma: allowlist secret + ], + ) + assert cfg.model.moe_token_dispatcher_type == "alltoall" + assert cfg.model.moe_flex_dispatcher_backend is None + assert cfg.model.moe_hybridep_num_sms is None + assert cfg.model.moe_flex_dispatcher_num_sms is None + assert cfg.model.moe_a2a_overlap is False + assert cfg.model.moe_shared_expert_overlap is False + assert cfg.comm_overlap is None + _assert_finetuning_optimizer_contract(cfg) + _assert_moonlight_router_identity(cfg, aux_loss_coeff=0.001) + + +def test_moonlight_16b_legacy_sft_contract_remains_available(monkeypatch: pytest.MonkeyPatch): + """Keep the existing TP4/PP2 support topology available to callers.""" + from megatron.bridge.recipes.moonlight.h100.moonlight_16b import ( + moonlight_16b_sft_8gpu_h100_bf16_config, + ) + + patch_recipe_module_global(monkeypatch, moonlight_16b_sft_8gpu_h100_bf16_config, "AutoBridge", _FakeBridge) + cfg = moonlight_16b_sft_8gpu_h100_bf16_config() + + assert cfg.model.tensor_model_parallel_size == 4 + assert cfg.model.pipeline_model_parallel_size == 2 + assert cfg.model.pipeline_model_parallel_layout == [ + ["embedding"] + ["decoder"] * 14, + ["decoder"] * 13 + ["loss"], + ] + assert cfg.model.expert_model_parallel_size == 4 + assert cfg.model.sequence_parallel is True + assert cfg.model.vocab_size == 163844 + assert cfg.dataset.offline_packing_specs.pad_seq_to_mult == 4 - # Check manual GC is enabled - assert cfg.train.manual_gc is True - assert cfg.train.manual_gc_interval == 5 - - -def test_moonlight_16b_sft_full_defaults(monkeypatch: pytest.MonkeyPatch): - """Test that Moonlight-16B full SFT has correct default parallelism.""" - from megatron.bridge.recipes.moonlight import moonlight_16b_sft_config - - mod = importlib.import_module("megatron.bridge.recipes.moonlight.moonlight_16b") - patch_recipe_module_global(monkeypatch, mod, "MLAModelProvider", _FakeMoonlightModelProvider) - cfg = moonlight_16b_sft_config() - _apply_test_overrides(cfg, "moonlight_16b_sft_config") +def test_moonlight_16b_peft_convergence_contract(monkeypatch: pytest.MonkeyPatch): + """Test the exact 4-GPU LoRA convergence contract and frozen base.""" + from megatron.bridge.peft.lora import LoRA + from megatron.bridge.peft.lora_layers import LinearAdapter + from megatron.bridge.recipes.moonlight import moonlight_16b_peft_config - _assert_basic_config(cfg) + patch_recipe_module_global(monkeypatch, moonlight_16b_peft_config, "AutoBridge", _FakeBridge) + cfg = moonlight_16b_peft_config(peft_scheme="lora") - # For full SFT, Moonlight-16B should use TP=2, PP=1, EP=8 + assert moonlight_16b_peft_config.__name__ == "moonlight_16b_peft_4gpu_h100_bf16_config" + assert cfg.model.tensor_model_parallel_size == 4 + assert cfg.model.pipeline_model_parallel_size == 1 + assert cfg.model.pipeline_model_parallel_layout is None + assert cfg.model.context_parallel_size == 1 + assert cfg.model.expert_model_parallel_size == 4 + assert cfg.model.expert_tensor_parallel_size == 1 + assert cfg.model.sequence_parallel is True + assert cfg.model.vocab_size == 163844 + assert cfg.model.seq_length == 2048 + assert cfg.dataset.seq_length == 2048 + assert cfg.dataset.offline_packing_specs.packed_sequence_size == 2048 + assert cfg.dataset.offline_packing_specs.pad_seq_to_mult == 4 + assert cfg.train.train_iters == 100 + assert cfg.train.global_batch_size == 32 + assert cfg.train.micro_batch_size == 1 + assert cfg.dataset.seed == 1234 + assert cfg.rng.seed == 5678 + assert cfg.optimizer.lr == 1e-4 + assert cfg.optimizer.min_lr == 0.0 + assert cfg.scheduler.lr_warmup_iters == 10 + assert cfg.scheduler.lr_decay_iters == 100 + assert cfg.checkpoint.save_interval == 100 + assert cfg.checkpoint.load is None + assert cfg.tokenizer.tokenizer_model == "moonshotai/Moonlight-16B-A3B" + assert cfg.tokenizer.hf_tokenizer_kwargs == { + "revision": "476b36a473d4467f94469414bef6cee75c9c8172", # pragma: allowlist secret + "trust_remote_code": True, + } + assert isinstance(cfg.peft, LoRA) + expected_targets = ["linear_q_proj", "linear_kv_down_proj", "linear_kv_up_proj", "linear_proj"] + assert cfg.peft.target_modules == expected_targets + assert cfg.peft.dim == 8 + assert cfg.peft.alpha == 16 + assert cfg.peft.dropout == 0.0 + base_model = torch.nn.Module() + for target in expected_targets + ["linear_qkv"]: + setattr(base_model, target, torch.nn.Linear(2, 2)) + cfg.peft(base_model) + assert all(isinstance(getattr(base_model, target), LinearAdapter) for target in expected_targets) + assert isinstance(base_model.linear_qkv, torch.nn.Linear) + assert all(not parameter.requires_grad for parameter in base_model.linear_qkv.parameters()) + _assert_finetuning_optimizer_contract(cfg) + _assert_moonlight_router_identity(cfg, aux_loss_coeff=0.001) + + +def test_moonlight_16b_sft_8k_contract_is_separate(monkeypatch: pytest.MonkeyPatch): + """Test that 8K/CP2 SFT retains its prior batch, topology, and precision.""" + from megatron.bridge.recipes.moonlight import moonlight_16b_sft_8k_config + + patch_recipe_module_global(monkeypatch, moonlight_16b_sft_8k_config, "AutoBridge", _FakeBridge) + cfg = moonlight_16b_sft_8k_config() + + assert moonlight_16b_sft_8k_config.__name__ == "moonlight_16b_sft_8gpu_h100_bf16_8k_config" assert cfg.model.tensor_model_parallel_size == 2 assert cfg.model.pipeline_model_parallel_size == 1 + assert cfg.model.pipeline_model_parallel_layout is None + assert cfg.model.context_parallel_size == 2 assert cfg.model.expert_model_parallel_size == 8 assert cfg.model.sequence_parallel is True - - # Check manual GC is enabled - assert cfg.train.manual_gc is True - assert cfg.train.manual_gc_interval == 5 - - -def test_moonlight_16b_sft_precision_aware_optimizer(monkeypatch: pytest.MonkeyPatch): - """Test that Moonlight-16B SFT uses precision-aware optimizer settings.""" - from megatron.bridge.recipes.moonlight import moonlight_16b_sft_config - - mod = importlib.import_module("megatron.bridge.recipes.moonlight.moonlight_16b") - patch_recipe_module_global(monkeypatch, mod, "MLAModelProvider", _FakeMoonlightModelProvider) - - cfg = moonlight_16b_sft_config() - _apply_test_overrides(cfg, "moonlight_16b_sft_config") - - _assert_basic_config(cfg) - - # Check precision-aware optimizer settings + assert cfg.model.vocab_size == 163842 + assert cfg.model.seq_length == 8192 + assert cfg.dataset.seq_length == 8192 + assert cfg.dataset.offline_packing_specs.packed_sequence_size == 8192 + assert cfg.dataset.offline_packing_specs.pad_seq_to_mult == 4 + assert cfg.train.train_iters == 20 + assert cfg.train.global_batch_size == 128 + assert cfg.train.micro_batch_size == 1 + assert cfg.optimizer.lr == 1e-6 + assert cfg.optimizer.min_lr == 0.0 + assert cfg.optimizer.adam_beta1 == 0.9 + assert cfg.optimizer.adam_beta2 == 0.98 + assert cfg.optimizer.adam_eps == 1e-5 assert cfg.optimizer.use_precision_aware_optimizer is True - import torch - assert cfg.optimizer.main_params_dtype == torch.float32 assert cfg.optimizer.main_grads_dtype == torch.bfloat16 assert cfg.optimizer.exp_avg_dtype == torch.bfloat16 assert cfg.optimizer.exp_avg_sq_dtype == torch.bfloat16 - - -def test_moonlight_16b_sft_tokenizer_with_trust_remote_code(monkeypatch: pytest.MonkeyPatch): - """Test that Moonlight-16B SFT uses HF tokenizer with trust_remote_code.""" - from megatron.bridge.recipes.moonlight import moonlight_16b_sft_config - - mod = importlib.import_module("megatron.bridge.recipes.moonlight.moonlight_16b") - patch_recipe_module_global(monkeypatch, mod, "MLAModelProvider", _FakeMoonlightModelProvider) - - cfg = moonlight_16b_sft_config() - - _assert_basic_config(cfg) - - # Check tokenizer settings - assert cfg.tokenizer.tokenizer_type == "HuggingFaceTokenizer" - assert cfg.tokenizer.tokenizer_model == "moonshotai/Moonlight-16B-A3B" - assert cfg.tokenizer.hf_tokenizer_kwargs == {"trust_remote_code": True} + assert cfg.scheduler.lr_warmup_iters == 2 + assert cfg.scheduler.lr_decay_iters == 20 + assert cfg.mixed_precision.grad_reduce_in_fp32 is False + assert cfg.ddp.grad_reduce_in_fp32 is False + assert cfg.model.cross_entropy_loss_fusion is False + assert cfg.model.calculate_per_token_loss is True + assert cfg.ddp.average_in_collective is False + _assert_moonlight_router_identity(cfg, aux_loss_coeff=0.001) + + +def test_moonlight_compatibility_recipes_remain_exported(monkeypatch: pytest.MonkeyPatch): + """Test that the prior explicit pretrain and PEFT entry points remain available.""" + from megatron.bridge.recipes.moonlight.h100 import ( + moonlight_16b_peft_2gpu_h100_bf16_config, + moonlight_16b_pretrain_8gpu_h100_bf16_config, + ) + + patch_recipe_module_global(monkeypatch, moonlight_16b_pretrain_8gpu_h100_bf16_config, "AutoBridge", _FakeBridge) + pretrain_cfg = moonlight_16b_pretrain_8gpu_h100_bf16_config() + peft_cfg = moonlight_16b_peft_2gpu_h100_bf16_config(peft_scheme="lora") + + assert pretrain_cfg.model.tensor_model_parallel_size == 2 + assert pretrain_cfg.model.pipeline_model_parallel_size == 1 + assert pretrain_cfg.model.expert_model_parallel_size == 8 + assert pretrain_cfg.train.global_batch_size == 2048 + assert pretrain_cfg.optimizer.use_precision_aware_optimizer is True + assert pretrain_cfg.optimizer.main_grads_dtype == torch.bfloat16 + assert peft_cfg.model.tensor_model_parallel_size == 1 + assert peft_cfg.model.pipeline_model_parallel_size == 1 + assert peft_cfg.model.expert_model_parallel_size == 2 + assert peft_cfg.train.global_batch_size == 128 + assert peft_cfg.optimizer.use_precision_aware_optimizer is True + assert peft_cfg.optimizer.main_grads_dtype == torch.bfloat16 diff --git a/tests/unit_tests/recipes/test_qwen_recipes.py b/tests/unit_tests/recipes/test_qwen_recipes.py index 32e9ef55ab..6bcb9a39cd 100644 --- a/tests/unit_tests/recipes/test_qwen_recipes.py +++ b/tests/unit_tests/recipes/test_qwen_recipes.py @@ -25,6 +25,7 @@ from typing import Callable import pytest +import torch from tests.unit_tests.recipes.recipe_test_utils import patch_recipe_module_global @@ -67,8 +68,13 @@ def to_megatron_provider(self, load_weights: bool = False): return _FakeModelCfg() @staticmethod - def from_hf_pretrained(hf_path: str): - # Ignore hf_path; return a bridge that yields a fake provider + def from_hf_pretrained(hf_path: str, **kwargs): + expected_revisions = { + "Qwen/Qwen3-8B": "b968826d9c46dd6066d109eabc6255188de91218", # pragma: allowlist secret + "Qwen/Qwen3-30B-A3B": "ad44e777bcd18fa416d9da3bd8f70d33ebb85d39", # pragma: allowlist secret + } + if hf_path in expected_revisions: + assert kwargs == {"revision": expected_revisions[hf_path]} return _FakeBridge() @@ -158,6 +164,181 @@ def test_each_qwen_recipe_builds_config(recipe_func: Callable, monkeypatch: pyte assert cfg.dataset is not None +def _patch_qwen3_dense_bridge(monkeypatch: pytest.MonkeyPatch) -> None: + qwen3_mod = importlib.import_module("megatron.bridge.recipes.qwen.h100.qwen3") + patch_recipe_module_global(monkeypatch, qwen3_mod, "AutoBridge", _FakeBridge) + + +def _assert_qwen_finetune_optimizer_contract(cfg, *, expected_lr: float) -> None: + assert cfg.optimizer.optimizer == "adam" + assert cfg.optimizer.lr == expected_lr + assert cfg.optimizer.min_lr == 0.0 + assert cfg.optimizer.adam_beta1 == 0.9 + assert cfg.optimizer.adam_beta2 == 0.95 + assert cfg.optimizer.adam_eps == 1.0e-8 + assert cfg.optimizer.clip_grad == 1.0 + assert cfg.scheduler.start_weight_decay == 0.033 + assert cfg.scheduler.end_weight_decay == 0.033 + assert cfg.scheduler.weight_decay_incr_style == "constant" + assert cfg.scheduler.lr_decay_style == "cosine" + assert cfg.scheduler.lr_warmup_init == 0.0 + assert cfg.optimizer.use_precision_aware_optimizer is False + assert cfg.optimizer.main_params_dtype == torch.float32 + assert cfg.optimizer.main_grads_dtype == torch.float32 + assert cfg.optimizer.exp_avg_dtype == torch.float32 + assert cfg.optimizer.exp_avg_sq_dtype == torch.float32 + assert cfg.optimizer.use_distributed_optimizer is True + assert cfg.mixed_precision.bf16 is True + assert cfg.mixed_precision.params_dtype == torch.bfloat16 + assert cfg.mixed_precision.grad_reduce_in_fp32 is True + assert cfg.ddp.grad_reduce_in_fp32 is True + assert cfg.ddp.use_distributed_optimizer is True + + +def test_qwen3_8b_pretrain_convergence_contract(monkeypatch: pytest.MonkeyPatch): + """The generic 8B pretrain recipe should select the 16-GPU convergence cohort.""" + from megatron.bridge.recipes.qwen import qwen3_8b_pretrain_config + from megatron.bridge.recipes.qwen.h100.qwen3 import qwen3_8b_pretrain_16gpu_h100_bf16_config + + _patch_qwen3_dense_bridge(monkeypatch) + + assert qwen3_8b_pretrain_config is qwen3_8b_pretrain_16gpu_h100_bf16_config + cfg = qwen3_8b_pretrain_config() + + _assert_basic_config(cfg) + assert cfg.model.tensor_model_parallel_size == 1 + assert cfg.model.pipeline_model_parallel_size == 1 + assert cfg.model.context_parallel_size == 1 + assert cfg.model.sequence_parallel is False + assert cfg.model.seq_length == 4096 + assert cfg.dataset.seq_length == 4096 + assert cfg.train.train_iters == 100 + assert cfg.train.global_batch_size == 1024 + assert cfg.train.micro_batch_size == 1 + assert cfg.dataset.random_seed == 1234 + assert cfg.rng.seed == 1234 + assert cfg.optimizer.optimizer == "adam" + assert cfg.optimizer.lr == 3.0e-4 + assert cfg.optimizer.min_lr == 3.0e-5 + assert cfg.optimizer.adam_beta1 == 0.9 + assert cfg.optimizer.adam_beta2 == 0.95 + assert cfg.optimizer.adam_eps == 1.0e-8 + assert cfg.optimizer.clip_grad == 1.0 + assert cfg.scheduler.start_weight_decay == 0.033 + assert cfg.scheduler.end_weight_decay == 0.033 + assert cfg.scheduler.weight_decay_incr_style == "constant" + assert cfg.scheduler.lr_decay_style == "cosine" + assert cfg.scheduler.lr_warmup_init == 0.0 + assert cfg.scheduler.lr_warmup_iters == 40 + assert cfg.scheduler.lr_decay_iters == 100 + assert cfg.optimizer.use_precision_aware_optimizer is True + assert cfg.optimizer.main_params_dtype == torch.float32 + assert cfg.optimizer.main_grads_dtype == torch.float32 + assert cfg.optimizer.exp_avg_dtype == torch.float32 + assert cfg.optimizer.exp_avg_sq_dtype == torch.float32 + assert cfg.optimizer.use_distributed_optimizer is True + assert cfg.mixed_precision.bf16 is True + assert cfg.mixed_precision.params_dtype == torch.bfloat16 + assert cfg.mixed_precision.grad_reduce_in_fp32 is False + assert cfg.ddp.grad_reduce_in_fp32 is False + assert cfg.ddp.use_distributed_optimizer is True + assert cfg.checkpoint.save_interval == 50 + assert cfg.checkpoint.load is None + assert cfg.tokenizer.hf_tokenizer_kwargs == { + "revision": "b968826d9c46dd6066d109eabc6255188de91218" # pragma: allowlist secret + } + from megatron.bridge.training.utils.omegaconf_utils import process_config_with_overrides + + process_config_with_overrides( + cfg.tokenizer, + cli_overrides=[ + '++hf_tokenizer_kwargs.revision="b968826d9c46dd6066d109eabc6255188de91218"' # pragma: allowlist secret + ], + ) + + +def test_qwen3_8b_sft_convergence_contract(monkeypatch: pytest.MonkeyPatch): + """The bounded 8B SFT recipe should own the shared finetuning contract.""" + from megatron.bridge.recipes.qwen import qwen3_8b_sft_config + + _patch_qwen3_dense_bridge(monkeypatch) + cfg = qwen3_8b_sft_config() + + _assert_basic_config(cfg) + assert cfg.model.seq_length == 2048 + assert cfg.train.train_iters == 100 + assert cfg.train.global_batch_size == 32 + assert cfg.train.micro_batch_size == 1 + assert cfg.dataset.seed == 1234 + assert cfg.rng.seed == 5678 + assert cfg.dataset.offline_packing_specs.pad_seq_to_mult == 1 + assert cfg.scheduler.lr_warmup_iters == 10 + assert cfg.scheduler.lr_decay_iters == 100 + assert cfg.checkpoint.save_interval == 100 + assert cfg.checkpoint.load is None + assert cfg.tokenizer.hf_tokenizer_kwargs == { + "revision": "b968826d9c46dd6066d109eabc6255188de91218" # pragma: allowlist secret + } + _assert_qwen_finetune_optimizer_contract(cfg, expected_lr=5.0e-6) + + +def test_qwen3_8b_peft_convergence_contract(monkeypatch: pytest.MonkeyPatch): + """The bounded 8B PEFT recipe should own the optimizer and LoRA contracts.""" + from megatron.bridge.recipes.qwen import qwen3_8b_peft_config + + _patch_qwen3_dense_bridge(monkeypatch) + cfg = qwen3_8b_peft_config() + + _assert_basic_config(cfg) + assert cfg.model.seq_length == 2048 + assert cfg.train.train_iters == 100 + assert cfg.train.global_batch_size == 32 + assert cfg.train.micro_batch_size == 1 + assert cfg.dataset.seed == 1234 + assert cfg.rng.seed == 5678 + assert cfg.dataset.offline_packing_specs.pad_seq_to_mult == 4 + assert cfg.scheduler.lr_warmup_iters == 10 + assert cfg.scheduler.lr_decay_iters == 100 + assert cfg.checkpoint.save_interval == 100 + assert cfg.checkpoint.load is None + assert cfg.tokenizer.hf_tokenizer_kwargs == { + "revision": "b968826d9c46dd6066d109eabc6255188de91218" # pragma: allowlist secret + } + _assert_qwen_finetune_optimizer_contract(cfg, expected_lr=1.0e-4) + assert cfg.peft is not None + assert cfg.peft.target_modules == ["linear_qkv", "linear_proj"] + assert cfg.peft.dim == 8 + assert cfg.peft.alpha == 16 + assert cfg.peft.dropout == 0.0 + + +def test_qwen3_8b_32k_sft_preserves_separate_batch_contract(monkeypatch: pytest.MonkeyPatch): + """The 32K SFT cohort should not inherit the bounded 2K recipe's GBS=32.""" + from megatron.bridge.recipes.qwen import qwen3_8b_sft_32k_config + from megatron.bridge.recipes.qwen.h100.qwen3 import qwen3_8b_sft_8gpu_h100_bf16_32k_config + + _patch_qwen3_dense_bridge(monkeypatch) + + assert qwen3_8b_sft_32k_config is qwen3_8b_sft_8gpu_h100_bf16_32k_config + cfg = qwen3_8b_sft_32k_config() + + _assert_basic_config(cfg) + assert cfg.model.tensor_model_parallel_size == 4 + assert cfg.model.pipeline_model_parallel_size == 1 + assert cfg.model.context_parallel_size == 2 + assert cfg.model.sequence_parallel is True + assert cfg.model.cp_comm_type == "a2a" + assert cfg.model.seq_length == 32768 + assert cfg.dataset.seq_length == 32768 + assert cfg.dataset.offline_packing_specs.packed_sequence_size == 32768 + assert cfg.dataset.offline_packing_specs.pad_seq_to_mult == 8 + assert cfg.train.global_batch_size == 8 + assert cfg.train.micro_batch_size == 1 + assert cfg.model.cross_entropy_loss_fusion is False + assert cfg.model.calculate_per_token_loss is True + assert cfg.ddp.average_in_collective is False + + # Qwen3 MoE SFT and PEFT-specific tests _QWEN3_MOE_SFT_FUNCS = [ getattr(_qwen_module, name) @@ -230,7 +411,23 @@ def test_qwen3_30b_a3b_lora_defaults(monkeypatch: pytest.MonkeyPatch): assert cfg.peft is not None assert cfg.peft.dim == 8 assert cfg.peft.alpha == 16 + assert cfg.peft.dropout == 0.0 assert cfg.peft.target_modules == ["linear_qkv", "linear_proj"] + assert cfg.model.seq_length == 2048 + assert cfg.train.train_iters == 100 + assert cfg.train.global_batch_size == 32 + assert cfg.train.micro_batch_size == 1 + assert cfg.dataset.seed == 1234 + assert cfg.rng.seed == 5678 + assert cfg.dataset.offline_packing_specs.pad_seq_to_mult == 4 + assert cfg.scheduler.lr_warmup_iters == 10 + assert cfg.scheduler.lr_decay_iters == 100 + assert cfg.checkpoint.save_interval == 100 + assert cfg.checkpoint.load is None + assert cfg.tokenizer.hf_tokenizer_kwargs == { + "revision": "ad44e777bcd18fa416d9da3bd8f70d33ebb85d39" # pragma: allowlist secret + } + _assert_qwen_finetune_optimizer_contract(cfg, expected_lr=1.0e-4) def test_qwen3_30b_a3b_dora_defaults(monkeypatch: pytest.MonkeyPatch): @@ -258,7 +455,7 @@ def test_qwen3_30b_a3b_dora_defaults(monkeypatch: pytest.MonkeyPatch): def test_qwen3_30b_a3b_full_sft_defaults(monkeypatch: pytest.MonkeyPatch): - """Test that 30B-A3B full SFT has correct default parallelism.""" + """Test that generic 30B-A3B full SFT uses the verified 16-GPU config.""" from megatron.bridge.recipes.qwen import qwen3_30b_a3b_sft_config mod = importlib.import_module("megatron.bridge.recipes.qwen.qwen3_moe") @@ -268,12 +465,210 @@ def test_qwen3_30b_a3b_full_sft_defaults(monkeypatch: pytest.MonkeyPatch): _assert_basic_config(cfg) - # For full SFT, 30B-A3B should use TP=4, PP=2, EP=4 + assert cfg.model.tensor_model_parallel_size == 1 + assert cfg.model.pipeline_model_parallel_size == 1 + assert cfg.model.context_parallel_size == 1 + assert cfg.model.expert_model_parallel_size == 16 + assert cfg.model.expert_tensor_parallel_size == 1 + assert cfg.model.sequence_parallel is False + assert cfg.train.global_batch_size == 32 + assert cfg.train.micro_batch_size == 1 + assert cfg.get_data_parallel_size(16) == 16 + assert cfg.train.global_batch_size // (cfg.train.micro_batch_size * cfg.get_data_parallel_size(16)) == 2 + assert cfg.model.seq_length == 2048 + assert cfg.model.moe_token_dispatcher_type == "alltoall" + assert cfg.model.moe_flex_dispatcher_backend is None + assert cfg.model.moe_hybridep_num_sms is None + assert cfg.model.moe_flex_dispatcher_num_sms is None + assert cfg.model.moe_a2a_overlap is False + assert cfg.model.moe_shared_expert_overlap is False + assert cfg.model.bias_activation_fusion is True + assert cfg.model.apply_rope_fusion is True + assert cfg.model.moe_router_fusion is True + assert cfg.model.cuda_graph_impl == "none" + assert cfg.comm_overlap is None + assert cfg.train.train_iters == 100 + assert cfg.dataset.seed == 1234 + assert cfg.rng.seed == 5678 + assert cfg.dataset.offline_packing_specs.pad_seq_to_mult == 1 + assert cfg.scheduler.lr_warmup_iters == 10 + assert cfg.scheduler.lr_decay_iters == 100 + assert cfg.checkpoint.save_interval == 100 + assert cfg.checkpoint.load is None + assert cfg.tokenizer.hf_tokenizer_kwargs == { + "revision": "ad44e777bcd18fa416d9da3bd8f70d33ebb85d39" # pragma: allowlist secret + } + assert cfg.peft is None + _assert_qwen_finetune_optimizer_contract(cfg, expected_lr=5.0e-6) + + +def test_qwen3_30b_a3b_legacy_8gpu_full_sft_defaults(monkeypatch: pytest.MonkeyPatch): + """Keep the explicit 8-GPU SFT topology available for existing callers.""" + from megatron.bridge.recipes.qwen.h100.qwen3_moe import qwen3_30b_a3b_sft_8gpu_h100_bf16_config + + mod = importlib.import_module("megatron.bridge.recipes.qwen.h100.qwen3_moe") + patch_recipe_module_global(monkeypatch, mod, "AutoBridge", _FakeBridge) + + cfg = qwen3_30b_a3b_sft_8gpu_h100_bf16_config() + assert cfg.model.tensor_model_parallel_size == 4 assert cfg.model.pipeline_model_parallel_size == 2 assert cfg.model.expert_model_parallel_size == 4 assert cfg.model.sequence_parallel is True + assert cfg.model.moe_flex_dispatcher_backend == "deepep" assert cfg.peft is None + assert cfg.model.seq_length == 2048 + assert cfg.train.train_iters == 100 + assert cfg.train.global_batch_size == 32 + assert cfg.train.micro_batch_size == 1 + assert cfg.dataset.seed == 1234 + assert cfg.rng.seed == 5678 + assert cfg.dataset.offline_packing_specs.pad_seq_to_mult == 4 + assert cfg.scheduler.lr_warmup_iters == 10 + assert cfg.scheduler.lr_decay_iters == 100 + assert cfg.checkpoint.save_interval == 100 + assert cfg.checkpoint.load is None + _assert_qwen_finetune_optimizer_contract(cfg, expected_lr=5.0e-6) + + +def test_qwen3_30b_a3b_sft_generic_alias_uses_16gpu_factory(): + """Keep the generic SFT alias attached to the verified 16-GPU factory.""" + from megatron.bridge.recipes.qwen import qwen3_30b_a3b_sft_config + from megatron.bridge.recipes.qwen.h100.qwen3_moe import qwen3_30b_a3b_sft_16gpu_h100_bf16_config + + assert qwen3_30b_a3b_sft_config is qwen3_30b_a3b_sft_16gpu_h100_bf16_config + + +def test_qwen3_30b_a3b_pretrain_defaults(monkeypatch: pytest.MonkeyPatch): + """Test that the generic 30B-A3B pretrain recipe uses the verified 16-GPU topology.""" + from megatron.bridge.recipes.qwen import qwen3_30b_a3b_pretrain_config + + mod = importlib.import_module("megatron.bridge.recipes.qwen.qwen3_moe") + patch_recipe_module_global(monkeypatch, mod, "AutoBridge", _FakeBridge) + + cfg = qwen3_30b_a3b_pretrain_config() + + _assert_basic_config(cfg) + assert cfg.model.tensor_model_parallel_size == 1 + assert cfg.model.pipeline_model_parallel_size == 1 + assert cfg.model.context_parallel_size == 1 + assert cfg.model.expert_model_parallel_size == 16 + assert cfg.model.expert_tensor_parallel_size == 1 + assert cfg.model.sequence_parallel is False + assert cfg.train.train_iters == 100 + assert cfg.train.global_batch_size == 1024 + assert cfg.train.micro_batch_size == 1 + assert cfg.dataset.random_seed == 1234 + assert cfg.rng.seed == 1234 + assert cfg.model.moe_flex_dispatcher_backend == "hybridep" + assert cfg.model.moe_token_dispatcher_type == "flex" + assert cfg.model.moe_shared_expert_overlap is False + assert cfg.model.moe_hybridep_num_sms == 32 + assert cfg.model.moe_router_force_load_balancing is False + assert cfg.model.recompute_granularity is None + assert cfg.model.recompute_method is None + assert cfg.model.recompute_num_layers is None + assert cfg.model.cuda_graph_impl == "transformer_engine" + assert cfg.model.cuda_graph_scope == ["moe_router", "moe_preprocess"] + assert cfg.model.use_te_rng_tracker is True + assert cfg.rng.te_rng_tracker is True + assert cfg.mixed_precision.grad_reduce_in_fp32 is False + assert cfg.ddp.grad_reduce_in_fp32 is False + assert cfg.optimizer.use_precision_aware_optimizer is True + assert cfg.optimizer.optimizer == "adam" + assert cfg.optimizer.lr == 3.0e-4 + assert cfg.optimizer.min_lr == 3.0e-5 + assert cfg.optimizer.adam_beta1 == 0.9 + assert cfg.optimizer.adam_beta2 == 0.95 + assert cfg.optimizer.adam_eps == 1.0e-8 + assert cfg.optimizer.weight_decay == 0.1 + assert cfg.optimizer.clip_grad == 1.0 + assert cfg.scheduler.start_weight_decay == 0.033 + assert cfg.scheduler.end_weight_decay == 0.033 + assert cfg.scheduler.weight_decay_incr_style == "constant" + assert cfg.scheduler.lr_decay_style == "cosine" + assert cfg.scheduler.lr_warmup_init == 0.0 + assert cfg.scheduler.lr_warmup_iters == 40 + assert cfg.scheduler.lr_decay_iters == 100 + assert cfg.optimizer.main_params_dtype == torch.float32 + assert cfg.optimizer.main_grads_dtype == torch.float32 + assert cfg.optimizer.exp_avg_dtype == torch.float32 + assert cfg.optimizer.exp_avg_sq_dtype == torch.float32 + assert cfg.optimizer.use_distributed_optimizer is True + assert cfg.mixed_precision.bf16 is True + assert cfg.tokenizer.hf_tokenizer_kwargs == { + "revision": "ad44e777bcd18fa416d9da3bd8f70d33ebb85d39" # pragma: allowlist secret + } + from megatron.bridge.training.utils.omegaconf_utils import process_config_with_overrides + + process_config_with_overrides( + cfg.tokenizer, + cli_overrides=[ + '++hf_tokenizer_kwargs.revision="ad44e777bcd18fa416d9da3bd8f70d33ebb85d39"' # pragma: allowlist secret + ], + ) + assert cfg.mixed_precision.params_dtype == torch.bfloat16 + assert cfg.checkpoint.save_interval == 50 + assert cfg.checkpoint.load is None + assert cfg.comm_overlap.tp_comm_overlap is True + + +def test_qwen3_30b_a3b_bf16_perf_recipe_uses_default_functional_config( + monkeypatch: pytest.MonkeyPatch, +): + """Test that the H100 BF16 perf recipe only adds benchmark-specific overrides.""" + from megatron.bridge.perf_recipes.qwen.h100.qwen3_moe import ( + qwen3_30b_a3b_pretrain_16gpu_h100_bf16_config, + ) + from megatron.bridge.recipes.qwen import qwen3_30b_a3b_pretrain_config + + mod = importlib.import_module("megatron.bridge.recipes.qwen.qwen3_moe") + patch_recipe_module_global(monkeypatch, mod, "AutoBridge", _FakeBridge) + + default_cfg = qwen3_30b_a3b_pretrain_config() + perf_cfg = qwen3_30b_a3b_pretrain_16gpu_h100_bf16_config() + + assert perf_cfg.model.tensor_model_parallel_size == default_cfg.model.tensor_model_parallel_size + assert perf_cfg.model.pipeline_model_parallel_size == default_cfg.model.pipeline_model_parallel_size + assert perf_cfg.model.expert_model_parallel_size == default_cfg.model.expert_model_parallel_size + assert perf_cfg.model.sequence_parallel == default_cfg.model.sequence_parallel + assert perf_cfg.train.global_batch_size == default_cfg.train.global_batch_size + assert perf_cfg.train.micro_batch_size == default_cfg.train.micro_batch_size + assert perf_cfg.model.moe_flex_dispatcher_backend == default_cfg.model.moe_flex_dispatcher_backend + assert perf_cfg.model.moe_token_dispatcher_type == default_cfg.model.moe_token_dispatcher_type + assert perf_cfg.model.cuda_graph_impl == default_cfg.model.cuda_graph_impl + assert perf_cfg.model.cuda_graph_scope == default_cfg.model.cuda_graph_scope + assert perf_cfg.comm_overlap.tp_comm_overlap == default_cfg.comm_overlap.tp_comm_overlap + assert perf_cfg.optimizer.use_precision_aware_optimizer == default_cfg.optimizer.use_precision_aware_optimizer + assert perf_cfg.model.moe_router_force_load_balancing is True + assert default_cfg.model.moe_router_force_load_balancing is False + + +def test_qwen3_30b_a3b_perf_base_remains_legacy_8gpu_recipe(): + """Keep non-H100-BF16 perf recipes isolated from the new generic default.""" + from megatron.bridge.perf_recipes.qwen.common import qwen3_30b_a3b_pretrain_config as perf_base + from megatron.bridge.recipes.qwen.h100.qwen3_moe import ( + qwen3_30b_a3b_pretrain_8gpu_h100_bf16_config as legacy_base, + ) + + assert perf_base is legacy_base + + +def test_qwen3_30b_a3b_h100_fp8_perf_recipe_keeps_cuda_graphs_disabled( + monkeypatch: pytest.MonkeyPatch, +): + """Test that changing the generic default does not alter the H100 FP8 recipe.""" + from megatron.bridge.perf_recipes.qwen.h100.qwen3_moe import ( + qwen3_30b_a3b_pretrain_16gpu_h100_fp8cs_config, + ) + + mod = importlib.import_module("megatron.bridge.recipes.qwen.qwen3_moe") + patch_recipe_module_global(monkeypatch, mod, "AutoBridge", _FakeBridge) + + cfg = qwen3_30b_a3b_pretrain_16gpu_h100_fp8cs_config() + + assert cfg.model.cuda_graph_impl == "none" + assert cfg.model.cuda_graph_scope == [] def test_qwen3_235b_a22b_lora_defaults(monkeypatch: pytest.MonkeyPatch): diff --git a/tests/unit_tests/recipes/test_recipe_environment_defaults.py b/tests/unit_tests/recipes/test_recipe_environment_defaults.py index c876cbcbf9..251103e393 100644 --- a/tests/unit_tests/recipes/test_recipe_environment_defaults.py +++ b/tests/unit_tests/recipes/test_recipe_environment_defaults.py @@ -125,7 +125,7 @@ def test_every_supported_hardware_recipe_declares_its_environment_inline(): ) assert isinstance(node.body[assignment_index + 1], ast.Return) - assert len(supported) == 247 + assert len(supported) == 259 assert unsupported == ["qwen/h100/qwen3_next.py:qwen3_next_80b_a3b_peft_1gpu_h100_bf16_config"] @@ -158,7 +158,7 @@ def test_explicit_recipe_environment_invariants(): assert environment["NVTE_FWD_LAYERNORM_SM_MARGIN"] == 20 assert environment["NVTE_BWD_LAYERNORM_SM_MARGIN"] == 20 - assert len(recipes) == 247 + assert len(recipes) == 259 assert hybrid_ep_count == 8 assert deepseek_v3_environment_recipe_names == _DEEPSEEK_V3_ENVIRONMENT_RECIPE_NAMES diff --git a/tests/unit_tests/scripts/training/test_setup_experiment.py b/tests/unit_tests/scripts/training/test_setup_experiment.py index 112e6aed79..74940e3154 100644 --- a/tests/unit_tests/scripts/training/test_setup_experiment.py +++ b/tests/unit_tests/scripts/training/test_setup_experiment.py @@ -271,8 +271,10 @@ def __init__(self, **kwargs): @pytest.mark.parametrize( ("extra_options", "expected_run", "expected_dryrun"), [ - ([], [{"detach": True}], 0), - (["--dry_run"], [{"detach": True}], 0), + ([], [{"detach": True, "tail_logs": False}], 0), + (["--wait"], [{"detach": False, "tail_logs": True}], 0), + (["--dry_run"], [{"detach": True, "tail_logs": False}], 0), + (["--wait", "--dry_run"], [{"detach": False, "tail_logs": True}], 0), (["--submission-dry-run"], [], 1), (["--dry-run"], [], 1), ], @@ -349,7 +351,7 @@ def dryrun(self): **module.TRAINING_LAUNCH_ENV, "PYTHONPATH": "/opt/Megatron-Bridge/src:/opt/Megatron-Bridge/3rdparty/Megatron-LM:$PYTHONPATH", } - submission_options = {"--submission-dry-run", "--dry-run"} + submission_options = {"--submission-dry-run", "--dry-run", "--wait"} expected_training_options = [option for option in extra_options if option not in submission_options] assert scripts[0].args == [*training_args, *expected_training_options] diff --git a/tests/unit_tests/test_compare_mask_handling.py b/tests/unit_tests/test_compare_mask_handling.py index a7f3edde1e..2134e33105 100644 --- a/tests/unit_tests/test_compare_mask_handling.py +++ b/tests/unit_tests/test_compare_mask_handling.py @@ -34,6 +34,9 @@ "megatron", "megatron.core", "megatron.core.parallel_state", + "megatron.core.inference", + "megatron.core.inference.contexts", + "megatron.core.inference.utils", "megatron.core.pipeline_parallel", "megatron.core.pipeline_parallel.schedules", "megatron.core.dist_checkpointing", @@ -51,6 +54,7 @@ "megatron.bridge.training.utils.checkpoint_utils", "megatron.bridge.utils", "megatron.bridge.utils.common_utils", + "megatron.bridge.utils.safe_url", "PIL", "PIL.Image", "requests", @@ -71,7 +75,11 @@ import compare # noqa: E402 from compare import ( # noqa: E402 SingleBatchIterator, + _broadcast_hf_results, + _maybe_gather_tensor_parallel_logits, _run_hf_inference, # noqa: E402 + _run_megatron_forward, + inference_forward_step, vlm_forward_step, ) @@ -108,6 +116,66 @@ def test_vlm_forward_step_passes_none_attention_mask(self): call_kwargs = mock_model.call_args.kwargs assert call_kwargs["attention_mask"] is None + assert "inference_context" not in call_kwargs + assert "runtime_gather_output" not in call_kwargs + + def test_text_inference_forward_step_passes_static_context(self): + """Test that the text path receives the cache context and gathered-logit request.""" + inference_context = object() + batch = { + "tokens": torch.tensor([[1, 2, 3]]), + "position_ids": torch.arange(3).unsqueeze(0), + "attention_mask": None, + "inference_context": inference_context, + } + mock_model = MagicMock(return_value=torch.randn(1, 1, 100)) + + inference_forward_step(iter([batch]), mock_model) + + assert mock_model.call_args.kwargs["inference_context"] is inference_context + assert mock_model.call_args.kwargs["runtime_gather_output"] is True + + def test_megatron_forward_activates_inference_mode(self): + """Test that the scheduled forward runs inside MCore inference mode.""" + mock_forward = MagicMock(return_value="output") + mock_context = MagicMock() + + with patch.object(compare.InferenceMode, "active", return_value=mock_context) as mock_active: + result = _run_megatron_forward(mock_forward, forward_only=True) + + assert result == "output" + mock_active.assert_called_once_with() + mock_context.__enter__.assert_called_once_with() + mock_context.__exit__.assert_called_once() + mock_forward.assert_called_once_with(forward_only=True) + + def test_tp_logits_skip_gather_when_runtime_output_is_already_full(self): + """Test that runtime-gathered text logits are not gathered a second time.""" + full_logits = torch.randn(1, 1, 128) + + with patch.object(compare.dist, "all_gather") as mock_all_gather: + result = _maybe_gather_tensor_parallel_logits(full_logits, 128, 2, object()) + + assert result is full_logits + mock_all_gather.assert_not_called() + + def test_tp_logits_gather_vlm_shards(self): + """Test that the existing sharded VLM path still gathers across TP ranks.""" + local_logits = torch.arange(64, dtype=torch.float32).reshape(1, 1, 64) + + def mock_all_gather(outputs, tensor, group): + assert group is tp_group + outputs[0].copy_(tensor) + outputs[1].copy_(tensor + 64) + + tp_group = object() + with patch.object(compare.dist, "all_gather", side_effect=mock_all_gather) as gather: + result = _maybe_gather_tensor_parallel_logits(local_logits, 128, 2, tp_group) + + assert result.shape == (1, 1, 128) + assert torch.equal(result[..., :64], local_logits) + assert torch.equal(result[..., 64:], local_logits + 64) + gather.assert_called_once() def test_hf_path_receives_ones_like_attention_mask(self): """Test that HF model receives torch.ones_like(input_ids, dtype=torch.bool) attention_mask.""" @@ -140,6 +208,26 @@ def test_hf_path_receives_ones_like_attention_mask(self): assert call_kwargs["attention_mask"].shape == input_ids.shape assert torch.equal(call_kwargs["attention_mask"], expected_mask) + def test_hf_broadcast_uses_model_output_vocab_size(self): + """Test that non-rank-0 buffers use the HF logits size instead of tokenizer vocab size.""" + broadcast_shapes = [] + + def mock_broadcast(tensor, _source_rank): + broadcast_shapes.append(tuple(tensor.shape)) + if len(broadcast_shapes) == 1: + tensor.fill_(163840) + + with ( + patch.object(torch.distributed, "broadcast", side_effect=mock_broadcast), + patch.object(torch.distributed, "barrier"), + ): + hf_logits, hf_next_token = _broadcast_hf_results(None, None, torch.device("cpu")) + + assert hf_logits.shape == (163840,) + assert hf_logits.dtype == torch.float32 + assert hf_next_token.shape == (1,) + assert broadcast_shapes == [(1,), (1,), (163840,)] + @pytest.mark.parametrize("flag", ["--trust_remote_code", "--trust-remote-code"]) def test_trust_remote_code_accepts_underscore_and_hyphen_flags(self, flag): """Test that compare.py accepts both trust_remote_code flag spellings.""" @@ -154,3 +242,21 @@ def test_trust_remote_code_accepts_underscore_and_hyphen_flags(self, flag): ) assert args.trust_remote_code is True + + def test_hf_revision_is_parsed_and_forwarded(self): + """Test that an immutable revision reaches every HF loader via shared kwargs.""" + revision = "0123456789abcdef0123456789abcdef01234567" # pragma: allowlist secret + args = compare.build_parser().parse_args( + [ + "--hf_model_path", + "org/model", + "--prompt", + "Hello", + "--hf-revision", + revision, + ] + ) + + assert args.hf_revision == revision + assert compare._hf_revision_kwargs(args.hf_revision) == {"revision": revision} + assert compare._hf_revision_kwargs(None) == {} diff --git a/tests/unit_tests/test_model_verification_card_validator.py b/tests/unit_tests/test_model_verification_card_validator.py new file mode 100644 index 0000000000..5218ed5f15 --- /dev/null +++ b/tests/unit_tests/test_model_verification_card_validator.py @@ -0,0 +1,270 @@ +import ast +import copy +import importlib.util +from pathlib import Path +from types import ModuleType +from typing import Any + +import pytest +import yaml + + +REPO_ROOT = Path(__file__).resolve().parents[2] +VALIDATOR_PATH = REPO_ROOT / "skills/create-model-verification-card/scripts/validate_card.py" +COMPARE_PATH = REPO_ROOT / "examples/conversion/compare_hf_and_megatron/compare.py" +CARD_PATH = REPO_ROOT / "model_cards/nemotron-3-nano-4b/card.yaml" +CORRELATION_CARD_PATH = REPO_ROOT / "model_cards/qwen3-8b/card.yaml" + + +def _load_validator() -> ModuleType: + spec = importlib.util.spec_from_file_location("model_card_validator", VALIDATOR_PATH) + assert spec is not None + assert spec.loader is not None + module = importlib.util.module_from_spec(spec) + spec.loader.exec_module(module) + return module + + +VALIDATOR = _load_validator() + + +def _card(path: Path = CARD_PATH) -> dict[str, Any]: + card = yaml.safe_load(path.read_text(encoding="utf-8")) + assert isinstance(card, dict) + return card + + +def _errors(card: dict[str, Any]) -> list[str]: + raw = yaml.safe_dump(card, sort_keys=False) + return VALIDATOR._validate_card(card, raw, ()) + + +def _assert_error(card: dict[str, Any], fragment: str) -> None: + errors = _errors(card) + assert any(fragment in error for error in errors), errors + + +def _assigned_float(path: Path, name: str) -> float: + tree = ast.parse(path.read_text(encoding="utf-8")) + for node in tree.body: + if isinstance(node, ast.Assign) and any( + isinstance(target, ast.Name) and target.id == name for target in node.targets + ): + value = ast.literal_eval(node.value) + assert isinstance(value, float) + return value + raise AssertionError(f"{name} is not assigned in {path}") + + +def test_repository_model_card_is_valid() -> None: + assert _errors(_card()) == [] + + +def test_manual_forward_requires_one_percent_correlation() -> None: + card = _card(CORRELATION_CARD_PATH) + manual = card["items"]["manual_forward_pass"] + expected = manual["expected_result"] + assert "0.999969" in expected + manual["expected_result"] = expected.replace("0.999969", "0.989999") + _assert_error(card, "cosine similarity must be at least 0.99") + + card = _card(CORRELATION_CARD_PATH) + manual = card["items"]["manual_forward_pass"] + manual["expected_result"] = manual["expected_result"].replace("0.999969", "0.990000") + assert _errors(card) == [] + + +def test_manual_forward_gate_matches_comparison_helper() -> None: + assert VALIDATOR.MANUAL_FORWARD_COSINE_THRESHOLD == 0.99 + assert _assigned_float(COMPARE_PATH, "SIMILARITY_THRESHOLD") == VALIDATOR.MANUAL_FORWARD_COSINE_THRESHOLD + + +def test_manual_forward_grandfathers_documented_historical_evidence() -> None: + card = _card(CORRELATION_CARD_PATH) + manual = card["items"]["manual_forward_pass"] + assert "--hf-revision" not in manual["command"] + assert "historical" in manual["expected_result"] + assert _errors(card) == [] + + manual["expected_result"] = ( + manual["expected_result"].replace("historical", "earlier").replace("predates", "preceded") + ) + _assert_error(card, "missing --hf-revision is allowed only for verified historical evidence") + + card = _card(CORRELATION_CARD_PATH) + card["items"]["manual_forward_pass"]["last_verified"] = "2026-07-20" + _assert_error(card, "evidence dated before 2026-07-20") + + +def test_unverified_manual_forward_command_requires_matching_revision() -> None: + card = _card(CORRELATION_CARD_PATH) + manual = card["items"]["manual_forward_pass"] + manual["status"] = "unverified" + manual["last_verified"] = None + _assert_error(card, "missing --hf-revision is allowed only for verified historical evidence") + + revision = card["model"]["hf_revision"] + manual["command"] += f" --hf-revision {revision}" + assert _errors(card) == [] + + manual["command"] = manual["command"].replace(revision, "1" * 40) + _assert_error(card, "--hf-revision must equal model.hf_revision") + + +def test_manual_forward_revision_must_match_card() -> None: + card = _card() + manual = card["items"]["manual_forward_pass"] + revision = card["model"]["hf_revision"] + assert f"--hf-revision {revision}" in manual["command"] + manual["command"] = manual["command"].replace(revision, "1" * 40) + _assert_error(card, "--hf-revision must equal model.hf_revision") + + +def test_manual_forward_revision_may_appear_only_once() -> None: + card = _card() + manual = card["items"]["manual_forward_pass"] + revision = card["model"]["hf_revision"] + manual["command"] += f" --hf-revision {revision}" + _assert_error(card, "specify --hf-revision at most once with a value") + + +def test_manual_forward_absolute_differences_are_report_only() -> None: + card = _card(CORRELATION_CARD_PATH) + manual = card["items"]["manual_forward_pass"] + expected = manual["expected_result"] + assert "0.187500" in expected + assert "0.030649" in expected + manual["expected_result"] = expected.replace("0.187500", "12.000000").replace("0.030649", "3.000000") + assert _errors(card) == [] + + +def test_manual_forward_requires_affirmative_next_token_match() -> None: + card = _card(CORRELATION_CARD_PATH) + manual = card["items"]["manual_forward_pass"] + manual["expected_result"] = manual["expected_result"].replace( + "next-token predictions match", + "next-token predictions do not match", + ) + _assert_error(card, "record that the next token matches") + + +@pytest.mark.parametrize("failure", ["failed to match", "never match"]) +def test_manual_forward_rejects_other_token_mismatch_wording(failure: str) -> None: + card = _card(CORRELATION_CARD_PATH) + manual = card["items"]["manual_forward_pass"] + manual["expected_result"] = manual["expected_result"].replace("predictions match", f"predictions {failure}") + _assert_error(card, "record that the next token matches") + + +def test_manual_forward_requires_numeric_absolute_differences() -> None: + card = _card(CORRELATION_CARD_PATH) + manual = card["items"]["manual_forward_pass"] + expected = manual["expected_result"] + manual["expected_result"] = expected.replace("0.187500", "not reported").replace("0.030649", "not reported") + _assert_error(card, "record the numeric max absolute logit difference") + _assert_error(card, "record the numeric mean absolute logit difference") + + +def test_manual_forward_accepts_helper_native_absolute_differences() -> None: + card = _card(CORRELATION_CARD_PATH) + manual = card["items"]["manual_forward_pass"] + expected = manual["expected_result"] + start = expected.index("The maximum and mean absolute logit differences") + manual["expected_result"] = expected[:start] + "Logits diff - max: 0.187500, mean: 0.030649." + assert _errors(card) == [] + + +def test_manual_forward_accepts_natural_absolute_differences() -> None: + card = _card(CORRELATION_CARD_PATH) + manual = card["items"]["manual_forward_pass"] + expected = manual["expected_result"] + start = expected.index("The maximum and mean absolute logit differences") + manual["expected_result"] = ( + expected[:start] + + "The maximum absolute logit difference was 0.187500 and the mean absolute logit difference was 0.030649." + ) + assert _errors(card) == [] + + +def test_manual_forward_rejects_negative_absolute_differences() -> None: + card = _card(CORRELATION_CARD_PATH) + manual = card["items"]["manual_forward_pass"] + expected = manual["expected_result"] + manual["expected_result"] = expected.replace("0.187500", "-0.187500").replace("0.030649", "-0.030649") + _assert_error(card, "record the numeric max absolute logit difference") + _assert_error(card, "record the numeric mean absolute logit difference") + + +def test_resume_must_load_the_reference_save_directory() -> None: + card = _card() + resume = card["items"]["checkpoint_resume"] + resume["command"] = resume["command"].replace( + "pretrain-reference-checkpoints", + "another-reference", + ) + _assert_error(card, "--load_dir must equal the pretrain --save_dir") + + +def test_resume_must_keep_reference_launch_settings() -> None: + card = _card() + resume = card["items"]["checkpoint_resume"] + resume["command"] = resume["command"].replace( + "nemotron_3_nano_4b_pretrain_8gpu_h100_bf16_config", + "another_recipe", + ) + _assert_error(card, "reference and resume launch settings must match") + + +def test_resume_cannot_disable_canonical_checkpoint_paths() -> None: + card = _card() + resume = card["items"]["checkpoint_resume"] + resume["command"] += " checkpoint.load=null checkpoint.save=null" + errors = _errors(card) + assert any("not checkpoint.load" in error for error in errors), errors + assert any("not checkpoint.save" in error for error in errors), errors + + +def test_reference_must_save_resume_state() -> None: + card = _card() + pretrain = card["items"]["pretrain"] + pretrain["command"] += " checkpoint.save_rng=false" + _assert_error(card, "checkpoint.save_rng must remain true") + + +def test_verified_resume_requires_matching_verified_commit() -> None: + card = _card() + pretrain = card["items"]["pretrain"] + pretrain["status"] = "unverified" + _assert_error(card, "pretrain must be verified first") + + card = _card() + resume = card["items"]["checkpoint_resume"] + resume["bridge_commit"] = "1" * 40 + _assert_error(card, "resume and pretrain must use the same Bridge commit") + + +def test_item_commit_override_must_be_verified_and_nonredundant() -> None: + card = _card() + manual = card["items"]["manual_forward_pass"] + manual["status"] = "unverified" + manual["bridge_commit"] = "1" * 40 + _assert_error(card, "item overrides are allowed only when verified") + + card = _card() + manual = card["items"]["manual_forward_pass"] + manual["bridge_commit"] = card["verification_environment"]["bridge_commit"] + _assert_error(card, "omit a redundant override") + + +def test_resume_setting_sort_handles_flag_and_flag_value() -> None: + settings = VALIDATOR._resume_reference_settings("./scripts/training/train.sh --deterministic --deterministic=true") + assert settings is not None + assert len(settings) == 2 + + +def test_missing_environment_commit_does_not_cascade_redundant_errors() -> None: + card = copy.deepcopy(_card()) + del card["verification_environment"]["bridge_commit"] + errors = _errors(card) + assert any("verification_environment/bridge_commit" in error for error in errors) + assert not any("omit a redundant override" in error for error in errors) diff --git a/tests/unit_tests/training/test_config.py b/tests/unit_tests/training/test_config.py index dfde228241..6c9354f6d4 100644 --- a/tests/unit_tests/training/test_config.py +++ b/tests/unit_tests/training/test_config.py @@ -3216,6 +3216,22 @@ def test_recipe_environment_variables_round_trip_through_yaml(self, tmp_path): finally: restore_get_world_size_safe(og_ws, cfg_mod) + def test_hf_model_revision_round_trip_through_yaml(self, tmp_path): + """Immutable Hugging Face model provenance should survive runtime config persistence.""" + revision = "b968826d9c46dd6066d109eabc6255188de91218" # pragma: allowlist secret + gpt_cfg = create_test_gpt_config(hf_model_id="Qwen/Qwen3-8B", hf_model_revision=revision) + full_cfg, og_ws, cfg_mod = create_test_config_container(world_size_override=1, model_config=gpt_cfg) + config_path = tmp_path / "recipe.yaml" + + try: + full_cfg.to_yaml(str(config_path)) + restored_cfg = ConfigContainer.from_yaml(str(config_path)) + + assert restored_cfg.model.hf_model_id == "Qwen/Qwen3-8B" + assert restored_cfg.model.hf_model_revision == revision + finally: + restore_get_world_size_safe(og_ws, cfg_mod) + def test_runtime_config_update_with_mixed_precision_string(self): """Test runtime_config_update with mixed precision as string.""" from megatron.bridge.training.config import runtime_config_update diff --git a/tests/unit_tests/training/test_tokenizer.py b/tests/unit_tests/training/test_tokenizer.py index dd374e3529..4db09c033f 100644 --- a/tests/unit_tests/training/test_tokenizer.py +++ b/tests/unit_tests/training/test_tokenizer.py @@ -12,7 +12,7 @@ # See the License for the specific language governing permissions and # limitations under the License. -from unittest.mock import patch +from unittest.mock import patch, sentinel import pytest from transformers import AutoTokenizer @@ -79,6 +79,79 @@ def test_build_hf_tokenizer(self, mock_hf_tokenizer_class, chat_template): assert tokenizer.library == "huggingface" assert tokenizer.chat_template == chat_template + @pytest.mark.parametrize( + ("model_id", "revision", "trust_remote_code"), + [ + ("Qwen/Qwen3-8B", "b968826d9c46dd6066d109eabc6255188de91218", False), # pragma: allowlist secret + ( + "moonshotai/Moonlight-16B-A3B", + "476b36a473d4467f94469414bef6cee75c9c8172", # pragma: allowlist secret + True, + ), + ], + ) + @patch("megatron.bridge.training.tokenizers.tokenizer.build_mcore_tokenizer") + @patch("huggingface_hub.snapshot_download") + def test_build_hf_tokenizer_resolves_immutable_revision( + self, + mock_snapshot_download, + mock_build_mcore_tokenizer, + model_id, + revision, + trust_remote_code, + ): + mock_snapshot_download.return_value = f"/cache/models--resolved/snapshots/{revision}" + mock_build_mcore_tokenizer.return_value = sentinel.tokenizer + hf_tokenizer_kwargs = {"revision": revision, "trust_remote_code": trust_remote_code} + config = TokenizerConfig( + tokenizer_type="HuggingFaceTokenizer", + tokenizer_model=model_id, + hf_tokenizer_kwargs=hf_tokenizer_kwargs, + ) + + tokenizer = build_tokenizer(config) + + assert tokenizer is sentinel.tokenizer + mock_snapshot_download.assert_called_once() + snapshot_kwargs = mock_snapshot_download.call_args.kwargs + assert snapshot_kwargs["repo_id"] == model_id + assert snapshot_kwargs["revision"] == revision + assert {"config.json", "tokenizer*", "*.py", "*.jinja"}.issubset(snapshot_kwargs["allow_patterns"]) + assert {"*.safetensors", "*.bin", "*.pt", "*.gguf"}.issubset(snapshot_kwargs["ignore_patterns"]) + + resolved_config = mock_build_mcore_tokenizer.call_args.args[0] + assert resolved_config is not config + assert resolved_config.tokenizer_model == mock_snapshot_download.return_value + assert resolved_config.trust_remote_code is trust_remote_code + assert config.tokenizer_model == model_id + assert config.hf_tokenizer_kwargs == hf_tokenizer_kwargs + + @patch("megatron.bridge.training.tokenizers.tokenizer.build_mcore_tokenizer") + @patch("huggingface_hub.snapshot_download") + def test_build_tokenizer_skips_snapshot_resolution_without_remote_hf_revision( + self, mock_snapshot_download, mock_build_mcore_tokenizer, tmp_path + ): + mock_build_mcore_tokenizer.return_value = sentinel.tokenizer + configs = [ + TokenizerConfig(tokenizer_type="HuggingFaceTokenizer", tokenizer_model="Qwen/Qwen3-8B"), + TokenizerConfig( + tokenizer_type="NullTokenizer", + vocab_size=32, + hf_tokenizer_kwargs={"revision": "not-used"}, + ), + TokenizerConfig( + tokenizer_type="HuggingFaceTokenizer", + tokenizer_model=tmp_path, + hf_tokenizer_kwargs={"revision": "local-path-is-already-resolved"}, + ), + ] + + for config in configs: + assert build_tokenizer(config) is sentinel.tokenizer + assert mock_build_mcore_tokenizer.call_args.args[0] is config + + mock_snapshot_download.assert_not_called() + @patch("megatron.core.tokenizers.text.libraries.SentencePieceTokenizer") @pytest.mark.parametrize("legacy", [True]) def test_build_sp_tokenizer(self, mock_sp_tokenizer, legacy):