diff --git a/docs/models/llm/qwen.md b/docs/models/llm/qwen.md index 852672243b9..cd291056d4d 100644 --- a/docs/models/llm/qwen.md +++ b/docs/models/llm/qwen.md @@ -1,6 +1,7 @@ # Qwen -Qwen2/2.5/3 models are supported via the Bridge with QK layernorm handling (Qwen3) and bias in QKV (Qwen2). +Qwen2/2.5/3/3-Next models are supported via the Bridge with QK layernorm handling (Qwen3) and bias in QKV (Qwen2). +Qwen3-Next models also support Gated-Delta Networks, Zero-Centered RMSNorm and QK Layernorm weight decay during training. ## Conversion with 🤗 Hugging Face @@ -47,6 +48,7 @@ cfg = qwen3_8b_pretrain_config( - Qwen2 recipes: [bridge.recipes.qwen.qwen2](../../apidocs/bridge/bridge.recipes.qwen.qwen2.md) - Qwen3 recipes: [bridge.recipes.qwen.qwen3](../../apidocs/bridge/bridge.recipes.qwen.qwen3.md) - Qwen3 MoE recipes: [bridge.recipes.qwen.qwen3_moe](../../apidocs/bridge/bridge.recipes.qwen.qwen3_moe.md) + - Qwen3-Next recipes: [bridge.recipes.qwen.qwen3_next](../../apidocs/bridge/bridge.recipes.qwen.qwen3_next.md) ## Finetuning recipes - Coming soon @@ -55,6 +57,7 @@ cfg = qwen3_8b_pretrain_config( - Qwen2: `https://huggingface.co/Qwen/Qwen2-7B` - Qwen2.5: `https://huggingface.co/Qwen/Qwen2.5-7B` - Qwen3: `https://huggingface.co/Qwen/Qwen3-7B` +- Qwen3-Next `https://huggingface.co/Qwen/Qwen3-Next-80B-A3B-Instruct` and `https://huggingface.co/Qwen/Qwen3-Next-80B-A3B-Thinking` ## Related docs - Recipe usage and customization: [Recipe usage](../../recipe-usage.md) diff --git a/src/megatron/bridge/models/qwen/qwen_provider.py b/src/megatron/bridge/models/qwen/qwen_provider.py index 68a05a74995..cb581423d45 100644 --- a/src/megatron/bridge/models/qwen/qwen_provider.py +++ b/src/megatron/bridge/models/qwen/qwen_provider.py @@ -15,15 +15,23 @@ import logging from dataclasses import dataclass from functools import partial -from typing import TYPE_CHECKING, Callable, List, Optional, Union +from typing import TYPE_CHECKING, Callable, Optional import torch import torch.nn.functional as F from megatron.core.models.gpt.gpt_layer_specs import get_gpt_decoder_block_spec +from megatron.core.transformer.spec_utils import ModuleSpec from megatron.bridge.models.gpt_provider import GPTModelProvider +try: + import transformer_engine # type: ignore # noqa: F401 + + HAVE_TE = True +except (ImportError, ModuleNotFoundError): + HAVE_TE = False + try: import transformer_engine # type: ignore # noqa: F401 @@ -425,34 +433,34 @@ class Qwen3MoEModelProvider235B_A22B(Qwen3MoEModelProvider): class Qwen3NextModelProvider(Qwen3MoEModelProvider): """Base provider for Qwen 3 Next Models.""" - transformer_layer_spec: Union["ModuleSpec", Callable[["GPTModelProvider"], "ModuleSpec"]] = partial( + transformer_layer_spec: ModuleSpec | Callable[["GPTModelProvider"], ModuleSpec] = partial( get_gpt_decoder_block_spec, use_transformer_engine=HAVE_TE ) - normalization: str = "RMSNorm" - layernorm_zero_centered_gamma: bool = True # Zero-Centered RMSNorm - rotary_percent: float = 0.25 # Qwen3-Next applies RoPE only to the first 25% of dims - attention_output_gate: bool = True # Qwen3-Next applies output gate to standard attention - kv_channels: Optional[int] = 256 - seq_length: int = 262144 # Qwen3-Next-80B-A3B supports up to 256k sequence length + layernorm_zero_centered_gamma: bool = True # Zero-centered RMSNorm + kv_channels: int | None = 256 + num_query_groups: int = 2 + seq_length: int = 262144 # 256k tokens rotary_base: float = 10000000.0 + rotary_percent: float = 0.25 # 25% of the hidden size is used for RoPE + attention_output_gate: bool = True # Gated Attention - # GDN specific parameters - linear_attention_type: str = "gated_delta_net" - linear_attention_freq: Union[int, List[int]] = 4 # 1 standard attention layer per 4 layers + # MoE specific parameters + num_moe_experts: int = 512 + moe_router_topk: int = 10 # 10 routed experts per token + moe_shared_expert_gate: bool = True # Qwen3-Next uses a gate for the shared expert + moe_router_dtype: str = "fp32" + moe_router_load_balancing_type: str = "global_aux_loss" # Qwen3-Next uses global aux loss for load balancing + + # Linear Attention specific parameters + linear_attention_type: str = "gated_delta_net" # Gated Delta Net used in 75% of the model layers + linear_attention_freq: int | list[int] = 4 # 1 gated standard attention layer per 4 layers linear_conv_kernel_dim: int = 4 linear_key_head_dim: int = 128 linear_value_head_dim: int = 128 linear_num_key_heads: int = 16 linear_num_value_heads: int = 32 - # MoE specific parameters - num_moe_experts: int = 512 - moe_router_topk: int = 10 - moe_ffn_hidden_size: int = 512 - moe_shared_expert_intermediate_size: int = 512 # 512 * 1 shared expert - moe_shared_expert_gate: bool = True - # Checkpointing hetereogenous_dist_checkpoint: bool = True @@ -460,7 +468,7 @@ class Qwen3NextModelProvider(Qwen3MoEModelProvider): @dataclass class Qwen3NextModelProvider80B_A3B(Qwen3NextModelProvider): """ - Provider for Qwen 3 Next: https://huggingface.co/Qwen/Qwen3-Next-80B-A3B-Instruct + Provider for Qwen 3 Next 80B-A3B: https://huggingface.co/Qwen/Qwen3-Next-80B-A3B-Instruct and https://huggingface.co/Qwen/Qwen3-Next-80B-A3B-Thinking """ num_layers: int = 48 @@ -469,4 +477,6 @@ class Qwen3NextModelProvider80B_A3B(Qwen3NextModelProvider): num_query_groups: int = 2 ffn_hidden_size: int = 5120 moe_ffn_hidden_size: int = 512 - mtp_num_layers: int = 0 # Set to 1 if need MTP + moe_shared_expert_intermediate_size: int = 512 + mtp_num_layers: Optional[int] = 0 + mtp_loss_scaling_factor: Optional[float] = None diff --git a/src/megatron/bridge/recipes/qwen/__init__.py b/src/megatron/bridge/recipes/qwen/__init__.py index 74bb946f33e..ca8856e5f6e 100644 --- a/src/megatron/bridge/recipes/qwen/__init__.py +++ b/src/megatron/bridge/recipes/qwen/__init__.py @@ -49,6 +49,11 @@ qwen3_235b_a22b_pretrain_config, ) +# Qwen3-Next models +from .qwen3_next import ( + qwen3_next_80b_a3b_pretrain_config, +) + __all__ = [ # Qwen2 models @@ -79,4 +84,6 @@ # Qwen3 MoE models "qwen3_30b_a3b_pretrain_config", "qwen3_235b_a22b_pretrain_config", + # Qwen3-Next models + "qwen3_next_80b_a3b_pretrain_config", ] diff --git a/src/megatron/bridge/recipes/qwen/qwen3_next.py b/src/megatron/bridge/recipes/qwen/qwen3_next.py new file mode 100644 index 00000000000..cadc25cfd76 --- /dev/null +++ b/src/megatron/bridge/recipes/qwen/qwen3_next.py @@ -0,0 +1,322 @@ +# Copyright (c) 2025, NVIDIA CORPORATION. All rights reserved. +# +# Licensed under the Apache License, Version 2.0 (the "License"); +# you may not use this file except in compliance with the License. +# You may obtain a copy of the License at +# +# http://www.apache.org/licenses/LICENSE-2.0 +# +# Unless required by applicable law or agreed to in writing, software +# distributed under the License is distributed on an "AS IS" BASIS, +# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. +# See the License for the specific language governing permissions and +# limitations under the License. + +import os +from typing import List, Optional, Union + +import torch +from megatron.core.distributed import DistributedDataParallelConfig +from typing_extensions import TypedDict, Unpack + +from megatron.bridge import AutoBridge +from megatron.bridge.recipes.utils.dataset_utils import get_blend_fields_from_data_paths +from megatron.bridge.recipes.utils.optimizer_utils import distributed_fused_adam_with_cosine_annealing +from megatron.bridge.recipes.utils.tokenizer_utils import DEFAULT_NULL_TOKENIZER_VOCAB_SIZE +from megatron.bridge.training.comm_overlap import CommOverlapConfig +from megatron.bridge.training.config import ( + CheckpointConfig, + ConfigContainer, + GPTDatasetConfig, + LoggerConfig, + RNGConfig, + TokenizerConfig, + TrainingConfig, +) +from megatron.bridge.training.mixed_precision import MixedPrecisionConfig, bf16_mixed + + +class Qwen3NextCommonKwargs(TypedDict, total=False): + """Typed options accepted by Qwen3-Next recipe helpers.""" + + # Core identifiers + hf_path: str + dir: Optional[str] + name: str + # Dataset configuration + data_paths: Optional[List[str]] + data_args_path: Optional[str] + train_data_path: Optional[List[str]] + valid_data_path: Optional[List[str]] + test_data_path: Optional[List[str]] + per_split_data_args_path: Optional[str] + mock: bool + # Model configuration + tensor_parallelism: int + pipeline_parallelism: int + pipeline_parallelism_dtype: Optional[torch.dtype] + virtual_pipeline_parallelism: Optional[int] + context_parallelism: int + expert_parallelism: Optional[int] + expert_tensor_parallelism: int + sequence_parallelism: bool + use_megatron_fsdp: bool + enable_recompute: bool + account_for_embedding_in_pipeline_split: bool + account_for_loss_in_pipeline_split: bool + # MTP support + mtp_num_layers: Optional[int] + mtp_loss_scaling_factor: Optional[float] + # Training hyperparameters + train_iters: int + global_batch_size: int + micro_batch_size: int + seq_length: int + lr: float + min_lr: float + lr_warmup_iters: int + lr_decay_iters: Optional[int] + eval_interval: int + save_interval: int + use_null_tokenizer: bool + # Precision / overlap configs + precision_config: Optional[Union[MixedPrecisionConfig, str]] + comm_overlap_config: Optional[CommOverlapConfig] + # Performance optimization knobs + enable_deepep: bool + + +def qwen3_next_80b_a3b_pretrain_config(**user_kwargs: Unpack[Qwen3NextCommonKwargs]) -> ConfigContainer: + """Return a pre-training config for Qwen3-Next 80B-A3B. + + See `_qwen3_next_common` for the full list of parameters. + """ + recommended_kwargs: Qwen3NextCommonKwargs = { + "hf_path": "Qwen/Qwen3-Next-80B-A3B-Instruct", + "tensor_parallelism": 1, + "pipeline_parallelism": 4, + "pipeline_parallelism_dtype": torch.bfloat16, + "context_parallelism": 1, + "expert_parallelism": 8, + "sequence_parallelism": False, + "enable_recompute": True, + } + # Combine defaults with user kwargs; user values take precedence. + combined_kwargs: Qwen3NextCommonKwargs = {**recommended_kwargs, **user_kwargs} + return _qwen3_next_common(**combined_kwargs) + + +def _qwen3_next_common( + hf_path: str, + dir: Optional[str] = None, + name: str = "default", + # Dataset configuration + data_paths: Optional[List[str]] = None, + data_args_path: Optional[str] = None, + train_data_path: Optional[List[str]] = None, + valid_data_path: Optional[List[str]] = None, + test_data_path: Optional[List[str]] = None, + per_split_data_args_path: Optional[str] = None, + mock: bool = False, + path_to_cache: Optional[str] = None, + # Model configuration + tensor_parallelism: int = 4, + pipeline_parallelism: int = 2, + pipeline_parallelism_dtype: Optional[torch.dtype] = torch.bfloat16, + virtual_pipeline_parallelism: Optional[int] = None, + context_parallelism: int = 1, + expert_parallelism: Optional[int] = 4, + expert_tensor_parallelism: int = 1, + sequence_parallelism: bool = True, + use_megatron_fsdp: bool = False, + enable_recompute: bool = False, + account_for_embedding_in_pipeline_split: bool = False, + account_for_loss_in_pipeline_split: bool = False, + # MTP support + mtp_num_layers: Optional[int] = 1, + mtp_loss_scaling_factor: Optional[float] = 0.1, + # Training hyperparameters + train_iters: int = 300000, + global_batch_size: int = 32, + micro_batch_size: int = 2, + seq_length: int = 4096, + lr: float = 3e-4, + min_lr: float = 3e-5, + lr_warmup_iters: int = 500, + lr_decay_iters: Optional[int] = None, + eval_interval: int = 500, + save_interval: int = 500, + use_null_tokenizer: bool = False, + # Precision recipe + precision_config: Optional[Union[MixedPrecisionConfig, str]] = None, + comm_overlap_config: Optional[CommOverlapConfig] = None, + enable_deepep: bool = False, +) -> ConfigContainer: + """ + Create a pre-training configuration for Qwen3 MoE models using a given HuggingFace path. + + Args: + hf_path (str): HuggingFace model path (e.g., "Qwen/Qwen3-30B-A3B", "Qwen/Qwen3-235B-A22B"). + dir (Optional[str]): Base directory for saving logs and checkpoints. + name (str): Name of the pre-training run. + data_paths (Optional[List[str]]): List of paths to dataset files. If None, mock data will be used. + data_args_path (Optional[str]): Path to file containing data arguments. + train_data_path (Optional[List[str]]): List of training data paths. + valid_data_path (Optional[List[str]]): List of validation data paths. + test_data_path (Optional[List[str]]): List of test data paths. + per_split_data_args_path (Optional[str]): Path to JSON file with per-split data configuration. + mock (bool): Whether to use mock data. If True, ignores data_paths. + tensor_parallelism (int): Degree of tensor model parallelism. + pipeline_parallelism (int): Degree of pipeline model parallelism. + pipeline_parallelism_dtype (Optional[torch.dtype]): Data type for pipeline parallelism. + virtual_pipeline_parallelism (Optional[int]): Size of virtual pipeline parallelism. + context_parallelism (int): Degree of context parallelism to be passed to model_config. + expert_parallelism (Optional[int]): Degree of expert parallelism for MoE. + expert_tensor_parallelism (int): Expert tensor parallelism for MoE. + sequence_parallelism (bool): Whether to use sequence parallelism. + use_megatron_fsdp (bool): Whether to use Megatron FSDP. + enable_recompute (bool): Whether to enable recompute for memory optimization. + account_for_embedding_in_pipeline_split (bool): Whether to account for embedding in pipeline split. + account_for_loss_in_pipeline_split (bool): Whether to account for loss in pipeline split. + mtp_num_layers (Optional[int]): Number of layers for MTP. + mtp_loss_scaling_factor (Optional[float]): Loss scaling factor for MTP. + train_iters (int): Total number of training iterations. + global_batch_size (int): Global batch size for training. + micro_batch_size (int): Micro batch size for training. + seq_length (int): Sequence length for training data. + lr (float): Learning rate. + min_lr (float): Minimum learning rate for cosine decay. + lr_warmup_iters (int): Number of warmup iterations for the learning rate. + lr_decay_iters (Optional[int]): Number of iterations over which to decay the LR. + precision_config (Optional[Union[MixedPrecisionConfig, str]]): Precision configuration for the model. + comm_overlap_config (Optional[CommOverlapConfig]): Communication overlap configuration. + enable_deepep (bool): Whether to enable DEEPEP for MoE. + + Returns: + ConfigContainer: Configuration for pre-training. + """ + base_output_dir = dir if dir is not None else os.path.join(os.getcwd(), "nemo_experiments") + run_output_dir = os.path.join(base_output_dir, name) + checkpoint_dir = os.path.join(run_output_dir, "checkpoints") + tensorboard_dir = os.path.join(run_output_dir, "tb_logs") + + blend, blend_per_split, split = get_blend_fields_from_data_paths( + data_paths, data_args_path, train_data_path, valid_data_path, test_data_path, per_split_data_args_path, mock + ) + + bridge = AutoBridge.from_hf_pretrained(hf_path) + model_cfg = bridge.to_megatron_provider(load_weights=False) + model_cfg.tensor_model_parallel_size = tensor_parallelism + model_cfg.sequence_parallel = sequence_parallelism + model_cfg.pipeline_model_parallel_size = pipeline_parallelism + model_cfg.pipeline_dtype = pipeline_parallelism_dtype + model_cfg.virtual_pipeline_model_parallel_size = virtual_pipeline_parallelism + model_cfg.context_parallel_size = context_parallelism + model_cfg.expert_model_parallel_size = expert_parallelism + model_cfg.expert_tensor_parallel_size = expert_tensor_parallelism + + model_cfg.mtp_num_layers = 0 if mtp_num_layers is None else mtp_num_layers + model_cfg.mtp_loss_scaling_factor = mtp_loss_scaling_factor + + # Performance optimization knobs + model_cfg.moe_permute_fusion = True + model_cfg.moe_grouped_gemm = True + if enable_deepep: + model_cfg.moe_token_dispatcher_type = "flex" + # TODO: Remove moe_enable_deepep since it is deprecated + model_cfg.moe_enable_deepep = True + model_cfg.moe_shared_expert_overlap = False + model_cfg.moe_flex_dispatcher_backend = "deepep" + + if precision_config is None: + precision_config = bf16_mixed() + if isinstance(precision_config, MixedPrecisionConfig): + precision_config.grad_reduce_in_fp32 = False + + # MoE-specific pipeline split configurations + if account_for_embedding_in_pipeline_split: + model_cfg.account_for_embedding_in_pipeline_split = True + if account_for_loss_in_pipeline_split: + model_cfg.account_for_loss_in_pipeline_split = True + + # Add recompute settings for memory optimization (used by some MoE models) + if enable_recompute: + model_cfg.recompute_granularity = "selective" + model_cfg.recompute_modules = ["layernorm", "moe", "moe_act"] + model_cfg.recompute_method = None + model_cfg.recompute_num_layers = None + model_cfg.seq_length = seq_length + + opt_config, scheduler = distributed_fused_adam_with_cosine_annealing( + lr_warmup_iters=lr_warmup_iters, + lr_decay_iters=lr_decay_iters, + max_lr=lr, + min_lr=min_lr, + ) + scheduler.no_weight_decay_cond_type = "qwen3_next" + + # Config Container + cfg = ConfigContainer( + model=model_cfg, + train=TrainingConfig( + train_iters=train_iters, + eval_interval=eval_interval, + eval_iters=32, + global_batch_size=global_batch_size, + micro_batch_size=micro_batch_size, + manual_gc=True, + manual_gc_interval=100, + manual_gc_eval=100, + ), + optimizer=opt_config, + scheduler=scheduler, + ddp=DistributedDataParallelConfig( + check_for_nan_in_grad=True, + grad_reduce_in_fp32=True, + overlap_grad_reduce=True, + overlap_param_gather=True, + average_in_collective=True, # Not supported for Megatron FSDP for now, need to be set to False if using Megatron FSDP + data_parallel_sharding_strategy="optim_grads_params", # For Megatron FSDP only + use_distributed_optimizer=True, + use_megatron_fsdp=use_megatron_fsdp, # need use_distributed_optimizer=True + ), + dataset=GPTDatasetConfig( + random_seed=1234, + reset_attention_mask=False, + reset_position_ids=False, + eod_mask_loss=False, + sequence_length=seq_length, + num_dataset_builder_threads=1, + blend=blend, + blend_per_split=blend_per_split, + split=split, + path_to_cache=path_to_cache, + mmap_bin_files=False, + # Dataloader config parameters + data_sharding=True, + dataloader_type="single", + skip_getting_attention_mask_from_dataset=True, + ), + logger=LoggerConfig( + log_interval=10, + tensorboard_dir=tensorboard_dir, + log_timers_to_tensorboard=True, + ), + tokenizer=TokenizerConfig( + tokenizer_type="NullTokenizer" if use_null_tokenizer else "HuggingFaceTokenizer", + tokenizer_model=hf_path if not use_null_tokenizer else None, + vocab_size=DEFAULT_NULL_TOKENIZER_VOCAB_SIZE if use_null_tokenizer else None, + ), + checkpoint=CheckpointConfig( + save_interval=save_interval, + save=checkpoint_dir, + load=checkpoint_dir, + ckpt_format="torch_dist", + fully_parallel_save=True, + ), + rng=RNGConfig(seed=1234), + comm_overlap=comm_overlap_config, + mixed_precision=precision_config, + ) + + return cfg diff --git a/src/megatron/bridge/training/config.py b/src/megatron/bridge/training/config.py index 1ca0a380823..1a7d69d1fab 100644 --- a/src/megatron/bridge/training/config.py +++ b/src/megatron/bridge/training/config.py @@ -416,6 +416,12 @@ class SchedulerConfig: weight_decay_incr_style: Literal["constant", "linear", "cosine"] = "constant" """Weight decay increment function.""" + no_weight_decay_cond_type: Optional[Literal["qwen3_next"]] = None + """Type of no weight decay condition. Choices: + None (default): param no weight decay if and only if it is 1D; or it is bias; + or it is embedding and embedding_init_method_std is not None. + "qwen3_next": In addition to the default rules, apply weight decay to qk layernorm as a special case.""" + lr_warmup_steps: Optional[int] = field(init=False, default=None) lr_decay_steps: Optional[int] = field(init=False, default=None) wd_incr_steps: Optional[int] = field(init=False, default=None) diff --git a/src/megatron/bridge/training/setup.py b/src/megatron/bridge/training/setup.py index 754c737e701..b62930a5825 100644 --- a/src/megatron/bridge/training/setup.py +++ b/src/megatron/bridge/training/setup.py @@ -42,6 +42,7 @@ from megatron.bridge.training.state import GlobalState from megatron.bridge.training.tokenizers.tokenizer import build_tokenizer from megatron.bridge.training.utils.log_utils import append_to_progress_log, barrier_and_log, setup_logging +from megatron.bridge.training.utils.weight_decay_utils import get_no_weight_decay_cond from megatron.bridge.utils.common_utils import print_rank_0, get_rank_safe from megatron.bridge.training.tensor_inspect import ( finalize_tensor_inspect_post_model_initialization, @@ -186,11 +187,16 @@ def setup( ) cfg.model.timers = timers cfg.optimizer.timers = timers + no_weight_decay_cond = get_no_weight_decay_cond( + cfg.scheduler.no_weight_decay_cond_type, + default_skip_embedding_weight_decay=cfg.model.embedding_init_method_std is not None, + ) optimizer, scheduler = setup_optimizer( optimizer_config=cfg.optimizer, scheduler_config=cfg.scheduler, model=model, use_gloo_process_groups=cfg.dist.use_gloo_process_groups, + no_weight_decay_cond=no_weight_decay_cond, ) timers("model-and-optimizer-setup").stop() barrier_and_log("after model, optimizer, and learning rate scheduler are built") diff --git a/src/megatron/bridge/training/utils/weight_decay_utils.py b/src/megatron/bridge/training/utils/weight_decay_utils.py new file mode 100644 index 00000000000..cacd476d297 --- /dev/null +++ b/src/megatron/bridge/training/utils/weight_decay_utils.py @@ -0,0 +1,45 @@ +# Copyright (c) 2025, 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 typing import Callable + +import torch + + +def get_no_weight_decay_cond( + no_weight_decay_cond_type: str, default_skip_embedding_weight_decay: bool +) -> Callable[[str, torch.Tensor], bool]: + """Get the no weight decay condition function.""" + + # Default case: no_wd_decay_cond_type is None + no_weight_decay_cond_fn = None + + if no_weight_decay_cond_type == "qwen3_next": + # Qwen3-Next applies weight decay to qk layernorm as a special case + def qwen3_next_no_weight_decay_cond(name, param): + if "q_layernorm" in name or "k_layernorm" in name: + no_wd = False + else: + no_wd = ( + name.endswith(".bias") + or len(param.shape) == 1 + or (default_skip_embedding_weight_decay and "embedding" in name) + ) + return no_wd + + no_weight_decay_cond_fn = qwen3_next_no_weight_decay_cond + elif no_weight_decay_cond_type is not None: + raise ValueError(f"Invalid no_weight_decay_cond_type: {no_weight_decay_cond_type}") + + return no_weight_decay_cond_fn diff --git a/tests/unit_tests/models/qwen/test_qwen3_next_model_provider.py b/tests/unit_tests/models/qwen/test_qwen3_next_model_provider.py new file mode 100644 index 00000000000..510d5b371f5 --- /dev/null +++ b/tests/unit_tests/models/qwen/test_qwen3_next_model_provider.py @@ -0,0 +1,490 @@ +# Copyright (c) 2025, NVIDIA CORPORATION. All rights reserved. +# +# Licensed under the Apache License, Version 2.0 (the "License"); +# you may not use this file except in compliance with the License. +# You may obtain a copy of the License at +# +# http://www.apache.org/licenses/LICENSE-2.0 +# +# Unless required by applicable law or agreed to in writing, software +# distributed under the License is distributed on an "AS IS" BASIS, +# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. +# See the License for the specific language governing permissions and +# limitations under the License. + +import torch +import torch.nn.functional as F + +from megatron.bridge.models.qwen import ( + Qwen3NextModelProvider, + Qwen3NextModelProvider80B_A3B, +) + + +class TestQwen3NextModelProvider: + """Test cases for base Qwen3NextModelProvider class.""" + + def test_qwen3_next_model_provider_initialization(self): + """Test Qwen3NextModelProvider can be initialized with default values.""" + provider = Qwen3NextModelProvider( + num_layers=32, + hidden_size=4096, + num_attention_heads=32, + ) + + # Check required transformer config fields + assert provider.num_layers == 32 + assert provider.hidden_size == 4096 + assert provider.num_attention_heads == 32 + + # Check Qwen3 MoE-specific defaults + assert provider.normalization == "RMSNorm" + assert provider.activation_func is F.silu + assert provider.gated_linear_unit is True + assert provider.layernorm_zero_centered_gamma is True + assert provider.add_bias_linear is False + assert provider.add_qkv_bias is False + assert provider.qk_layernorm is True + assert provider.kv_channels == 256 + assert provider.num_query_groups == 2 + assert provider.seq_length == 262144 # 256k tokens + assert provider.init_method_std == 0.02 + assert provider.hidden_dropout == 0.0 + assert provider.attention_dropout == 0.0 + assert provider.vocab_size == 151936 + assert provider.share_embeddings_and_output_weights is False + assert provider.layernorm_epsilon == 1e-6 + assert provider.rotary_base == 10000000.0 + assert provider.rotary_percent == 0.25 + assert provider.attention_output_gate is True + assert provider.position_embedding_type == "rope" + assert provider.autocast_dtype == torch.bfloat16 + assert provider.params_dtype == torch.bfloat16 + assert provider.bf16 is True + + # Check MoE-specific defaults + assert provider.num_moe_experts == 512 + assert provider.moe_router_load_balancing_type == "global_aux_loss" + assert provider.moe_aux_loss_coeff == 1e-3 + assert provider.moe_router_topk == 10 + assert provider.moe_router_pre_softmax is False + assert provider.moe_grouped_gemm is True + assert provider.moe_token_dispatcher_type == "alltoall" + assert provider.moe_permute_fusion is True + assert provider.moe_shared_expert_gate is True + assert provider.moe_router_dtype == "fp32" + + # Linear Attention specific defaults + assert provider.linear_attention_type == "gated_delta_net" + assert provider.linear_attention_freq == 4 + assert provider.linear_conv_kernel_dim == 4 + assert provider.linear_key_head_dim == 128 + assert provider.linear_value_head_dim == 128 + assert provider.linear_num_key_heads == 16 + assert provider.linear_num_value_heads == 32 + + # Checkpointing specific defaults + assert provider.hetereogenous_dist_checkpoint is True + + def test_qwen3_next_model_provider_with_custom_moe_config(self): + """Test Qwen3NextModelProvider with custom MoE configuration.""" + provider = Qwen3NextModelProvider( + num_layers=32, + hidden_size=4096, + num_attention_heads=32, + num_moe_experts=64, + moe_router_topk=4, + moe_aux_loss_coeff=1e-2, + moe_ffn_hidden_size=256, + ffn_hidden_size=1024, + moe_shared_expert_gate=False, + ) + + assert provider.num_moe_experts == 64 + assert provider.moe_router_topk == 4 + assert provider.moe_aux_loss_coeff == 1e-2 + assert provider.moe_ffn_hidden_size == 256 + assert provider.ffn_hidden_size == 1024 + assert provider.moe_shared_expert_gate is False + + def test_qwen3_next_model_provider_with_custom_linear_attention_config(self): + """Test Qwen3NextModelProvider with custom Linear Attention configuration.""" + provider = Qwen3NextModelProvider( + num_layers=32, + hidden_size=4096, + num_attention_heads=32, + linear_attention_type="gated_delta_net", + linear_attention_freq=2, + linear_conv_kernel_dim=8, + linear_key_head_dim=64, + linear_value_head_dim=64, + linear_num_key_heads=8, + linear_num_value_heads=16, + ) + + assert provider.linear_attention_type == "gated_delta_net" + assert provider.linear_attention_freq == 2 + assert provider.linear_conv_kernel_dim == 8 + assert provider.linear_key_head_dim == 64 + assert provider.linear_value_head_dim == 64 + assert provider.linear_num_key_heads == 8 + assert provider.linear_num_value_heads == 16 + + def test_qwen3_next_model_provider_with_custom_rope(self): + """Test Qwen3NextModelProvider with custom RoPE configuration.""" + provider = Qwen3NextModelProvider( + num_layers=32, + hidden_size=4096, + num_attention_heads=32, + rotary_base=500000.0, + rotary_percent=0.5, + ) + + assert provider.rotary_base == 500000.0 + assert provider.rotary_percent == 0.5 + + def test_qwen3_next_model_provider_ffn_hidden_size(self): + """Test Qwen3NextModelProvider FFN hidden size configuration.""" + provider = Qwen3NextModelProvider( + num_layers=32, + hidden_size=4096, + num_attention_heads=32, + ffn_hidden_size=11008, + moe_ffn_hidden_size=1536, + moe_shared_expert_intermediate_size=1024, + ) + + assert provider.ffn_hidden_size == 11008 + assert provider.moe_ffn_hidden_size == 1536 + assert provider.moe_shared_expert_intermediate_size == 1024 + + def test_qwen3_next_model_provider_group_query_attention(self): + """Test Qwen3NextModelProvider with group query attention.""" + provider = Qwen3NextModelProvider( + num_layers=32, + hidden_size=4096, + num_attention_heads=32, + num_query_groups=4, + ) + + assert provider.num_query_groups == 4 + + def test_qwen3_next_model_provider_custom_vocab_size(self): + """Test Qwen3NextModelProvider with custom vocabulary size.""" + provider = Qwen3NextModelProvider( + num_layers=32, + hidden_size=4096, + num_attention_heads=32, + vocab_size=32000, + ) + + assert provider.vocab_size == 32000 + + def test_qwen3_next_model_provider_custom_sequence_length(self): + """Test Qwen3NextModelProvider with custom sequence length.""" + provider = Qwen3NextModelProvider( + num_layers=32, + hidden_size=4096, + num_attention_heads=32, + seq_length=8192, + ) + + assert provider.seq_length == 8192 + + def test_qwen3_next_model_provider_qk_layernorm(self): + """Test Qwen3NextModelProvider has Zero-Centered QK layernorm enabled by default.""" + provider = Qwen3NextModelProvider( + num_layers=32, + hidden_size=4096, + num_attention_heads=32, + ) + + assert provider.qk_layernorm is True + assert provider.layernorm_zero_centered_gamma is True + + def test_qwen3_next_model_provider_dtype_configuration(self): + """Test Qwen3NextModelProvider dtype configuration.""" + provider = Qwen3NextModelProvider( + num_layers=32, + hidden_size=4096, + num_attention_heads=32, + fp16=True, + bf16=False, + params_dtype=torch.float16, + ) + + assert provider.fp16 is True + assert provider.bf16 is False + assert provider.params_dtype == torch.float16 + + def test_qwen3_next_model_provider_custom_mtp_configuration(self): + """Test Qwen3NextModelProvider with custom MTP configuration.""" + provider = Qwen3NextModelProvider( + num_layers=32, + hidden_size=4096, + num_attention_heads=32, + mtp_num_layers=2, + mtp_loss_scaling_factor=0.2, + ) + assert provider.mtp_num_layers == 2 + assert provider.mtp_loss_scaling_factor == 0.2 + + +class TestQwen3NextModelProvider80B_A3B: + """Test cases for Qwen3NextModelProvider80B_A3B class.""" + + def test_qwen3_next_80b_a3b_default_configuration(self): + """Test Qwen3 Next 80B-A3B model has correct default configuration.""" + provider = Qwen3NextModelProvider80B_A3B() + + # Check Qwen3 Next 80B-A3B specific configuration + assert provider.num_layers == 48 + assert provider.hidden_size == 2048 + assert provider.num_attention_heads == 16 + assert provider.num_query_groups == 2 + assert provider.ffn_hidden_size == 5120 + assert provider.moe_ffn_hidden_size == 512 + assert provider.moe_shared_expert_intermediate_size == 512 + + # Check inherited MoE defaults + assert provider.num_moe_experts == 512 + assert provider.moe_router_topk == 10 + assert provider.moe_shared_expert_gate is True + assert provider.moe_router_dtype == "fp32" + assert provider.moe_router_pre_softmax is False + + # Check inherited Linear Attention defaults + assert provider.linear_attention_type == "gated_delta_net" + assert provider.linear_attention_freq == 4 + assert provider.linear_conv_kernel_dim == 4 + assert provider.linear_key_head_dim == 128 + assert provider.linear_value_head_dim == 128 + assert provider.linear_num_key_heads == 16 + assert provider.linear_num_value_heads == 32 + + # Check inherited base defaults + assert provider.layernorm_zero_centered_gamma is True + assert provider.qk_layernorm is True + assert provider.kv_channels == 256 + assert provider.vocab_size == 151936 + assert provider.seq_length == 262144 # 256k tokens + assert provider.rotary_base == 10000000.0 + assert provider.rotary_percent == 0.25 + assert provider.attention_output_gate is True + assert provider.normalization == "RMSNorm" + assert provider.activation_func is F.silu + assert provider.gated_linear_unit is True + assert provider.hetereogenous_dist_checkpoint is True + + def test_qwen3_next_80b_a3b_override_configuration(self): + """Test Qwen3 Next 80B-A3B model with overridden configuration.""" + provider = Qwen3NextModelProvider80B_A3B( + seq_length=8192, + num_moe_experts=64, + moe_router_topk=4, + ) + + # Check overridden values + assert provider.seq_length == 8192 + assert provider.num_moe_experts == 64 + assert provider.moe_router_topk == 4 + + # Check defaults remain + assert provider.num_layers == 48 + assert provider.hidden_size == 2048 + assert provider.num_attention_heads == 16 + assert provider.num_query_groups == 2 + assert provider.ffn_hidden_size == 5120 + assert provider.moe_ffn_hidden_size == 512 + assert provider.moe_shared_expert_intermediate_size == 512 + assert provider.moe_router_pre_softmax is False + + +class TestQwen3NextProviderInheritance: + """Test inheritance relationships between Qwen3 Next providers.""" + + def test_qwen3_next_models_inherit_from_base(self): + """Test Qwen3 Next providers inherit from Qwen3NextModelProvider.""" + assert issubclass(Qwen3NextModelProvider80B_A3B, Qwen3NextModelProvider) + + def test_provide_method_inherited(self): + """Test that provide method works correctly in inherited classes.""" + # Test with Qwen3 Next 80B-A3B + provider = Qwen3NextModelProvider80B_A3B() + + # The provide method should be inherited from GPTModelProvider + assert hasattr(provider, "provide") + assert callable(provider.provide) + + +class TestQwen3NextProviderEdgeCases: + """Test edge cases and error conditions for Qwen3 Next providers.""" + + def test_valid_num_query_groups(self): + """Test that valid num_query_groups configuration works.""" + # num_attention_heads must be divisible by num_query_groups + provider = Qwen3NextModelProvider( + num_layers=32, + hidden_size=4096, + num_attention_heads=32, + num_query_groups=8, # 32 divisible by 8 + ) + assert provider.num_query_groups == 8 + + def test_valid_linear_attention_num_key_heads(self): + """Test that valid linear_num_key_heads configuration works.""" + provider = Qwen3NextModelProvider( + num_layers=32, + hidden_size=4096, + linear_num_value_heads=64, + linear_num_key_heads=8, # 64 divisible by 8 + ) + assert provider.linear_num_key_heads == 8 + + def test_valid_linear_attention_frequency(self): + """Test that valid linear_attention_frequency configuration works.""" + provider = Qwen3NextModelProvider( + num_layers=32, + hidden_size=4096, + linear_attention_freq=2, # 32 divisible by 2 + ) + assert provider.linear_attention_freq == 2 + + def test_valid_linear_attention_frequency_list(self): + """Test that valid linear_attention_frequency configuration works.""" + provider = Qwen3NextModelProvider( + num_layers=8, + hidden_size=4096, + linear_attention_freq=[1, 1, 1, 0, 1, 1, 1, 0], # 8 layers, 6 linear attention layers + ) + assert provider.linear_attention_freq == [1, 1, 1, 0, 1, 1, 1, 0] + + def test_moe_configuration_validity(self): + """Test MoE configuration parameters.""" + provider = Qwen3NextModelProvider( + num_layers=32, + hidden_size=4096, + num_attention_heads=32, + num_moe_experts=64, + moe_router_topk=4, + ) + + # moe_router_topk should be <= num_moe_experts + assert provider.moe_router_topk <= provider.num_moe_experts + assert provider.num_moe_experts == 64 + assert provider.moe_router_topk == 4 + + def test_vocabulary_size_divisibility(self): + """Test vocabulary size divisibility configuration.""" + provider = Qwen3NextModelProvider( + num_layers=32, + hidden_size=4096, + num_attention_heads=32, + vocab_size=152064, + make_vocab_size_divisible_by=128, + ) + + # The actual vocab size should be adjusted if needed + assert provider.make_vocab_size_divisible_by == 128 + + def test_seq_length_override(self): + """Test sequence length configuration.""" + provider = Qwen3NextModelProvider( + num_layers=32, + hidden_size=4096, + num_attention_heads=32, + seq_length=8192, + ) + + assert provider.seq_length == 8192 + + def test_rotary_base_configuration(self): + """Test rotary base configuration.""" + provider = Qwen3NextModelProvider( + num_layers=32, + hidden_size=4096, + num_attention_heads=32, + rotary_base=500000.0, + ) + + assert provider.rotary_base == 500000.0 + + def test_layernorm_epsilon_override(self): + """Test layernorm epsilon configuration.""" + provider = Qwen3NextModelProvider( + num_layers=32, + hidden_size=4096, + num_attention_heads=32, + layernorm_epsilon=1e-5, + ) + + assert provider.layernorm_epsilon == 1e-5 + + +class TestQwen3NextProviderArchitecturalFeatures: + """Test cases to verify Qwen3 MoE architectural features.""" + + def test_qwen3_next_qk_layernorm_feature(self): + """Test that Qwen3 Next models have Zero-Centered QK layernorm enabled.""" + provider = Qwen3NextModelProvider80B_A3B() + + assert provider.qk_layernorm is True + assert provider.layernorm_zero_centered_gamma is True + + def test_qwen3_next_dtype_defaults(self): + """Test that Qwen3 Next models have correct dtype defaults.""" + provider = Qwen3NextModelProvider80B_A3B() + + assert provider.autocast_dtype == torch.bfloat16 + assert provider.params_dtype == torch.bfloat16 + assert provider.bf16 is True + + def test_qwen3_next_kv_channels(self): + """Test that Qwen3 Next models have correct KV channels configuration.""" + provider = Qwen3NextModelProvider80B_A3B() + + assert provider.kv_channels == 256 + + def test_qwen3_next_rope_percent(self): + """Test that Qwen3 Next models have correct RoPE percent configuration.""" + provider = Qwen3NextModelProvider80B_A3B() + + assert provider.rotary_percent == 0.25 + + def test_qwen3_next_bias_configuration(self): + """Test that Qwen3 Next models have correct bias configuration.""" + provider = Qwen3NextModelProvider80B_A3B() + + assert provider.add_bias_linear is False + assert provider.add_qkv_bias is False + + def test_qwen3_next_attention_output_gate(self): + """Test that Qwen3 Next models have correct attention output gate configuration.""" + provider = Qwen3NextModelProvider80B_A3B() + + assert provider.attention_output_gate is True + + def test_qwen3_next_shared_expert(self): + """Test that Qwen3 Next models have correct shared expert configuration.""" + provider = Qwen3NextModelProvider80B_A3B() + + assert provider.moe_shared_expert_gate is True + assert provider.moe_shared_expert_intermediate_size == 512 + + def test_qwen3_next_linear_attention_type(self): + """Test that Qwen3 Next models have correct linear attention type (Gated Delta Net) configuration.""" + provider = Qwen3NextModelProvider80B_A3B() + + assert provider.linear_attention_type == "gated_delta_net" + + def test_qwen3_next_linear_attention_freq(self): + """Test that Qwen3 Next models have correct linear attention frequency configuration.""" + provider = Qwen3NextModelProvider80B_A3B() + + assert provider.linear_attention_freq == 4 + + def test_qwen3_next_moe_router_load_balancing_type(self): + """Test that Qwen3 Next models have correct Moe Router Load Balancing Type configuration.""" + provider = Qwen3NextModelProvider80B_A3B() + + assert provider.moe_router_load_balancing_type == "global_aux_loss"