diff --git a/docs/models/qwen/qwen35-vl.md b/docs/models/qwen/qwen35-vl.md index 7b3cee0dea..999640b9f3 100644 --- a/docs/models/qwen/qwen35-vl.md +++ b/docs/models/qwen/qwen35-vl.md @@ -52,6 +52,29 @@ Please upgrade to `transformers` >= 5.2.0 in order to use the Qwen 3.5 models. For checkpoint conversion, inference, finetuning recipes, and step-by-step training guides, see the [Qwen 3.5 Examples](https://github.com/NVIDIA-NeMo/Megatron-Bridge/blob/main/examples/models/qwen/qwen35_vl/README.md). +### Text-only pretraining + +The Qwen3.5 9B and 35B-A3B recipes can pretrain only the language-model +component of the unified models. They derive the registered language-model +provider from the nested Hugging Face `text_config`; no vision model, +projection, processor, or multimodal dataset is created. + +```python +from megatron.bridge.recipes.qwen import qwen35_text_9b_pretrain_config + +config = qwen35_text_9b_pretrain_config() +``` + +The canonical aliases select eight-GPU GB200 BF16 library recipes. The dense +9B recipe, `qwen35_text_9b_pretrain_8gpu_gb200_bf16_config`, uses the same data +parallel topology as the Llama 3 8B GB200 performance recipe, with +module-scoped CUDA graphs so library correctness checks remain enabled. The MoE recipe, +`qwen35_text_35b_a3b_pretrain_8gpu_gb200_bf16_config`, uses the applicable Qwen3.5-VL +GB200 HybridEP settings with learned routing. Both retain library-recipe +training, evaluation, logging, checkpointing, and correctness defaults. Set +`config.dataset.blend` (or `config.dataset.data_path`) to use a prepared +Megatron indexed text dataset. + ## Hugging Face Model Cards - Qwen3.5 0.8B: https://huggingface.co/Qwen/Qwen3.5-0.8B diff --git a/src/megatron/bridge/recipes/nemotronh/__init__.py b/src/megatron/bridge/recipes/nemotronh/__init__.py index 4a3ab0a63c..ed1375a11a 100644 --- a/src/megatron/bridge/recipes/nemotronh/__init__.py +++ b/src/megatron/bridge/recipes/nemotronh/__init__.py @@ -14,6 +14,10 @@ # Nemotron Nano v2 models # Nemotron 3 Nano models +from megatron.bridge.recipes.nemotronh.gb200 import ( + nemotron_3_nano_gb200_pretrain_config, + nemotron_3_nano_pretrain_8gpu_gb200_bf16_config, +) from megatron.bridge.recipes.nemotronh.nemotron_3_nano import ( nemotron_3_nano_peft_config, nemotron_3_nano_pretrain_config, @@ -88,6 +92,8 @@ "nemotron_3_nano_pretrain_config", "nemotron_3_nano_sft_config", "nemotron_3_nano_peft_config", + "nemotron_3_nano_gb200_pretrain_config", + "nemotron_3_nano_pretrain_8gpu_gb200_bf16_config", # Nemotron 3 Nano 4B model "nemotron_3_nano_4b_pretrain_config", "nemotron_3_nano_4b_sft_config", diff --git a/src/megatron/bridge/recipes/nemotronh/gb200/__init__.py b/src/megatron/bridge/recipes/nemotronh/gb200/__init__.py new file mode 100644 index 0000000000..c4b08f394a --- /dev/null +++ b/src/megatron/bridge/recipes/nemotronh/gb200/__init__.py @@ -0,0 +1,24 @@ +# 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. + +from megatron.bridge.recipes.nemotronh.gb200.nemotron_3_nano import ( + nemotron_3_nano_gb200_pretrain_config, + nemotron_3_nano_pretrain_8gpu_gb200_bf16_config, +) + + +__all__ = [ + "nemotron_3_nano_gb200_pretrain_config", + "nemotron_3_nano_pretrain_8gpu_gb200_bf16_config", +] diff --git a/src/megatron/bridge/recipes/nemotronh/gb200/nemotron_3_nano.py b/src/megatron/bridge/recipes/nemotronh/gb200/nemotron_3_nano.py new file mode 100644 index 0000000000..fcb5070d76 --- /dev/null +++ b/src/megatron/bridge/recipes/nemotronh/gb200/nemotron_3_nano.py @@ -0,0 +1,148 @@ +# 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. + +"""GB200 pretraining recipe for Nemotron 3 Nano.""" + +import torch + +from megatron.bridge import AutoBridge +from megatron.bridge.recipes.common import _pretrain_common +from megatron.bridge.training.comm_overlap import CommOverlapConfig +from megatron.bridge.training.config import ConfigContainer +from megatron.bridge.training.mixed_precision import get_mixed_precision_config +from megatron.bridge.utils.cuda_graph import set_cuda_graph_modules + + +_NEMOTRON_3_NANO_MODEL_ID = "nvidia/NVIDIA-Nemotron-3-Nano-30B-A3B-BF16" + + +def nemotron_3_nano_pretrain_8gpu_gb200_bf16_config() -> ConfigContainer: + """Return the Nemotron 3 Nano BF16 pretraining config for eight GB200 GPUs. + + The recipe retains the established optimizer, scheduler, routing, and BF16 + contracts. It applies the validated GB200 TP1/EP8 HybridEP topology and + uses a 4,096-token sequence length for the paired NeMo-CI convergence + workload. + + Returns: + GB200 BF16 pretraining configuration. + """ + cfg = _pretrain_common() + + cfg.model = AutoBridge.from_hf_pretrained(_NEMOTRON_3_NANO_MODEL_ID).to_megatron_provider(load_weights=False) + # Pretraining may use a tokenizer other than the HF checkpoint tokenizer. + # Defer the model vocabulary size to the runtime tokenizer, matching the + # pre-migration MambaModelProvider recipe behavior. + cfg.model.vocab_size = None + cfg.tokenizer.tokenizer_model = _NEMOTRON_3_NANO_MODEL_ID + + cfg.model.seq_length = 4096 + cfg.dataset.seq_length = 4096 + cfg.dataset.blend = None + cfg.dataset.num_workers = 8 + cfg.dataset.mmap_bin_files = False + + cfg.model.tensor_model_parallel_size = 1 + cfg.model.pipeline_model_parallel_size = 1 + cfg.model.pipeline_model_parallel_layout = None + cfg.model.pipeline_dtype = torch.bfloat16 + cfg.model.virtual_pipeline_model_parallel_size = None + cfg.model.context_parallel_size = 1 + cfg.model.sequence_parallel = False + cfg.model.expert_tensor_parallel_size = 1 + cfg.model.expert_model_parallel_size = 8 + + cfg.model.moe_token_dispatcher_type = "flex" + cfg.model.moe_flex_dispatcher_backend = "hybridep" + cfg.model.moe_flex_dispatcher_num_sms = 16 + cfg.model.moe_hybridep_num_sms = None + cfg.model.moe_shared_expert_overlap = False + cfg.model.moe_router_force_load_balancing = False + + cfg.train.train_iters = 39735 + cfg.train.global_batch_size = 3072 + cfg.train.micro_batch_size = 2 + cfg.train.manual_gc = False + cfg.train.manual_gc_interval = 0 + + cfg.model.transformer_impl = "transformer_engine" + + # Match the validated GB200 performance recipe's TE-scoped graph set. + cfg.model.cuda_graph_impl = "transformer_engine" + set_cuda_graph_modules(cfg.model, ["attn", "mamba", "moe_router", "moe_preprocess"]) + cfg.model.cuda_graph_warmup_steps = 3 + cfg.model.use_te_rng_tracker = True + cfg.rng.te_rng_tracker = True + + cfg.model.attention_backend = "fused" + cfg.model.moe_router_fusion = False + cfg.model.moe_permute_fusion = True + cfg.model.moe_grouped_gemm = True + cfg.model.cross_entropy_loss_fusion = True + cfg.model.apply_rope_fusion = True + cfg.model.cross_entropy_fusion_impl = "native" + cfg.model.recompute_granularity = None + cfg.model.recompute_modules = None + cfg.model.fine_grained_activation_offloading = False + cfg.model.offload_modules = None + cfg.model.moe_router_padding_for_fp8 = False + cfg.rerun_state_machine.check_for_nan_in_loss = False + + 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.lr = 1.6e-3 + cfg.optimizer.weight_decay = 0.1 + cfg.optimizer.min_lr = 1.6e-5 + cfg.scheduler.lr_warmup_iters = 333 + + # Keep BF16 compute while reducing gradients in BF16 instead of FP32. + cfg.mixed_precision = get_mixed_precision_config(cfg.mixed_precision) + cfg.mixed_precision.grad_reduce_in_fp32 = False + + cfg.comm_overlap = CommOverlapConfig( + tp_comm_bootstrap_backend="nccl", + tp_comm_overlap=False, + ) + cfg.comm_overlap.delay_wgrad_compute = False + cfg.comm_overlap.overlap_moe_expert_parallel_comm = False + + cfg.checkpoint.save_interval = 200 + cfg.checkpoint.ckpt_assume_constant_structure = True + cfg.checkpoint.dist_ckpt_strictness = "log_all" + + cfg.ddp.overlap_grad_reduce = True + cfg.ddp.overlap_param_gather = True + cfg.ddp.check_for_nan_in_grad = False + cfg.ddp.use_distributed_optimizer = True + cfg.ddp.grad_reduce_in_fp32 = False + + cfg.model.init_method_std = 0.0173 + cfg.model.use_fused_weighted_squared_relu = True + + return cfg + + +# NeMo-CI appends ``_pretrain_config`` to MODEL_RECIPE_NAME. This explicit +# alias lets the GB200 release case select the hardware recipe without changing +# the legacy ``nemotron_3_nano_pretrain_config`` default. +nemotron_3_nano_gb200_pretrain_config = nemotron_3_nano_pretrain_8gpu_gb200_bf16_config + + +__all__ = [ + "nemotron_3_nano_gb200_pretrain_config", + "nemotron_3_nano_pretrain_8gpu_gb200_bf16_config", +] diff --git a/src/megatron/bridge/recipes/qwen/__init__.py b/src/megatron/bridge/recipes/qwen/__init__.py index b228b95ab4..93588f4e14 100644 --- a/src/megatron/bridge/recipes/qwen/__init__.py +++ b/src/megatron/bridge/recipes/qwen/__init__.py @@ -12,6 +12,12 @@ # See the License for the specific language governing permissions and # limitations under the License. +# Qwen3.5 GB200 models +from .gb200.qwen35 import ( + qwen35_text_9b_pretrain_8gpu_gb200_bf16_config, + qwen35_text_35b_a3b_pretrain_8gpu_gb200_bf16_config, +) + # Qwen2 models from .qwen2 import ( qwen2_1p5b_peft_config, @@ -88,6 +94,9 @@ qwen3_next_80b_a3b_sft_config, ) +# Qwen3.5 text models +from .qwen35 import qwen35_text_9b_pretrain_config, qwen35_text_35b_a3b_pretrain_config + __all__ = [ # Qwen2 models @@ -155,4 +164,9 @@ "qwen3_next_80b_a3b_pretrain_config", "qwen3_next_80b_a3b_sft_config", "qwen3_next_80b_a3b_peft_config", + # Qwen3.5 text models + "qwen35_text_9b_pretrain_config", + "qwen35_text_9b_pretrain_8gpu_gb200_bf16_config", + "qwen35_text_35b_a3b_pretrain_config", + "qwen35_text_35b_a3b_pretrain_8gpu_gb200_bf16_config", ] diff --git a/src/megatron/bridge/recipes/qwen/gb200/__init__.py b/src/megatron/bridge/recipes/qwen/gb200/__init__.py new file mode 100644 index 0000000000..2d5eeb765f --- /dev/null +++ b/src/megatron/bridge/recipes/qwen/gb200/__init__.py @@ -0,0 +1,24 @@ +# 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. + +from megatron.bridge.recipes.qwen.gb200.qwen35 import ( + qwen35_text_9b_pretrain_8gpu_gb200_bf16_config, + qwen35_text_35b_a3b_pretrain_8gpu_gb200_bf16_config, +) + + +__all__ = [ + "qwen35_text_9b_pretrain_8gpu_gb200_bf16_config", + "qwen35_text_35b_a3b_pretrain_8gpu_gb200_bf16_config", +] diff --git a/src/megatron/bridge/recipes/qwen/gb200/qwen35.py b/src/megatron/bridge/recipes/qwen/gb200/qwen35.py new file mode 100644 index 0000000000..f151a90b01 --- /dev/null +++ b/src/megatron/bridge/recipes/qwen/gb200/qwen35.py @@ -0,0 +1,191 @@ +# 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. + +"""GB200 text-only pretraining recipes for Qwen3.5 dense and MoE models.""" + +from __future__ import annotations + +import torch +from transformers import AutoConfig + +from megatron.bridge import AutoBridge +from megatron.bridge.recipes.common import _pretrain_common +from megatron.bridge.training.comm_overlap import CommOverlapConfig +from megatron.bridge.training.config import ConfigContainer +from megatron.bridge.training.mixed_precision import bf16_mixed + + +_QWEN35_9B_BASE = "Qwen/Qwen3.5-9B-Base" +_QWEN35_35B_A3B_BASE = "Qwen/Qwen3.5-35B-A3B-Base" + + +def qwen35_text_9b_pretrain_8gpu_gb200_bf16_config() -> ConfigContainer: + """Return a text-only Qwen3.5-9B pretraining config for eight GB200 GPUs.""" + cfg = _pretrain_common() + + text_config = AutoConfig.from_pretrained(_QWEN35_9B_BASE).text_config + # The nested text config intentionally omits ``architectures``. AutoBridge + # needs it to select the registered causal-LM bridge instead of the VLM. + text_config.architectures = ["Qwen3_5ForCausalLM"] + cfg.model = AutoBridge.from_hf_config(text_config).to_megatron_provider(load_weights=False) + cfg.tokenizer.tokenizer_model = _QWEN35_9B_BASE + cfg.dataset.seq_length = 4096 + cfg.dataset.blend = None + cfg.dataset.num_workers = 8 + + # Follow the Llama 3 8B GB200 topology: keep model parallelism at one and + # use all eight GPUs for data parallelism. + cfg.model.tensor_model_parallel_size = 1 + cfg.model.pipeline_model_parallel_size = 1 + cfg.model.pipeline_model_parallel_layout = None + cfg.model.pipeline_dtype = torch.bfloat16 + cfg.model.virtual_pipeline_model_parallel_size = None + cfg.model.context_parallel_size = 1 + cfg.model.expert_model_parallel_size = 1 + cfg.model.expert_tensor_parallel_size = 1 + cfg.model.sequence_parallel = False + cfg.model.seq_length = 4096 + cfg.model.init_method_std = 0.02 + cfg.train.global_batch_size = 128 + cfg.train.micro_batch_size = 2 + + cfg.model.transformer_impl = "transformer_engine" + cfg.model.bias_activation_fusion = True + cfg.model.cross_entropy_loss_fusion = True + cfg.model.cross_entropy_fusion_impl = "native" + cfg.model.apply_rope_fusion = True + + cfg.model.recompute_granularity = None + cfg.model.recompute_method = None + cfg.model.recompute_num_layers = None + cfg.model.recompute_modules = None + cfg.model.fine_grained_activation_offloading = False + cfg.model.offload_modules = None + + # Capture the dense attention and MLP modules. Keep cross entropy on the + # native fused path validated by the 64-GPU GB200 performance run. + cfg.model.cuda_graph_impl = "transformer_engine" + cfg.model.cuda_graph_scope = None + cfg.model.cuda_graph_modules = ["attn", "mlp"] + cfg.model.cuda_graph_warmup_steps = 3 + cfg.model.use_te_rng_tracker = True + cfg.rng.te_rng_tracker = True + + 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 = False + + cfg.ddp.overlap_grad_reduce = True + cfg.ddp.overlap_param_gather = True + cfg.ddp.grad_reduce_in_fp32 = False + cfg.ddp.check_for_nan_in_grad = False + cfg.ddp.use_distributed_optimizer = True + cfg.ddp.use_megatron_fsdp = False + cfg.rerun_state_machine.check_for_nan_in_loss = False + + cfg.comm_overlap = CommOverlapConfig(tp_comm_overlap=False) + return cfg + + +def qwen35_text_35b_a3b_pretrain_8gpu_gb200_bf16_config() -> ConfigContainer: + """Return a text-only Qwen3.5-35B-A3B pretraining config for eight GB200 GPUs.""" + cfg = _pretrain_common() + + text_config = AutoConfig.from_pretrained(_QWEN35_35B_A3B_BASE).text_config + # The nested text config intentionally omits ``architectures``. AutoBridge + # needs it to select the registered causal-LM bridge instead of the VLM. + text_config.architectures = ["Qwen3_5MoeForCausalLM"] + cfg.model = AutoBridge.from_hf_config(text_config).to_megatron_provider(load_weights=False) + cfg.tokenizer.tokenizer_model = _QWEN35_35B_A3B_BASE + cfg.dataset.seq_length = 4096 + cfg.dataset.blend = None + cfg.dataset.num_workers = 8 + + # Match the Qwen3.5-VL GB200 topology while training only the text model. + cfg.model.tensor_model_parallel_size = 1 + cfg.model.pipeline_model_parallel_size = 1 + cfg.model.pipeline_model_parallel_layout = None + cfg.model.pipeline_dtype = torch.bfloat16 + 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.seq_length = 4096 + cfg.model.init_method_std = 0.02 + cfg.train.global_batch_size = 512 + # MBS4 is suitable for force-balanced throughput benchmarking, but OOMs + # with learned routing. MBS1 was validated with real RP2 data on GB200. + cfg.train.micro_batch_size = 1 + + cfg.model.transformer_impl = "transformer_engine" + cfg.model.bias_activation_fusion = True + cfg.model.moe_router_fusion = True + cfg.model.moe_permute_fusion = True + cfg.model.moe_grouped_gemm = True + cfg.model.cross_entropy_loss_fusion = True + # Keep the library-safe native implementation instead of the performance + # harness's TE cross-entropy path, which currently warns about stability. + cfg.model.cross_entropy_fusion_impl = "native" + cfg.model.apply_rope_fusion = True + + cfg.model.recompute_granularity = None + cfg.model.recompute_method = None + cfg.model.recompute_num_layers = None + cfg.model.recompute_modules = None + cfg.model.fine_grained_activation_offloading = False + cfg.model.offload_modules = None + + # Fixed-length text batches can use the scopes that the VLM recipe must + # disable for variable-length multimodal inputs. + cfg.model.cuda_graph_impl = "transformer_engine" + cfg.model.cuda_graph_scope = None + cfg.model.cuda_graph_modules = ["attn", "moe_router", "moe_preprocess"] + cfg.model.cuda_graph_warmup_steps = 3 + cfg.model.use_te_rng_tracker = True + cfg.rng.te_rng_tracker = True + + cfg.model.moe_token_dispatcher_type = "flex" + cfg.model.moe_flex_dispatcher_backend = "hybridep" + cfg.model.moe_flex_dispatcher_num_sms = 32 + cfg.model.moe_hybridep_num_sms = None + cfg.model.moe_router_dtype = "fp32" + cfg.model.moe_shared_expert_overlap = False + cfg.model.moe_router_force_load_balancing = False + cfg.model.moe_router_padding_for_fp8 = False + + 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.overlap_param_gather_with_optimizer_step = False + + 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 = True + cfg.ddp.use_megatron_fsdp = False + + cfg.comm_overlap = CommOverlapConfig( + tp_comm_overlap=True, + overlap_grad_reduce=False, + overlap_param_gather=False, + ) + return cfg diff --git a/src/megatron/bridge/recipes/qwen/qwen35.py b/src/megatron/bridge/recipes/qwen/qwen35.py new file mode 100644 index 0000000000..51025d0074 --- /dev/null +++ b/src/megatron/bridge/recipes/qwen/qwen35.py @@ -0,0 +1,27 @@ +# 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 alias for the canonical Qwen3.5 text recipe.""" + +from __future__ import annotations + +from megatron.bridge.recipes.qwen.gb200.qwen35 import ( + qwen35_text_9b_pretrain_8gpu_gb200_bf16_config as qwen35_text_9b_pretrain_config, +) +from megatron.bridge.recipes.qwen.gb200.qwen35 import ( + qwen35_text_35b_a3b_pretrain_8gpu_gb200_bf16_config as qwen35_text_35b_a3b_pretrain_config, +) + + +__all__ = ["qwen35_text_9b_pretrain_config", "qwen35_text_35b_a3b_pretrain_config"] diff --git a/tests/unit_tests/recipes/recipe_test_utils.py b/tests/unit_tests/recipes/recipe_test_utils.py index b8c483c241..334062bba9 100644 --- a/tests/unit_tests/recipes/recipe_test_utils.py +++ b/tests/unit_tests/recipes/recipe_test_utils.py @@ -152,6 +152,11 @@ def finalize(self) -> None: class _OfflineAutoBridge: """Build a local provider without reading a Hugging Face configuration.""" + @classmethod + def from_hf_config(cls, *args: object, **kwargs: object) -> "_OfflineAutoBridge": + del args, kwargs + return cls() + @classmethod def from_hf_pretrained(cls, *args: object, **kwargs: object) -> "_OfflineAutoBridge": del args, kwargs @@ -215,7 +220,17 @@ def skip_flex_dispatcher_hardware_probe(*args: object, **kwargs: object) -> None deepseek_v4_recipe_module = importlib.import_module("megatron.bridge.recipes.deepseek.h100.deepseek_v4") monkeypatch.setattr(deepseek_v4_recipe_module, "deepseek_v4_supports_blackwell_fused_kernels", lambda: False) - from transformers import AutoTokenizer + from transformers import AutoConfig, AutoTokenizer + + def load_offline_auto_config(*args: object, **kwargs: object) -> SimpleNamespace: + del args, kwargs + return SimpleNamespace(text_config=SimpleNamespace(architectures=None)) + + monkeypatch.setattr( + AutoConfig, + "from_pretrained", + staticmethod(load_offline_auto_config), + ) def load_offline_tokenizer(*args: object, **kwargs: object) -> _OfflineTokenizer: del args, kwargs diff --git a/tests/unit_tests/recipes/test_nemotronh_recipes.py b/tests/unit_tests/recipes/test_nemotronh_recipes.py index 98177dbdc9..c60a18cfd2 100644 --- a/tests/unit_tests/recipes/test_nemotronh_recipes.py +++ b/tests/unit_tests/recipes/test_nemotronh_recipes.py @@ -59,8 +59,9 @@ def to_megatron_provider(self, *args, **kwargs): @pytest.fixture(autouse=True) def _patch_hf_backed_recipe_providers(monkeypatch: pytest.MonkeyPatch) -> None: - """Keep Super and Ultra recipe construction deterministic and offline.""" + """Keep AutoBridge-backed recipe construction deterministic and offline.""" for module_name in ( + "megatron.bridge.recipes.nemotronh.gb200.nemotron_3_nano", "megatron.bridge.recipes.nemotronh.nemotron_3_super", "megatron.bridge.recipes.nemotronh.nemotron_3_ultra", ): @@ -138,6 +139,13 @@ def test_nemotronh_recipe_rejects_unknown_cli_override(): assert not hasattr(cfg.model, "not_a_real_field") +def test_nemotron_3_nano_gb200_defers_vocab_size_to_training_tokenizer(): + """The GB200 pretraining model vocabulary must follow its runtime tokenizer.""" + cfg = _nemotronh_module.nemotron_3_nano_pretrain_8gpu_gb200_bf16_config() + + assert cfg.model.vocab_size is None + + def test_nemotron_nano_9b_v2_lora_defaults(): """Test that Nemotron Nano 9B v2 LoRA has correct default parallelism.""" from megatron.bridge.recipes.nemotronh import nemotron_nano_9b_v2_peft_config diff --git a/tests/unit_tests/recipes/test_qwen_recipes.py b/tests/unit_tests/recipes/test_qwen_recipes.py index 6bcb9a39cd..5f965b65c7 100644 --- a/tests/unit_tests/recipes/test_qwen_recipes.py +++ b/tests/unit_tests/recipes/test_qwen_recipes.py @@ -77,6 +77,26 @@ def from_hf_pretrained(hf_path: str, **kwargs): assert kwargs == {"revision": expected_revisions[hf_path]} return _FakeBridge() + @staticmethod + def from_hf_config(hf_config): + # Ignore hf_config; return a bridge that yields a fake provider + return _FakeBridge() + + +class _FakeTextConfig: + architectures = None + + +class _FakeRootConfig: + text_config = _FakeTextConfig() + + +class _FakeAutoConfig: + @staticmethod + def from_pretrained(hf_path: str): + # Ignore hf_path; return a unified config with a nested text config. + return _FakeRootConfig() + def _assert_basic_config(cfg): from megatron.bridge.training.config import ConfigContainer @@ -114,6 +134,8 @@ def test_each_qwen_recipe_builds_config(recipe_func: Callable, monkeypatch: pyte mod = importlib.import_module(module_name) if hasattr(mod, "AutoBridge"): patch_recipe_module_global(monkeypatch, mod, "AutoBridge", _FakeBridge) + if hasattr(mod, "AutoConfig"): + patch_recipe_module_global(monkeypatch, mod, "AutoConfig", _FakeAutoConfig) overrides = _safe_overrides_for(recipe_func.__name__) @@ -150,7 +172,12 @@ def test_each_qwen_recipe_builds_config(recipe_func: Callable, monkeypatch: pyte assert getattr(cfg.model, "tensor_model_parallel_size", 1) >= 1 assert getattr(cfg.model, "pipeline_model_parallel_size", 1) >= 1 - if "qwen3" in recipe_name and "pretrain" in recipe_name and "next" not in recipe_name: + if ( + "qwen3" in recipe_name + and "pretrain" in recipe_name + and "next" not in recipe_name + and "qwen35" not in recipe_name + ): assert cfg.model.cross_entropy_fusion_impl == "te" # SFT and PEFT-specific assertions