diff --git a/src/megatron/bridge/recipes/mamba/__init__.py b/src/megatron/bridge/recipes/mamba/__init__.py index e69de29bb2..03e02c0a8c 100644 --- a/src/megatron/bridge/recipes/mamba/__init__.py +++ b/src/megatron/bridge/recipes/mamba/__init__.py @@ -0,0 +1,34 @@ +# 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 .mamba2 import ( + mamba2_1p3b_pretrain_config, + mamba2_2p7b_pretrain_config, + mamba2_8b_pretrain_config, + mamba2_130m_pretrain_config, + mamba2_370m_pretrain_config, + mamba2_780m_pretrain_config, + mamba2_hybrid_8b_pretrain_config, +) + + +__all__ = [ + "mamba2_130m_pretrain_config", + "mamba2_370m_pretrain_config", + "mamba2_780m_pretrain_config", + "mamba2_1p3b_pretrain_config", + "mamba2_2p7b_pretrain_config", + "mamba2_8b_pretrain_config", + "mamba2_hybrid_8b_pretrain_config", +] diff --git a/src/megatron/bridge/recipes/mamba/mamba2.py b/src/megatron/bridge/recipes/mamba/mamba2.py new file mode 100644 index 0000000000..0272b061b1 --- /dev/null +++ b/src/megatron/bridge/recipes/mamba/mamba2.py @@ -0,0 +1,340 @@ +# 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 + +import torch +from typing_extensions import TypedDict, Unpack + +from megatron.bridge.models.mamba import ( + MambaModelProvider1P3B, + MambaModelProvider2P7B, + MambaModelProvider130M, + MambaModelProvider370M, + MambaModelProvider780M, + NVIDIAMambaHybridProvider8B, + NVIDIAMambaModelProvider8B, +) +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, + DistributedDataParallelConfig, + GPTDatasetConfig, + LoggerConfig, + RNGConfig, + TokenizerConfig, + TrainingConfig, +) +from megatron.bridge.training.mixed_precision import MixedPrecisionConfig + + +class Mamba2CommonKwargs(TypedDict, total=False): + """Typed options accepted by Mamba2 recipe helper functions.""" + + # Core identifiers + model_provider: ( + type[MambaModelProvider130M] + | type[MambaModelProvider370M] + | type[MambaModelProvider780M] + | type[MambaModelProvider1P3B] + | type[MambaModelProvider2P7B] + | type[NVIDIAMambaModelProvider8B] + | type[NVIDIAMambaHybridProvider8B] + ) + tokenizer_model: str | None + dir: str | None + name: str + # Dataset configuration + data_paths: list[str] | None + data_args_path: str | None + train_data_path: list[str] | None + valid_data_path: list[str] | None + test_data_path: list[str] | None + per_split_data_args_path: str | None + mock: bool + # Model configuration + tensor_parallelism: int + pipeline_parallelism: int + pipeline_parallelism_dtype: torch.dtype | None + virtual_pipeline_parallelism: int | None + context_parallelism: int + sequence_parallelism: bool + # 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: int | None + # Tokenizer selection + use_null_tokenizer: bool + # Precision / overlap configs + precision_config: MixedPrecisionConfig | str | None + comm_overlap_config: CommOverlapConfig | None + + +def mamba2_130m_pretrain_config(**user_kwargs: Unpack[Mamba2CommonKwargs]) -> ConfigContainer: + """Return a pre-training config for Mamba2 130M.""" + recommended: Mamba2CommonKwargs = { + "model_provider": MambaModelProvider130M, + "tensor_parallelism": 1, + "pipeline_parallelism": 1, + "sequence_parallelism": False, + "precision_config": "bf16_mixed", + "use_null_tokenizer": False, + } + kwargs: Mamba2CommonKwargs = {**recommended, **user_kwargs} + return _mamba2_common(tokenizer_model=kwargs.get("tokenizer_model"), **kwargs) + + +def mamba2_370m_pretrain_config(**user_kwargs: Unpack[Mamba2CommonKwargs]) -> ConfigContainer: + """Return a pre-training config for Mamba2 370M.""" + recommended: Mamba2CommonKwargs = { + "model_provider": MambaModelProvider370M, + "tensor_parallelism": 1, + "pipeline_parallelism": 1, + "sequence_parallelism": False, + "precision_config": "bf16_mixed", + "use_null_tokenizer": False, + } + kwargs: Mamba2CommonKwargs = {**recommended, **user_kwargs} + return _mamba2_common(tokenizer_model=kwargs.get("tokenizer_model"), **kwargs) + + +def mamba2_780m_pretrain_config(**user_kwargs: Unpack[Mamba2CommonKwargs]) -> ConfigContainer: + """Return a pre-training config for Mamba2 780M.""" + recommended: Mamba2CommonKwargs = { + "model_provider": MambaModelProvider780M, + "tensor_parallelism": 1, + "pipeline_parallelism": 1, + "sequence_parallelism": False, + "precision_config": "bf16_mixed", + "use_null_tokenizer": False, + } + kwargs: Mamba2CommonKwargs = {**recommended, **user_kwargs} + return _mamba2_common(tokenizer_model=kwargs.get("tokenizer_model"), **kwargs) + + +def mamba2_1p3b_pretrain_config(**user_kwargs: Unpack[Mamba2CommonKwargs]) -> ConfigContainer: + """Return a pre-training config for Mamba2 1.3B.""" + recommended: Mamba2CommonKwargs = { + "model_provider": MambaModelProvider1P3B, + "tensor_parallelism": 1, + "pipeline_parallelism": 1, + "sequence_parallelism": False, + "precision_config": "bf16_mixed", + "use_null_tokenizer": False, + } + kwargs: Mamba2CommonKwargs = {**recommended, **user_kwargs} + return _mamba2_common(tokenizer_model=kwargs.get("tokenizer_model"), **kwargs) + + +def mamba2_2p7b_pretrain_config(**user_kwargs: Unpack[Mamba2CommonKwargs]) -> ConfigContainer: + """Return a pre-training config for Mamba2 2.7B.""" + recommended: Mamba2CommonKwargs = { + "model_provider": MambaModelProvider2P7B, + "tensor_parallelism": 1, + "pipeline_parallelism": 1, + "sequence_parallelism": False, + "precision_config": "bf16_mixed", + "use_null_tokenizer": False, + } + kwargs: Mamba2CommonKwargs = {**recommended, **user_kwargs} + return _mamba2_common(tokenizer_model=kwargs.get("tokenizer_model"), **kwargs) + + +def mamba2_8b_pretrain_config(**user_kwargs: Unpack[Mamba2CommonKwargs]) -> ConfigContainer: + """Return a pre-training config for Mamba2 8B.""" + recommended: Mamba2CommonKwargs = { + "model_provider": NVIDIAMambaModelProvider8B, + "tensor_parallelism": 8, + "pipeline_parallelism": 1, + "sequence_parallelism": False, + "precision_config": "bf16_mixed", + "use_null_tokenizer": True, + } + kwargs: Mamba2CommonKwargs = {**recommended, **user_kwargs} + return _mamba2_common(tokenizer_model=None, **kwargs) + + +def mamba2_hybrid_8b_pretrain_config(**user_kwargs: Unpack[Mamba2CommonKwargs]) -> ConfigContainer: + """Return a pre-training config for Mamba2 Hybrid 8B.""" + recommended: Mamba2CommonKwargs = { + "model_provider": NVIDIAMambaHybridProvider8B, + "tensor_parallelism": 8, + "pipeline_parallelism": 1, + "sequence_parallelism": False, + "precision_config": "bf16_mixed", + "use_null_tokenizer": True, + } + kwargs: Mamba2CommonKwargs = {**recommended, **user_kwargs} + return _mamba2_common(tokenizer_model=None, **kwargs) + + +def _mamba2_common( + model_provider: ( + type[MambaModelProvider130M] + | type[MambaModelProvider370M] + | type[MambaModelProvider780M] + | type[MambaModelProvider1P3B] + | type[MambaModelProvider2P7B] + | type[NVIDIAMambaModelProvider8B] + | type[NVIDIAMambaHybridProvider8B] + ), + tokenizer_model: str | None = None, + dir: str | None = None, + name: str = "default", + # Dataset configuration + data_paths: list[str] | None = None, + data_args_path: str | None = None, + train_data_path: list[str] | None = None, + valid_data_path: list[str] | None = None, + test_data_path: list[str] | None = None, + per_split_data_args_path: str | None = None, + mock: bool = False, + # Model configuration + tensor_parallelism: int = 1, + pipeline_parallelism: int = 1, + pipeline_parallelism_dtype: torch.dtype | None = None, + virtual_pipeline_parallelism: int | None = None, + context_parallelism: int = 1, + sequence_parallelism: bool = False, + # Training hyperparameters + train_iters: int = 1_168_251, + global_batch_size: int = 8, + micro_batch_size: int = 1, + seq_length: int = 4096, + lr: float = 3e-4, + min_lr: float = 3e-5, + lr_warmup_iters: int = 2000, + lr_decay_iters: int | None = None, + # Tokenizer selection + use_null_tokenizer: bool = False, + # Precision recipe + precision_config: MixedPrecisionConfig | str | None = "bf16_mixed", + comm_overlap_config: CommOverlapConfig | None = None, +) -> ConfigContainer: + """ + Create a pre-training configuration for Mamba 2.x models. + + Args mirror the individual recipe helpers; see those functions for recommended defaults. + """ + 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 + ) + + model_cfg = model_provider( + tensor_model_parallel_size=tensor_parallelism, + pipeline_model_parallel_size=pipeline_parallelism, + pipeline_dtype=pipeline_parallelism_dtype, + virtual_pipeline_model_parallel_size=virtual_pipeline_parallelism, + context_parallel_size=context_parallelism, + sequence_parallel=sequence_parallelism, + ) + + opt_config, scheduler = distributed_fused_adam_with_cosine_annealing( + lr_warmup_iters=lr_warmup_iters, + lr_decay_iters=lr_decay_iters, + adam_beta1=0.9, + adam_beta2=0.95, + adam_eps=1e-5, + weight_decay=0.1, + max_lr=lr, + min_lr=min_lr, + ) + + cfg = ConfigContainer( + model=model_cfg, + train=TrainingConfig( + train_iters=train_iters, + eval_interval=100, + eval_iters=32, + global_batch_size=global_batch_size, + micro_batch_size=micro_batch_size, + ), + 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, + 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, + data_sharding=True, + dataloader_type="single", + num_workers=8, + skip_getting_attention_mask_from_dataset=True, + ), + logger=LoggerConfig( + log_interval=10, + tensorboard_dir=tensorboard_dir, + ), + tokenizer=( + TokenizerConfig( + tokenizer_type="NullTokenizer", + tokenizer_model=None, + vocab_size=DEFAULT_NULL_TOKENIZER_VOCAB_SIZE, + ) + if use_null_tokenizer + else TokenizerConfig( + tokenizer_type="HuggingFaceTokenizer", + tokenizer_model=tokenizer_model or "EleutherAI/gpt-neox-20b", + ) + ), + checkpoint=CheckpointConfig( + save_interval=2000, + save=checkpoint_dir, + load=checkpoint_dir, + ckpt_format="torch_dist", + fully_parallel_load=True, + ), + rng=RNGConfig(seed=1234), + comm_overlap=comm_overlap_config, + mixed_precision=precision_config, + ) + + return cfg + + +__all__ = [ + "mamba2_130m_pretrain_config", + "mamba2_370m_pretrain_config", + "mamba2_780m_pretrain_config", + "mamba2_1p3b_pretrain_config", + "mamba2_2p7b_pretrain_config", + "mamba2_8b_pretrain_config", + "mamba2_hybrid_8b_pretrain_config", +] diff --git a/src/megatron/bridge/recipes/mamba/mamba2_130m.py b/src/megatron/bridge/recipes/mamba/mamba2_130m.py deleted file mode 100644 index 740eaea7fb..0000000000 --- a/src/megatron/bridge/recipes/mamba/mamba2_130m.py +++ /dev/null @@ -1,214 +0,0 @@ -# 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 Optional, Union - -import torch - -from megatron.bridge.models.mamba import MambaModelProvider130M -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.training.comm_overlap import CommOverlapConfig -from megatron.bridge.training.config import ( - CheckpointConfig, - ConfigContainer, - DistributedDataParallelConfig, - GPTDatasetConfig, - LoggerConfig, - RNGConfig, - TokenizerConfig, - TrainingConfig, -) -from megatron.bridge.training.mixed_precision import MixedPrecisionConfig - - -def model_config( - tensor_parallelism: int = 1, - pipeline_parallelism: int = 1, - pipeline_parallelism_dtype: Optional[torch.dtype] = None, - virtual_pipeline_parallelism: Optional[int] = None, - context_parallelism: int = 1, - sequence_parallelism: bool = False, -) -> MambaModelProvider130M: - """ - Configure the Mamba 130M model. - - Args: - tensor_parallelism: Degree of tensor model parallelism. - pipeline_parallelism: Degree of pipeline model parallelism. - pipeline_parallelism_dtype: Data type for pipeline parallelism. - virtual_pipeline_parallelism: Size of virtual pipeline parallelism. - context_parallelism: Degree of context parallelism. - sequence_parallelism: Whether to use sequence parallelism. - - Returns: - MambaModelProvider130M: Configuration for the Mamba 130M model. - """ - return MambaModelProvider130M( - tensor_model_parallel_size=tensor_parallelism, - pipeline_model_parallel_size=pipeline_parallelism, - pipeline_dtype=pipeline_parallelism_dtype, - virtual_pipeline_model_parallel_size=virtual_pipeline_parallelism, - context_parallel_size=context_parallelism, - sequence_parallel=sequence_parallelism, - ) - - -def pretrain_config( - 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, - # Model configuration - tensor_parallelism: int = 1, - pipeline_parallelism: int = 1, - pipeline_parallelism_dtype: Optional[torch.dtype] = None, - virtual_pipeline_parallelism: Optional[int] = None, - context_parallelism: int = 1, - sequence_parallelism: bool = False, - # Training hyperparameters - train_iters: int = 1_168_251, - global_batch_size: int = 8, - micro_batch_size: int = 1, - seq_length: int = 4096, - lr: float = 3e-4, - min_lr: float = 3e-5, - lr_warmup_iters: int = 2000, - lr_decay_iters: Optional[int] = None, - # Precision recipe - precision_config: Optional[Union[MixedPrecisionConfig, str]] = "bf16_mixed", - comm_overlap_config: Optional[CommOverlapConfig] = None, -) -> ConfigContainer: - """ - Create a pre-training configuration for Mamba 130M model. - - Args: - 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. - sequence_parallelism (bool): Whether to use sequence parallelism. - 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 for learning rate decay. - precision_config (Optional[Union[MixedPrecisionConfig, str]]): Precision configuration for the model. - comm_overlap_config (Optional[CommOverlapConfig]): Communication overlap configuration for the model. - - 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 - ) - - model_cfg = model_config( - tensor_parallelism=tensor_parallelism, - pipeline_parallelism=pipeline_parallelism, - pipeline_parallelism_dtype=pipeline_parallelism_dtype, - virtual_pipeline_parallelism=virtual_pipeline_parallelism, - context_parallelism=context_parallelism, - sequence_parallelism=sequence_parallelism, - ) - - opt_config, scheduler = distributed_fused_adam_with_cosine_annealing( - lr_warmup_iters=lr_warmup_iters, - lr_decay_iters=lr_decay_iters, - adam_beta1=0.9, - adam_beta2=0.95, - adam_eps=1e-5, - weight_decay=0.1, - max_lr=lr, - min_lr=min_lr, - ) - - # Config Container - cfg = ConfigContainer( - model=model_cfg, - train=TrainingConfig( - train_iters=train_iters, - eval_interval=100, - eval_iters=32, - global_batch_size=global_batch_size, - micro_batch_size=micro_batch_size, - ), - 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, - 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, - # Dataloader config parameters - data_sharding=True, - dataloader_type="single", - num_workers=8, - skip_getting_attention_mask_from_dataset=True, - ), - logger=LoggerConfig( - log_interval=10, - tensorboard_dir=tensorboard_dir, - ), - tokenizer=TokenizerConfig(tokenizer_type="HuggingFaceTokenizer", tokenizer_model="EleutherAI/gpt-neox-20b"), - checkpoint=CheckpointConfig( - save_interval=2000, - save=checkpoint_dir, - load=checkpoint_dir, - ckpt_format="torch_dist", - fully_parallel_load=True, - ), - rng=RNGConfig(seed=1234), - comm_overlap=comm_overlap_config, - mixed_precision=precision_config, - ) - - return cfg diff --git a/src/megatron/bridge/recipes/mamba/mamba2_1_3b.py b/src/megatron/bridge/recipes/mamba/mamba2_1_3b.py deleted file mode 100644 index 60e8691865..0000000000 --- a/src/megatron/bridge/recipes/mamba/mamba2_1_3b.py +++ /dev/null @@ -1,214 +0,0 @@ -# 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 Optional, Union - -import torch - -from megatron.bridge.models.mamba import MambaModelProvider1P3B -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.training.comm_overlap import CommOverlapConfig -from megatron.bridge.training.config import ( - CheckpointConfig, - ConfigContainer, - DistributedDataParallelConfig, - GPTDatasetConfig, - LoggerConfig, - RNGConfig, - TokenizerConfig, - TrainingConfig, -) -from megatron.bridge.training.mixed_precision import MixedPrecisionConfig - - -def model_config( - tensor_parallelism: int = 1, - pipeline_parallelism: int = 1, - pipeline_parallelism_dtype: Optional[torch.dtype] = None, - virtual_pipeline_parallelism: Optional[int] = None, - context_parallelism: int = 1, - sequence_parallelism: bool = False, -) -> MambaModelProvider1P3B: - """ - Configure the Mamba 1.3B model. - - Args: - tensor_parallelism: Degree of tensor model parallelism. - pipeline_parallelism: Degree of pipeline model parallelism. - pipeline_parallelism_dtype: Data type for pipeline parallelism. - virtual_pipeline_parallelism: Size of virtual pipeline parallelism. - context_parallelism: Degree of context parallelism. - sequence_parallelism: Whether to use sequence parallelism. - - Returns: - MambaModelProvider1P3B: Configuration for the Mamba 1.3B model. - """ - return MambaModelProvider1P3B( - tensor_model_parallel_size=tensor_parallelism, - pipeline_model_parallel_size=pipeline_parallelism, - pipeline_dtype=pipeline_parallelism_dtype, - virtual_pipeline_model_parallel_size=virtual_pipeline_parallelism, - context_parallel_size=context_parallelism, - sequence_parallel=sequence_parallelism, - ) - - -def pretrain_config( - 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, - # Model configuration - tensor_parallelism: int = 1, - pipeline_parallelism: int = 1, - pipeline_parallelism_dtype: Optional[torch.dtype] = None, - virtual_pipeline_parallelism: Optional[int] = None, - context_parallelism: int = 1, - sequence_parallelism: bool = False, - # Training hyperparameters - train_iters: int = 1_168_251, - global_batch_size: int = 8, - micro_batch_size: int = 1, - seq_length: int = 4096, - lr: float = 3e-4, - min_lr: float = 3e-5, - lr_warmup_iters: int = 2000, - lr_decay_iters: Optional[int] = None, - # Precision recipe - precision_config: Optional[Union[MixedPrecisionConfig, str]] = "bf16_mixed", - comm_overlap_config: Optional[CommOverlapConfig] = None, -) -> ConfigContainer: - """ - Create a pre-training configuration for Mamba 1.3B model. - - Args: - 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. - sequence_parallelism (bool): Whether to use sequence parallelism. - 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 for learning rate decay. - precision_config (Optional[Union[MixedPrecisionConfig, str]]): Precision configuration for the model. - comm_overlap_config (Optional[CommOverlapConfig]): Communication overlap configuration for the model. - - 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 - ) - - model_cfg = model_config( - tensor_parallelism=tensor_parallelism, - pipeline_parallelism=pipeline_parallelism, - pipeline_parallelism_dtype=pipeline_parallelism_dtype, - virtual_pipeline_parallelism=virtual_pipeline_parallelism, - context_parallelism=context_parallelism, - sequence_parallelism=sequence_parallelism, - ) - - opt_config, scheduler = distributed_fused_adam_with_cosine_annealing( - lr_warmup_iters=lr_warmup_iters, - lr_decay_iters=lr_decay_iters, - adam_beta1=0.9, - adam_beta2=0.95, - adam_eps=1e-5, - weight_decay=0.1, - max_lr=lr, - min_lr=min_lr, - ) - - # Config Container - cfg = ConfigContainer( - model=model_cfg, - train=TrainingConfig( - train_iters=train_iters, - eval_interval=100, - eval_iters=32, - global_batch_size=global_batch_size, - micro_batch_size=micro_batch_size, - ), - 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, - 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, - # Dataloader config parameters - data_sharding=True, - dataloader_type="single", - num_workers=8, - skip_getting_attention_mask_from_dataset=True, - ), - logger=LoggerConfig( - log_interval=10, - tensorboard_dir=tensorboard_dir, - ), - tokenizer=TokenizerConfig(tokenizer_type="HuggingFaceTokenizer", tokenizer_model="EleutherAI/gpt-neox-20b"), - checkpoint=CheckpointConfig( - save_interval=2000, - save=checkpoint_dir, - load=checkpoint_dir, - ckpt_format="torch_dist", - fully_parallel_load=True, - ), - rng=RNGConfig(seed=1234), - comm_overlap=comm_overlap_config, - mixed_precision=precision_config, - ) - - return cfg diff --git a/src/megatron/bridge/recipes/mamba/mamba2_2_7b.py b/src/megatron/bridge/recipes/mamba/mamba2_2_7b.py deleted file mode 100644 index 3266cc51f3..0000000000 --- a/src/megatron/bridge/recipes/mamba/mamba2_2_7b.py +++ /dev/null @@ -1,214 +0,0 @@ -# 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 Optional, Union - -import torch - -from megatron.bridge.models.mamba import MambaModelProvider2P7B -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.training.comm_overlap import CommOverlapConfig -from megatron.bridge.training.config import ( - CheckpointConfig, - ConfigContainer, - DistributedDataParallelConfig, - GPTDatasetConfig, - LoggerConfig, - RNGConfig, - TokenizerConfig, - TrainingConfig, -) -from megatron.bridge.training.mixed_precision import MixedPrecisionConfig - - -def model_config( - tensor_parallelism: int = 1, - pipeline_parallelism: int = 1, - pipeline_parallelism_dtype: Optional[torch.dtype] = None, - virtual_pipeline_parallelism: Optional[int] = None, - context_parallelism: int = 1, - sequence_parallelism: bool = False, -) -> MambaModelProvider2P7B: - """ - Configure the Mamba 2.7B model. - - Args: - tensor_parallelism: Degree of tensor model parallelism. - pipeline_parallelism: Degree of pipeline model parallelism. - pipeline_parallelism_dtype: Data type for pipeline parallelism. - virtual_pipeline_parallelism: Size of virtual pipeline parallelism. - context_parallelism: Degree of context parallelism. - sequence_parallelism: Whether to use sequence parallelism. - - Returns: - MambaModelProvider2P7B: Configuration for the Mamba 2.7B model. - """ - return MambaModelProvider2P7B( - tensor_model_parallel_size=tensor_parallelism, - pipeline_model_parallel_size=pipeline_parallelism, - pipeline_dtype=pipeline_parallelism_dtype, - virtual_pipeline_model_parallel_size=virtual_pipeline_parallelism, - context_parallel_size=context_parallelism, - sequence_parallel=sequence_parallelism, - ) - - -def pretrain_config( - 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, - # Model configuration - tensor_parallelism: int = 1, - pipeline_parallelism: int = 1, - pipeline_parallelism_dtype: Optional[torch.dtype] = None, - virtual_pipeline_parallelism: Optional[int] = None, - context_parallelism: int = 1, - sequence_parallelism: bool = False, - # Training hyperparameters - train_iters: int = 1_168_251, - global_batch_size: int = 8, - micro_batch_size: int = 1, - seq_length: int = 4096, - lr: float = 3e-4, - min_lr: float = 3e-5, - lr_warmup_iters: int = 2000, - lr_decay_iters: Optional[int] = None, - # Precision recipe - precision_config: Optional[Union[MixedPrecisionConfig, str]] = "bf16_mixed", - comm_overlap_config: Optional[CommOverlapConfig] = None, -) -> ConfigContainer: - """ - Create a pre-training configuration for Mamba 2.7B model. - - Args: - 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. - sequence_parallelism (bool): Whether to use sequence parallelism. - 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 for learning rate decay. - precision_config (Optional[Union[MixedPrecisionConfig, str]]): Precision configuration for the model. - comm_overlap_config (Optional[CommOverlapConfig]): Communication overlap configuration for the model. - - 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 - ) - - model_cfg = model_config( - tensor_parallelism=tensor_parallelism, - pipeline_parallelism=pipeline_parallelism, - pipeline_parallelism_dtype=pipeline_parallelism_dtype, - virtual_pipeline_parallelism=virtual_pipeline_parallelism, - context_parallelism=context_parallelism, - sequence_parallelism=sequence_parallelism, - ) - - opt_config, scheduler = distributed_fused_adam_with_cosine_annealing( - lr_warmup_iters=lr_warmup_iters, - lr_decay_iters=lr_decay_iters, - adam_beta1=0.9, - adam_beta2=0.95, - adam_eps=1e-5, - weight_decay=0.1, - max_lr=lr, - min_lr=min_lr, - ) - - # Config Container - cfg = ConfigContainer( - model=model_cfg, - train=TrainingConfig( - train_iters=train_iters, - eval_interval=100, - eval_iters=32, - global_batch_size=global_batch_size, - micro_batch_size=micro_batch_size, - ), - 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, - 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, - # Dataloader config parameters - data_sharding=True, - dataloader_type="single", - num_workers=8, - skip_getting_attention_mask_from_dataset=True, - ), - logger=LoggerConfig( - log_interval=10, - tensorboard_dir=tensorboard_dir, - ), - tokenizer=TokenizerConfig(tokenizer_type="HuggingFaceTokenizer", tokenizer_model="EleutherAI/gpt-neox-20b"), - checkpoint=CheckpointConfig( - save_interval=2000, - save=checkpoint_dir, - load=checkpoint_dir, - ckpt_format="torch_dist", - fully_parallel_load=True, - ), - rng=RNGConfig(seed=1234), - comm_overlap=comm_overlap_config, - mixed_precision=precision_config, - ) - - return cfg diff --git a/src/megatron/bridge/recipes/mamba/mamba2_370m.py b/src/megatron/bridge/recipes/mamba/mamba2_370m.py deleted file mode 100644 index fb1b796376..0000000000 --- a/src/megatron/bridge/recipes/mamba/mamba2_370m.py +++ /dev/null @@ -1,214 +0,0 @@ -# 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 Optional, Union - -import torch - -from megatron.bridge.models.mamba import MambaModelProvider370M -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.training.comm_overlap import CommOverlapConfig -from megatron.bridge.training.config import ( - CheckpointConfig, - ConfigContainer, - DistributedDataParallelConfig, - GPTDatasetConfig, - LoggerConfig, - RNGConfig, - TokenizerConfig, - TrainingConfig, -) -from megatron.bridge.training.mixed_precision import MixedPrecisionConfig - - -def model_config( - tensor_parallelism: int = 1, - pipeline_parallelism: int = 1, - pipeline_parallelism_dtype: Optional[torch.dtype] = None, - virtual_pipeline_parallelism: Optional[int] = None, - context_parallelism: int = 1, - sequence_parallelism: bool = False, -) -> MambaModelProvider370M: - """ - Configure the Mamba 370M model. - - Args: - tensor_parallelism: Degree of tensor model parallelism. - pipeline_parallelism: Degree of pipeline model parallelism. - pipeline_parallelism_dtype: Data type for pipeline parallelism. - virtual_pipeline_parallelism: Size of virtual pipeline parallelism. - context_parallelism: Degree of context parallelism. - sequence_parallelism: Whether to use sequence parallelism. - - Returns: - MambaModelProvider370M: Configuration for the Mamba 370M model. - """ - return MambaModelProvider370M( - tensor_model_parallel_size=tensor_parallelism, - pipeline_model_parallel_size=pipeline_parallelism, - pipeline_dtype=pipeline_parallelism_dtype, - virtual_pipeline_model_parallel_size=virtual_pipeline_parallelism, - context_parallel_size=context_parallelism, - sequence_parallel=sequence_parallelism, - ) - - -def pretrain_config( - 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, - # Model configuration - tensor_parallelism: int = 1, - pipeline_parallelism: int = 1, - pipeline_parallelism_dtype: Optional[torch.dtype] = None, - virtual_pipeline_parallelism: Optional[int] = None, - context_parallelism: int = 1, - sequence_parallelism: bool = False, - # Training hyperparameters - train_iters: int = 1_168_251, - global_batch_size: int = 8, - micro_batch_size: int = 1, - seq_length: int = 4096, - lr: float = 3e-4, - min_lr: float = 3e-5, - lr_warmup_iters: int = 2000, - lr_decay_iters: Optional[int] = None, - # Precision recipe - precision_config: Optional[Union[MixedPrecisionConfig, str]] = "bf16_mixed", - comm_overlap_config: Optional[CommOverlapConfig] = None, -) -> ConfigContainer: - """ - Create a pre-training configuration for Mamba 370M model. - - Args: - 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. - sequence_parallelism (bool): Whether to use sequence parallelism. - 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 for learning rate decay. - precision_config (Optional[Union[MixedPrecisionConfig, str]]): Precision configuration for the model. - comm_overlap_config (Optional[CommOverlapConfig]): Communication overlap configuration for the model. - - 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 - ) - - model_cfg = model_config( - tensor_parallelism=tensor_parallelism, - pipeline_parallelism=pipeline_parallelism, - pipeline_parallelism_dtype=pipeline_parallelism_dtype, - virtual_pipeline_parallelism=virtual_pipeline_parallelism, - context_parallelism=context_parallelism, - sequence_parallelism=sequence_parallelism, - ) - - opt_config, scheduler = distributed_fused_adam_with_cosine_annealing( - lr_warmup_iters=lr_warmup_iters, - lr_decay_iters=lr_decay_iters, - adam_beta1=0.9, - adam_beta2=0.95, - adam_eps=1e-5, - weight_decay=0.1, - max_lr=lr, - min_lr=min_lr, - ) - - # Config Container - cfg = ConfigContainer( - model=model_cfg, - train=TrainingConfig( - train_iters=train_iters, - eval_interval=100, - eval_iters=32, - global_batch_size=global_batch_size, - micro_batch_size=micro_batch_size, - ), - 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, - 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, - # Dataloader config parameters - data_sharding=True, - dataloader_type="single", - num_workers=8, - skip_getting_attention_mask_from_dataset=True, - ), - logger=LoggerConfig( - log_interval=10, - tensorboard_dir=tensorboard_dir, - ), - tokenizer=TokenizerConfig(tokenizer_type="HuggingFaceTokenizer", tokenizer_model="EleutherAI/gpt-neox-20b"), - checkpoint=CheckpointConfig( - save_interval=2000, - save=checkpoint_dir, - load=checkpoint_dir, - ckpt_format="torch_dist", - fully_parallel_load=True, - ), - rng=RNGConfig(seed=1234), - comm_overlap=comm_overlap_config, - mixed_precision=precision_config, - ) - - return cfg diff --git a/src/megatron/bridge/recipes/mamba/mamba2_780m.py b/src/megatron/bridge/recipes/mamba/mamba2_780m.py deleted file mode 100644 index 800459be7f..0000000000 --- a/src/megatron/bridge/recipes/mamba/mamba2_780m.py +++ /dev/null @@ -1,214 +0,0 @@ -# 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 Optional, Union - -import torch - -from megatron.bridge.models.mamba import MambaModelProvider780M -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.training.comm_overlap import CommOverlapConfig -from megatron.bridge.training.config import ( - CheckpointConfig, - ConfigContainer, - DistributedDataParallelConfig, - GPTDatasetConfig, - LoggerConfig, - RNGConfig, - TokenizerConfig, - TrainingConfig, -) -from megatron.bridge.training.mixed_precision import MixedPrecisionConfig - - -def model_config( - tensor_parallelism: int = 1, - pipeline_parallelism: int = 1, - pipeline_parallelism_dtype: Optional[torch.dtype] = None, - virtual_pipeline_parallelism: Optional[int] = None, - context_parallelism: int = 1, - sequence_parallelism: bool = False, -) -> MambaModelProvider780M: - """ - Configure the Mamba 780M model. - - Args: - tensor_parallelism: Degree of tensor model parallelism. - pipeline_parallelism: Degree of pipeline model parallelism. - pipeline_parallelism_dtype: Data type for pipeline parallelism. - virtual_pipeline_parallelism: Size of virtual pipeline parallelism. - context_parallelism: Degree of context parallelism. - sequence_parallelism: Whether to use sequence parallelism. - - Returns: - MambaModelProvider780M: Configuration for the Mamba 780M model. - """ - return MambaModelProvider780M( - tensor_model_parallel_size=tensor_parallelism, - pipeline_model_parallel_size=pipeline_parallelism, - pipeline_dtype=pipeline_parallelism_dtype, - virtual_pipeline_model_parallel_size=virtual_pipeline_parallelism, - context_parallel_size=context_parallelism, - sequence_parallel=sequence_parallelism, - ) - - -def pretrain_config( - 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, - # Model configuration - tensor_parallelism: int = 1, - pipeline_parallelism: int = 1, - pipeline_parallelism_dtype: Optional[torch.dtype] = None, - virtual_pipeline_parallelism: Optional[int] = None, - context_parallelism: int = 1, - sequence_parallelism: bool = False, - # Training hyperparameters - train_iters: int = 1_168_251, - global_batch_size: int = 8, - micro_batch_size: int = 1, - seq_length: int = 4096, - lr: float = 3e-4, - min_lr: float = 3e-5, - lr_warmup_iters: int = 2000, - lr_decay_iters: Optional[int] = None, - # Precision recipe - precision_config: Optional[Union[MixedPrecisionConfig, str]] = "bf16_mixed", - comm_overlap_config: Optional[CommOverlapConfig] = None, -) -> ConfigContainer: - """ - Create a pre-training configuration for Mamba 780M model. - - Args: - 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. - sequence_parallelism (bool): Whether to use sequence parallelism. - 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 for learning rate decay. - precision_config (Optional[Union[MixedPrecisionConfig, str]]): Precision configuration for the model. - comm_overlap_config (Optional[CommOverlapConfig]): Communication overlap configuration for the model. - - 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 - ) - - model_cfg = model_config( - tensor_parallelism=tensor_parallelism, - pipeline_parallelism=pipeline_parallelism, - pipeline_parallelism_dtype=pipeline_parallelism_dtype, - virtual_pipeline_parallelism=virtual_pipeline_parallelism, - context_parallelism=context_parallelism, - sequence_parallelism=sequence_parallelism, - ) - - opt_config, scheduler = distributed_fused_adam_with_cosine_annealing( - lr_warmup_iters=lr_warmup_iters, - lr_decay_iters=lr_decay_iters, - adam_beta1=0.9, - adam_beta2=0.95, - adam_eps=1e-5, - weight_decay=0.1, - max_lr=lr, - min_lr=min_lr, - ) - - # Config Container - cfg = ConfigContainer( - model=model_cfg, - train=TrainingConfig( - train_iters=train_iters, - eval_interval=100, - eval_iters=32, - global_batch_size=global_batch_size, - micro_batch_size=micro_batch_size, - ), - 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, - 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, - # Dataloader config parameters - data_sharding=True, - dataloader_type="single", - num_workers=8, - skip_getting_attention_mask_from_dataset=True, - ), - logger=LoggerConfig( - log_interval=10, - tensorboard_dir=tensorboard_dir, - ), - tokenizer=TokenizerConfig(tokenizer_type="HuggingFaceTokenizer", tokenizer_model="EleutherAI/gpt-neox-20b"), - checkpoint=CheckpointConfig( - save_interval=2000, - save=checkpoint_dir, - load=checkpoint_dir, - ckpt_format="torch_dist", - fully_parallel_load=True, - ), - rng=RNGConfig(seed=1234), - comm_overlap=comm_overlap_config, - mixed_precision=precision_config, - ) - - return cfg diff --git a/src/megatron/bridge/recipes/mamba/mamba2_8b.py b/src/megatron/bridge/recipes/mamba/mamba2_8b.py deleted file mode 100644 index d263ec3551..0000000000 --- a/src/megatron/bridge/recipes/mamba/mamba2_8b.py +++ /dev/null @@ -1,215 +0,0 @@ -# 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 Optional, Union - -import torch - -from megatron.bridge.models.mamba import NVIDIAMambaModelProvider8B -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, - DistributedDataParallelConfig, - GPTDatasetConfig, - LoggerConfig, - RNGConfig, - TokenizerConfig, - TrainingConfig, -) -from megatron.bridge.training.mixed_precision import MixedPrecisionConfig - - -def model_config( - tensor_parallelism: int = 8, - pipeline_parallelism: int = 1, - pipeline_parallelism_dtype: Optional[torch.dtype] = None, - virtual_pipeline_parallelism: Optional[int] = None, - context_parallelism: int = 1, - sequence_parallelism: bool = False, -) -> NVIDIAMambaModelProvider8B: - """ - Configure the Mamba 8B model. - - Args: - tensor_parallelism: Degree of tensor model parallelism. - pipeline_parallelism: Degree of pipeline model parallelism. - pipeline_parallelism_dtype: Data type for pipeline parallelism. - virtual_pipeline_parallelism: Size of virtual pipeline parallelism. - context_parallelism: Degree of context parallelism. - sequence_parallelism: Whether to use sequence parallelism. - - Returns: - NVIDIAMambaModelProvider8B: Configuration for the Mamba 8B model. - """ - return NVIDIAMambaModelProvider8B( - tensor_model_parallel_size=tensor_parallelism, - pipeline_model_parallel_size=pipeline_parallelism, - pipeline_dtype=pipeline_parallelism_dtype, - virtual_pipeline_model_parallel_size=virtual_pipeline_parallelism, - context_parallel_size=context_parallelism, - sequence_parallel=sequence_parallelism, - ) - - -def pretrain_config( - 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, - # Model configuration - tensor_parallelism: int = 8, - pipeline_parallelism: int = 1, - pipeline_parallelism_dtype: Optional[torch.dtype] = None, - virtual_pipeline_parallelism: Optional[int] = None, - context_parallelism: int = 1, - sequence_parallelism: bool = False, - # Training hyperparameters - train_iters: int = 1_168_251, - global_batch_size: int = 8, - micro_batch_size: int = 1, - seq_length: int = 4096, - lr: float = 3e-4, - min_lr: float = 3e-5, - lr_warmup_iters: int = 2000, - lr_decay_iters: Optional[int] = None, - # Precision recipe - precision_config: Optional[Union[MixedPrecisionConfig, str]] = "bf16_mixed", - comm_overlap_config: Optional[CommOverlapConfig] = None, -) -> ConfigContainer: - """ - Create a pre-training configuration for Mamba 8B model. - - Args: - 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. - sequence_parallelism (bool): Whether to use sequence parallelism. - 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 for learning rate decay. - precision_config (Optional[Union[MixedPrecisionConfig, str]]): Precision configuration for the model. - comm_overlap_config (Optional[CommOverlapConfig]): Communication overlap configuration for the model. - - 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 - ) - - model_cfg = model_config( - tensor_parallelism=tensor_parallelism, - pipeline_parallelism=pipeline_parallelism, - pipeline_parallelism_dtype=pipeline_parallelism_dtype, - virtual_pipeline_parallelism=virtual_pipeline_parallelism, - context_parallelism=context_parallelism, - sequence_parallelism=sequence_parallelism, - ) - - opt_config, scheduler = distributed_fused_adam_with_cosine_annealing( - lr_warmup_iters=lr_warmup_iters, - lr_decay_iters=lr_decay_iters, - adam_beta1=0.9, - adam_beta2=0.95, - adam_eps=1e-5, - weight_decay=0.1, - max_lr=lr, - min_lr=min_lr, - ) - - # Config Container - cfg = ConfigContainer( - model=model_cfg, - train=TrainingConfig( - train_iters=train_iters, - eval_interval=100, - eval_iters=32, - global_batch_size=global_batch_size, - micro_batch_size=micro_batch_size, - ), - 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, - 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, - # Dataloader config parameters - data_sharding=True, - dataloader_type="single", - num_workers=8, - skip_getting_attention_mask_from_dataset=True, - ), - logger=LoggerConfig( - log_interval=10, - tensorboard_dir=tensorboard_dir, - ), - tokenizer=TokenizerConfig(tokenizer_type="NullTokenizer", vocab_size=DEFAULT_NULL_TOKENIZER_VOCAB_SIZE), - checkpoint=CheckpointConfig( - save_interval=2000, - save=checkpoint_dir, - load=checkpoint_dir, - ckpt_format="torch_dist", - fully_parallel_load=True, - ), - rng=RNGConfig(seed=1234), - comm_overlap=comm_overlap_config, - mixed_precision=precision_config, - ) - - return cfg diff --git a/src/megatron/bridge/recipes/mamba/mamba2_hybrid_8b.py b/src/megatron/bridge/recipes/mamba/mamba2_hybrid_8b.py deleted file mode 100644 index 6a2fb6b78e..0000000000 --- a/src/megatron/bridge/recipes/mamba/mamba2_hybrid_8b.py +++ /dev/null @@ -1,215 +0,0 @@ -# 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 Optional, Union - -import torch - -from megatron.bridge.models.mamba import NVIDIAMambaHybridProvider8B -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, - DistributedDataParallelConfig, - GPTDatasetConfig, - LoggerConfig, - RNGConfig, - TokenizerConfig, - TrainingConfig, -) -from megatron.bridge.training.mixed_precision import MixedPrecisionConfig - - -def model_config( - tensor_parallelism: int = 8, - pipeline_parallelism: int = 1, - pipeline_parallelism_dtype: Optional[torch.dtype] = None, - virtual_pipeline_parallelism: Optional[int] = None, - context_parallelism: int = 1, - sequence_parallelism: bool = False, -) -> NVIDIAMambaHybridProvider8B: - """ - Configure the Mamba Hybrid 8B model. - - Args: - tensor_parallelism: Degree of tensor model parallelism. - pipeline_parallelism: Degree of pipeline model parallelism. - pipeline_parallelism_dtype: Data type for pipeline parallelism. - virtual_pipeline_parallelism: Size of virtual pipeline parallelism. - context_parallelism: Degree of context parallelism. - sequence_parallelism: Whether to use sequence parallelism. - - Returns: - NVIDIAMambaHybridProvider8B: Configuration for the Mamba Hybrid 8B model. - """ - return NVIDIAMambaHybridProvider8B( - tensor_model_parallel_size=tensor_parallelism, - pipeline_model_parallel_size=pipeline_parallelism, - pipeline_dtype=pipeline_parallelism_dtype, - virtual_pipeline_model_parallel_size=virtual_pipeline_parallelism, - context_parallel_size=context_parallelism, - sequence_parallel=sequence_parallelism, - ) - - -def pretrain_config( - 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, - # Model configuration - tensor_parallelism: int = 8, - pipeline_parallelism: int = 1, - pipeline_parallelism_dtype: Optional[torch.dtype] = None, - virtual_pipeline_parallelism: Optional[int] = None, - context_parallelism: int = 1, - sequence_parallelism: bool = False, - # Training hyperparameters - train_iters: int = 1_168_251, - global_batch_size: int = 8, - micro_batch_size: int = 1, - seq_length: int = 4096, - lr: float = 3e-4, - min_lr: float = 3e-5, - lr_warmup_iters: int = 2000, - lr_decay_iters: Optional[int] = None, - # Precision recipe - precision_config: Optional[Union[MixedPrecisionConfig, str]] = "bf16_mixed", - comm_overlap_config: Optional[CommOverlapConfig] = None, -) -> ConfigContainer: - """ - Create a pre-training configuration for Mamba Hybrid 8B model. - - Args: - 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. - sequence_parallelism (bool): Whether to use sequence parallelism. - 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 for learning rate decay. - precision_config (Optional[Union[MixedPrecisionConfig, str]]): Precision configuration for the model. - comm_overlap_config (Optional[CommOverlapConfig]): Communication overlap configuration for the model. - - 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 - ) - - model_cfg = model_config( - tensor_parallelism=tensor_parallelism, - pipeline_parallelism=pipeline_parallelism, - pipeline_parallelism_dtype=pipeline_parallelism_dtype, - virtual_pipeline_parallelism=virtual_pipeline_parallelism, - context_parallelism=context_parallelism, - sequence_parallelism=sequence_parallelism, - ) - - opt_config, scheduler = distributed_fused_adam_with_cosine_annealing( - lr_warmup_iters=lr_warmup_iters, - lr_decay_iters=lr_decay_iters, - adam_beta1=0.9, - adam_beta2=0.95, - adam_eps=1e-5, - weight_decay=0.1, - max_lr=lr, - min_lr=min_lr, - ) - - # Config Container - cfg = ConfigContainer( - model=model_cfg, - train=TrainingConfig( - train_iters=train_iters, - eval_interval=100, - eval_iters=32, - global_batch_size=global_batch_size, - micro_batch_size=micro_batch_size, - ), - 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, - 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, - # Dataloader config parameters - data_sharding=True, - dataloader_type="single", - num_workers=8, - skip_getting_attention_mask_from_dataset=True, - ), - logger=LoggerConfig( - log_interval=10, - tensorboard_dir=tensorboard_dir, - ), - tokenizer=TokenizerConfig(tokenizer_type="NullTokenizer", vocab_size=DEFAULT_NULL_TOKENIZER_VOCAB_SIZE), - checkpoint=CheckpointConfig( - save_interval=2000, - save=checkpoint_dir, - load=checkpoint_dir, - ckpt_format="torch_dist", - fully_parallel_load=True, - ), - rng=RNGConfig(seed=1234), - comm_overlap=comm_overlap_config, - mixed_precision=precision_config, - ) - - return cfg diff --git a/src/megatron/bridge/recipes/qwen_vl/__init__.py b/src/megatron/bridge/recipes/qwen_vl/__init__.py new file mode 100644 index 0000000000..341a77c5bc --- /dev/null +++ b/src/megatron/bridge/recipes/qwen_vl/__init__.py @@ -0,0 +1,13 @@ +# 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. diff --git a/tests/functional_tests/recipes/test_mamba_recipes_pretrain.py b/tests/functional_tests/recipes/test_mamba_recipes_pretrain.py index baa27d8af4..11b24a7537 100644 --- a/tests/functional_tests/recipes/test_mamba_recipes_pretrain.py +++ b/tests/functional_tests/recipes/test_mamba_recipes_pretrain.py @@ -16,9 +16,15 @@ import pytest -from megatron.bridge.recipes.mamba.mamba2_130m import pretrain_config as mamba2_130m_config -from megatron.bridge.recipes.mamba.mamba2_370m import pretrain_config as mamba2_370m_config -from megatron.bridge.recipes.mamba.mamba2_780m import pretrain_config as mamba2_780m_config +from megatron.bridge.recipes.mamba import ( + mamba2_130m_pretrain_config as mamba2_130m_config, +) +from megatron.bridge.recipes.mamba import ( + mamba2_370m_pretrain_config as mamba2_370m_config, +) +from megatron.bridge.recipes.mamba import ( + mamba2_780m_pretrain_config as mamba2_780m_config, +) from tests.functional_tests.recipes.utils import run_pretrain_config_override_test, run_pretrain_recipe_test diff --git a/tests/functional_tests/recipes/utils.py b/tests/functional_tests/recipes/utils.py index 8f69467645..32d340574f 100644 --- a/tests/functional_tests/recipes/utils.py +++ b/tests/functional_tests/recipes/utils.py @@ -94,18 +94,12 @@ def run_pretrain_recipe_test( if tensor_parallelism is not None: if hasattr(config.model, "tensor_model_parallel_size"): config.model.tensor_model_parallel_size = tensor_parallelism - else: - setattr(config.model, "tensor_parallelism", tensor_parallelism) if pipeline_parallelism is not None: if hasattr(config.model, "pipeline_model_parallel_size"): config.model.pipeline_model_parallel_size = pipeline_parallelism - else: - setattr(config.model, "pipeline_parallelism", pipeline_parallelism) if expert_parallelism is not None: if hasattr(config.model, "expert_model_parallel_size"): config.model.expert_model_parallel_size = expert_parallelism - else: - setattr(config.model, "expert_parallelism", expert_parallelism) # Apply any model-specific overrides provided by the caller if model_overrides: diff --git a/tests/unit_tests/recipes/mamba/test_mamba2.py b/tests/unit_tests/recipes/mamba/test_mamba2.py new file mode 100644 index 0000000000..8e1fb4356c --- /dev/null +++ b/tests/unit_tests/recipes/mamba/test_mamba2.py @@ -0,0 +1,305 @@ +# 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 +import tempfile +from unittest.mock import patch + +import pytest +import torch + +from megatron.bridge.models.mamba import ( + MambaModelProvider1P3B, + MambaModelProvider2P7B, + MambaModelProvider130M, + MambaModelProvider370M, + MambaModelProvider780M, + NVIDIAMambaHybridProvider8B, + NVIDIAMambaModelProvider8B, +) +from megatron.bridge.recipes.mamba import ( + mamba2_1p3b_pretrain_config, + mamba2_2p7b_pretrain_config, + mamba2_8b_pretrain_config, + mamba2_130m_pretrain_config, + mamba2_370m_pretrain_config, + mamba2_780m_pretrain_config, + mamba2_hybrid_8b_pretrain_config, +) +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 ConfigContainer + + +RECIPE_CASES = [ + ("mamba2_130m", mamba2_130m_pretrain_config, MambaModelProvider130M, False), + ("mamba2_370m", mamba2_370m_pretrain_config, MambaModelProvider370M, False), + ("mamba2_780m", mamba2_780m_pretrain_config, MambaModelProvider780M, False), + ("mamba2_1p3b", mamba2_1p3b_pretrain_config, MambaModelProvider1P3B, False), + ("mamba2_2p7b", mamba2_2p7b_pretrain_config, MambaModelProvider2P7B, False), + ("mamba2_8b", mamba2_8b_pretrain_config, NVIDIAMambaModelProvider8B, True), + ("mamba2_hybrid_8b", mamba2_hybrid_8b_pretrain_config, NVIDIAMambaHybridProvider8B, True), +] + +SMALL_PAR_COMBOS = [ + (1, 1, 1), + (1, 4, 2), + (1, 2, 4), + (1, 2, 2), + (1, 4, 1), +] + +LARGE_PAR_COMBOS = [ + (1, 1, 1), + (2, 1, 4), + (4, 2, 2), + (8, 2, 2), +] + + +class TestMamba2Recipes: + @pytest.mark.unit + @pytest.mark.parametrize("name,config_func,provider,uses_null_tokenizer", RECIPE_CASES) + def test_pretrain_config_default_parameters(self, name, config_func, provider, uses_null_tokenizer): + cfg = config_func() + + assert isinstance(cfg, ConfigContainer) + assert isinstance(cfg.model, provider) + + # Training configuration + assert cfg.train.train_iters == 1_168_251 + assert cfg.train.global_batch_size == 8 + assert cfg.train.micro_batch_size == 1 + assert cfg.train.eval_interval == 100 + assert cfg.train.eval_iters == 32 + + # Optimizer + assert cfg.optimizer.optimizer == "adam" + assert cfg.optimizer.lr == 3e-4 + assert cfg.optimizer.min_lr == 3e-5 + assert cfg.optimizer.weight_decay == 0.1 + assert cfg.optimizer.bf16 is True + assert cfg.optimizer.fp16 is False + + # Dataset in mock mode by default + assert cfg.dataset.sequence_length == 4096 + assert cfg.dataset.split == "1,1,1" + assert cfg.dataset.blend is None + assert cfg.dataset.blend_per_split is None + + # Tokenizer + if uses_null_tokenizer: + assert cfg.tokenizer.tokenizer_type == "NullTokenizer" + assert cfg.tokenizer.vocab_size == DEFAULT_NULL_TOKENIZER_VOCAB_SIZE + else: + assert cfg.tokenizer.tokenizer_type == "HuggingFaceTokenizer" + assert cfg.tokenizer.tokenizer_model == "EleutherAI/gpt-neox-20b" + + @pytest.mark.unit + @pytest.mark.parametrize("name,config_func,provider,uses_null_tokenizer", RECIPE_CASES) + def test_pretrain_config_custom_training_parameters(self, name, config_func, provider, uses_null_tokenizer): + cfg = config_func( + train_iters=10000, + global_batch_size=256, + micro_batch_size=2, + seq_length=4096, + lr=1e-4, + min_lr=1e-5, + lr_warmup_iters=1000, + ) + + assert cfg.train.train_iters == 10000 + assert cfg.train.global_batch_size == 256 + assert cfg.train.micro_batch_size == 2 + assert cfg.dataset.sequence_length == 4096 + assert cfg.optimizer.lr == 1e-4 + assert cfg.optimizer.min_lr == 1e-5 + assert cfg.scheduler.lr_warmup_iters == 1000 + assert cfg.scheduler.lr_decay_iters is None + + @pytest.mark.unit + @pytest.mark.parametrize("name,config_func,provider,uses_null_tokenizer", RECIPE_CASES) + def test_pretrain_config_with_custom_directory(self, name, config_func, provider, uses_null_tokenizer): + with tempfile.TemporaryDirectory() as temp_dir: + cfg = config_func(dir=temp_dir, name=f"{name}_run") + + expected_run_dir = os.path.join(temp_dir, f"{name}_run") + expected_checkpoint_dir = os.path.join(expected_run_dir, "checkpoints") + expected_tensorboard_dir = os.path.join(expected_run_dir, "tb_logs") + + assert cfg.checkpoint.save == expected_checkpoint_dir + assert cfg.logger.tensorboard_dir == expected_tensorboard_dir + + @pytest.mark.unit + @pytest.mark.parametrize("name,config_func,provider,uses_null_tokenizer", RECIPE_CASES) + def test_pretrain_config_mock_toggle(self, name, config_func, provider, uses_null_tokenizer): + cfg = config_func(mock=True) + assert cfg.dataset.blend is None + assert cfg.dataset.blend_per_split is None + assert cfg.dataset.split == "1,1,1" + + @pytest.mark.unit + @pytest.mark.parametrize("name,config_func,provider,uses_null_tokenizer", RECIPE_CASES) + def test_pretrain_config_with_data_paths(self, name, config_func, provider, uses_null_tokenizer): + data_paths = ["/path/to/data1", "/path/to/data2", "/path/to/data3"] + cfg = config_func(data_paths=data_paths) + assert cfg.dataset.split == "9999,8,2" + assert cfg.dataset.blend is not None + assert cfg.dataset.blend_per_split is None + + @pytest.mark.unit + @pytest.mark.parametrize("name,config_func,provider,uses_null_tokenizer", RECIPE_CASES) + def test_pretrain_config_with_train_valid_test_paths(self, name, config_func, provider, uses_null_tokenizer): + cfg = config_func( + train_data_path=["/path/to/train1", "/path/to/train2", "/path/to/train3"], + valid_data_path=["/path/to/valid1", "/path/to/valid2", "/path/to/valid3"], + test_data_path=["/path/to/test1", "/path/to/test2", "/path/to/test3"], + ) + assert cfg.dataset.split is None + assert cfg.dataset.blend is None + assert cfg.dataset.blend_per_split is not None + + @pytest.mark.unit + @pytest.mark.parametrize("name,config_func,provider,uses_null_tokenizer", RECIPE_CASES) + def test_pretrain_config_prioritizes_blend(self, name, config_func, provider, uses_null_tokenizer): + cfg = config_func( + train_data_path=["/path/to/train1", "/path/to/train2"], + valid_data_path=["/path/to/valid1", "/path/to/valid2"], + test_data_path=["/path/to/test1", "/path/to/test2"], + data_paths=["/path/to/data1", "/path/to/data2"], + ) + assert cfg.dataset.split == "9999,8,2" + assert cfg.dataset.blend is not None + assert cfg.dataset.blend_per_split is None + + @pytest.mark.unit + @pytest.mark.parametrize("name,config_func,provider,uses_null_tokenizer", RECIPE_CASES) + def test_pretrain_config_checkpoint_configuration(self, name, config_func, provider, uses_null_tokenizer): + cfg = config_func() + assert cfg.checkpoint.save_interval == 2000 + assert cfg.checkpoint.ckpt_format == "torch_dist" + assert cfg.checkpoint.fully_parallel_load is True + + @pytest.mark.unit + @pytest.mark.parametrize("name,config_func,provider,uses_null_tokenizer", RECIPE_CASES) + def test_pretrain_config_ddp_configuration(self, name, config_func, provider, uses_null_tokenizer): + cfg = config_func() + assert cfg.ddp.check_for_nan_in_grad is True + assert cfg.ddp.grad_reduce_in_fp32 is True + assert cfg.ddp.overlap_grad_reduce is True + assert cfg.ddp.overlap_param_gather is True + assert cfg.ddp.use_distributed_optimizer is True + + @pytest.mark.unit + @pytest.mark.parametrize("name,config_func,provider,uses_null_tokenizer", RECIPE_CASES) + def test_pretrain_config_comm_overlap(self, name, config_func, provider, uses_null_tokenizer): + cfg = config_func() + assert cfg.comm_overlap is None + + custom_overlap = CommOverlapConfig( + tp_comm_overlap=True, + defer_embedding_wgrad_compute=True, + wgrad_deferral_limit=50, + data_parallel_size=1, + ) + cfg2 = config_func(comm_overlap_config=custom_overlap) + assert cfg2.comm_overlap is not None + assert cfg2.comm_overlap.tp_comm_overlap is True + + @pytest.mark.unit + @pytest.mark.parametrize("name,config_func,provider,uses_null_tokenizer", RECIPE_CASES) + def test_pretrain_config_scheduler_configuration(self, name, config_func, provider, uses_null_tokenizer): + cfg = config_func(train_iters=50000) + 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_iters == 2000 + assert cfg.scheduler.lr_warmup_init == 0.0 + assert cfg.scheduler.lr_decay_iters is None + assert cfg.scheduler.override_opt_param_scheduler is True + + @pytest.mark.unit + @pytest.mark.parametrize("name,config_func,provider,uses_null_tokenizer", RECIPE_CASES) + def test_pretrain_config_rng_and_dataset(self, name, config_func, provider, uses_null_tokenizer): + cfg = config_func() + assert cfg.rng.seed == 1234 + assert cfg.dataset.random_seed == 1234 + assert cfg.dataset.reset_attention_mask is False + assert cfg.dataset.reset_position_ids is False + assert cfg.dataset.eod_mask_loss is False + assert cfg.dataset.num_dataset_builder_threads == 1 + assert cfg.dataset.data_sharding is True + assert cfg.dataset.dataloader_type == "single" + assert cfg.dataset.num_workers == 8 + + @pytest.mark.unit + @pytest.mark.parametrize("name,config_func,provider,uses_null_tokenizer", RECIPE_CASES) + def test_pretrain_config_logger(self, name, config_func, provider, uses_null_tokenizer): + cfg = config_func() + assert cfg.logger.log_interval == 10 + assert "tb_logs" in cfg.logger.tensorboard_dir + + @pytest.mark.unit + @pytest.mark.parametrize("name,config_func,provider,uses_null_tokenizer", RECIPE_CASES) + def test_pretrain_config_parallelism_combinations(self, name, config_func, provider, uses_null_tokenizer): + is_large = name in {"mamba2_8b", "mamba2_hybrid_8b"} + combos = LARGE_PAR_COMBOS if is_large else SMALL_PAR_COMBOS + for tp, pp, cp in combos: + cfg = config_func( + tensor_parallelism=tp, + pipeline_parallelism=pp, + context_parallelism=cp, + pipeline_parallelism_dtype=torch.bfloat16, + sequence_parallelism=(is_large and tp > 1), + ) + assert cfg.model.tensor_model_parallel_size == tp + assert cfg.model.pipeline_model_parallel_size == pp + assert cfg.model.context_parallel_size == cp + + @pytest.mark.unit + @pytest.mark.parametrize("name,config_func,provider,uses_null_tokenizer", RECIPE_CASES) + @pytest.mark.parametrize("global_batch_size,micro_batch_size", [(8, 1), (16, 2), (32, 4), (64, 8)]) + def test_pretrain_config_batch_sizes( + self, name, config_func, provider, uses_null_tokenizer, global_batch_size, micro_batch_size + ): + cfg = config_func(global_batch_size=global_batch_size, micro_batch_size=micro_batch_size) + assert cfg.train.global_batch_size == global_batch_size + assert cfg.train.micro_batch_size == micro_batch_size + + @pytest.mark.unit + @pytest.mark.parametrize("name,config_func,provider,uses_null_tokenizer", RECIPE_CASES) + @pytest.mark.parametrize("seq_length", [1024, 2048, 4096, 8192, 16384]) + def test_pretrain_config_sequence_lengths(self, name, config_func, provider, uses_null_tokenizer, seq_length): + cfg = config_func(seq_length=seq_length) + assert cfg.dataset.sequence_length == seq_length + + @pytest.mark.unit + @pytest.mark.parametrize("name,config_func,provider,uses_null_tokenizer", RECIPE_CASES) + @pytest.mark.parametrize("precision", ["fp16_mixed", "bf16_mixed"]) + def test_precision_recipes(self, name, config_func, provider, uses_null_tokenizer, precision): + cfg = config_func(mixed_precision=precision) if False else config_func(precision_config=precision) + assert cfg.mixed_precision == precision + + @pytest.mark.unit + @pytest.mark.parametrize("name,config_func,provider,uses_null_tokenizer", RECIPE_CASES) + @patch("megatron.bridge.recipes.utils.dataset_utils.get_blend_and_blend_per_split") + def test_pretrain_config_fallback_to_mock_when_no_weights( + self, mock_get_blend, name, config_func, provider, uses_null_tokenizer + ): + mock_get_blend.return_value = (None, None) + cfg = config_func(data_paths=["/some/path"]) + assert cfg.dataset.blend is None + assert cfg.dataset.blend_per_split is None + assert cfg.dataset.split == "1,1,1" diff --git a/tests/unit_tests/recipes/mamba/test_mamba2_130m.py b/tests/unit_tests/recipes/mamba/test_mamba2_130m.py deleted file mode 100644 index 2c4eec093a..0000000000 --- a/tests/unit_tests/recipes/mamba/test_mamba2_130m.py +++ /dev/null @@ -1,362 +0,0 @@ -# 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 -import tempfile -from unittest.mock import patch - -import pytest -import torch - -from megatron.bridge.models.mamba import MambaModelProvider130M -from megatron.bridge.recipes.mamba.mamba2_130m import model_config, pretrain_config -from megatron.bridge.training.comm_overlap import CommOverlapConfig -from megatron.bridge.training.config import ConfigContainer - - -@pytest.mark.unit -class TestModelConfig: - """Test cases for the model_config function.""" - - def test_model_config_default_parameters(self): - """Test model_config with default parameters.""" - config = model_config() - - assert isinstance(config, MambaModelProvider130M) - assert config.tensor_model_parallel_size == 1 - assert config.pipeline_model_parallel_size == 1 - assert config.pipeline_dtype is None - assert config.virtual_pipeline_model_parallel_size is None - assert config.context_parallel_size == 1 - assert config.sequence_parallel is False - - def test_model_config_custom_pipeline_parallelism(self): - """Test model_config with custom pipeline parallelism.""" - config = model_config(pipeline_parallelism=8, pipeline_parallelism_dtype=torch.float16) - - assert config.tensor_model_parallel_size == 1 # default - assert config.pipeline_model_parallel_size == 8 - assert config.pipeline_dtype is torch.float16 - - def test_model_config_with_pipeline_dtype(self): - """Test model_config with pipeline dtype specified.""" - config = model_config(pipeline_parallelism=2, pipeline_parallelism_dtype=torch.float16) - - assert config.pipeline_model_parallel_size == 2 - assert config.pipeline_dtype == torch.float16 - - def test_model_config_virtual_pipeline_parallelism(self): - """Test model_config with virtual pipeline parallelism.""" - config = model_config(virtual_pipeline_parallelism=4) - - assert config.virtual_pipeline_model_parallel_size == 4 - - def test_model_config_context_parallelism(self): - """Test model_config with custom context parallelism.""" - config = model_config(context_parallelism=8) - - assert config.context_parallel_size == 8 - - def test_model_config_all_custom_parameters(self): - """Test model_config with all parameters customized.""" - config = model_config( - tensor_parallelism=1, # Mamba 130M has only 1 attention head - pipeline_parallelism=4, - pipeline_parallelism_dtype=torch.bfloat16, - virtual_pipeline_parallelism=8, - context_parallelism=16, - ) - - assert config.tensor_model_parallel_size == 1 - assert config.pipeline_model_parallel_size == 4 - assert config.pipeline_dtype == torch.bfloat16 - assert config.virtual_pipeline_model_parallel_size == 8 - assert config.context_parallel_size == 16 - - -@pytest.mark.unit -class TestPretrainConfig: - """Test cases for the pretrain_config function.""" - - def test_pretrain_config_default_parameters(self): - """Test pretrain_config with default parameters (mock mode).""" - config = pretrain_config() - - assert isinstance(config, ConfigContainer) - assert isinstance(config.model, MambaModelProvider130M) - - # Check training configuration - assert config.train.train_iters == 1_168_251 - assert config.train.global_batch_size == 8 - assert config.train.micro_batch_size == 1 - assert config.train.eval_interval == 100 - assert config.train.eval_iters == 32 - - # Check optimizer configuration - assert config.optimizer.optimizer == "adam" - assert config.optimizer.lr == 3e-4 - assert config.optimizer.min_lr == 3e-5 - assert config.optimizer.weight_decay == 0.1 - assert config.optimizer.bf16 is True - assert config.optimizer.fp16 is False - - # Check dataset configuration (should be in mock mode) - assert config.dataset.sequence_length == 4096 - assert config.dataset.split == "1,1,1" - assert config.dataset.blend is None - assert config.dataset.blend_per_split is None - - def test_pretrain_config_custom_training_parameters(self): - """Test pretrain_config with custom training parameters.""" - config = pretrain_config( - train_iters=10000, - global_batch_size=256, - micro_batch_size=2, - seq_length=4096, - lr=1e-4, - min_lr=1e-5, - lr_warmup_iters=1000, - ) - - assert config.train.train_iters == 10000 - assert config.train.global_batch_size == 256 - assert config.train.micro_batch_size == 2 - assert config.dataset.sequence_length == 4096 - assert config.optimizer.lr == 1e-4 - assert config.optimizer.min_lr == 1e-5 - assert config.scheduler.lr_warmup_iters == 1000 - assert config.scheduler.lr_decay_iters is None # Will be set to train_iters during validation - - def test_pretrain_config_custom_model_parameters(self): - """Test pretrain_config with custom model parameters.""" - config = pretrain_config( - tensor_parallelism=1, # Mamba 130M has only 1 attention head - pipeline_parallelism=2, - context_parallelism=8, - pipeline_parallelism_dtype=torch.bfloat16, - ) - - assert config.model.tensor_model_parallel_size == 1 - assert config.model.pipeline_model_parallel_size == 2 - assert config.model.context_parallel_size == 8 - assert config.model.pipeline_dtype == torch.bfloat16 - - def test_pretrain_config_with_custom_directory(self): - """Test pretrain_config with custom directory.""" - with tempfile.TemporaryDirectory() as temp_dir: - config = pretrain_config(dir=temp_dir, name="test_run") - - expected_run_dir = os.path.join(temp_dir, "test_run") - expected_checkpoint_dir = os.path.join(expected_run_dir, "checkpoints") - expected_tensorboard_dir = os.path.join(expected_run_dir, "tb_logs") - - assert config.checkpoint.save == expected_checkpoint_dir - assert config.logger.tensorboard_dir == expected_tensorboard_dir - - def test_pretrain_config_explicit_mock_mode(self): - """Test pretrain_config with explicit mock=True.""" - config = pretrain_config(mock=True) - - assert config.dataset.blend is None - assert config.dataset.blend_per_split is None - assert config.dataset.split == "1,1,1" - - def test_pretrain_config_with_data_paths(self): - """Test pretrain_config with data paths provided.""" - data_paths = ["/path/to/data1", "/path/to/data2", "/path/to/data3"] - config = pretrain_config(data_paths=data_paths) - - # Check that non-mock mode is configured - assert config.dataset.split == "9999,8,2" - assert config.dataset.blend is not None - assert config.dataset.blend_per_split is None - - def test_pretrain_config_with_train_valid_test_paths(self): - """Test pretrain_config with separate train/valid/test paths.""" - config = pretrain_config( - train_data_path=["/path/to/train1", "/path/to/train2", "/path/to/train3"], - valid_data_path=["/path/to/valid1", "/path/to/valid2", "/path/to/valid3"], - test_data_path=["/path/to/test1", "/path/to/test2", "/path/to/test3"], - ) - - # When blend_per_split is used, split should be None - assert config.dataset.split is None - assert config.dataset.blend is None - assert config.dataset.blend_per_split is not None - - def test_pretrain_config_prioritizes_blend(self): - """Test that blend takes priority over blend_per_split when both are provided.""" - config = pretrain_config( - train_data_path=["/path/to/train1", "/path/to/train2"], - valid_data_path=["/path/to/valid1", "/path/to/valid2"], - test_data_path=["/path/to/test1", "/path/to/test2"], - data_paths=["/path/to/data1", "/path/to/data2"], - ) - - # Should prioritize blend over blend_per_split - assert config.dataset.split == "9999,8,2" - assert config.dataset.blend is not None - assert config.dataset.blend_per_split is None - - @patch("megatron.bridge.recipes.utils.dataset_utils.get_blend_and_blend_per_split") - def test_pretrain_config_fallback_to_mock_when_no_weights(self, mock_get_blend): - """Test pretrain_config falls back to mock when no weights are returned.""" - # Mock function returns None for both weights - mock_get_blend.return_value = (None, None) - - config = pretrain_config(data_paths=["/some/path"]) - - # Should fall back to mock mode - assert config.dataset.blend is None - assert config.dataset.blend_per_split is None - assert config.dataset.split == "1,1,1" - - def test_pretrain_config_checkpoint_configuration(self): - """Test checkpoint configuration in pretrain_config.""" - config = pretrain_config() - - assert config.checkpoint.save_interval == 2000 - assert config.checkpoint.ckpt_format == "torch_dist" - assert config.checkpoint.fully_parallel_load is True - - def test_pretrain_config_ddp_configuration(self): - """Test distributed data parallel configuration.""" - config = pretrain_config() - - assert config.ddp.check_for_nan_in_grad is True - assert config.ddp.grad_reduce_in_fp32 is True - assert config.ddp.overlap_grad_reduce is True - assert config.ddp.overlap_param_gather is True - assert config.ddp.use_distributed_optimizer is True - - def test_pretrain_config_default_comm_overlap(self): - """Test default CommOverlapConfig setup.""" - config = pretrain_config() - - # Default setup should have comm overlap as None - assert config.comm_overlap is None - - def test_pretrain_config_custom_comm_overlap(self): - """Test custom CommOverlapConfig.""" - custom_overlap = CommOverlapConfig( - tp_comm_overlap=True, - defer_embedding_wgrad_compute=True, - wgrad_deferral_limit=50, - data_parallel_size=1, - ) - config = pretrain_config(comm_overlap_config=custom_overlap) - - # Should use the custom config - assert config.comm_overlap is not None - assert config.comm_overlap.tp_comm_overlap is True - - def test_pretrain_config_scheduler_configuration(self): - """Test scheduler configuration.""" - config = pretrain_config(train_iters=50000) - - assert config.scheduler.start_weight_decay == 0.033 - assert config.scheduler.end_weight_decay == 0.033 - assert config.scheduler.weight_decay_incr_style == "constant" - assert config.scheduler.lr_decay_style == "cosine" - assert config.scheduler.lr_warmup_iters == 2000 - assert config.scheduler.lr_warmup_init == 0.0 - assert config.scheduler.lr_decay_iters is None # Will be set to train_iters during validation - assert config.scheduler.override_opt_param_scheduler is True - - def test_pretrain_config_tokenizer_configuration(self): - """Test tokenizer configuration.""" - config = pretrain_config() - - assert config.tokenizer.tokenizer_type == "HuggingFaceTokenizer" - assert config.tokenizer.tokenizer_model == "EleutherAI/gpt-neox-20b" - - def test_pretrain_config_rng_configuration(self): - """Test RNG configuration.""" - config = pretrain_config() - - assert config.rng.seed == 1234 - assert config.dataset.random_seed == 1234 - - def test_pretrain_config_dataset_configuration(self): - """Test dataset configuration details.""" - config = pretrain_config() - - assert config.dataset.reset_attention_mask is False - assert config.dataset.reset_position_ids is False - assert config.dataset.eod_mask_loss is False - assert config.dataset.num_dataset_builder_threads == 1 - assert config.dataset.data_sharding is True - assert config.dataset.dataloader_type == "single" - assert config.dataset.num_workers == 8 - - def test_pretrain_config_logger_configuration(self): - """Test logger configuration.""" - config = pretrain_config() - - assert config.logger.log_interval == 10 - assert "tb_logs" in config.logger.tensorboard_dir - - @pytest.mark.parametrize( - "tensor_parallelism,pipeline_parallelism,context_parallelism", - [ - (1, 1, 1), - (1, 4, 2), - (1, 2, 4), - (1, 2, 2), - (1, 4, 1), - ], - ) - def test_pretrain_config_parallelism_combinations( - self, tensor_parallelism, pipeline_parallelism, context_parallelism - ): - """Test various parallelism combinations.""" - config = pretrain_config( - tensor_parallelism=tensor_parallelism, - pipeline_parallelism=pipeline_parallelism, - context_parallelism=context_parallelism, - pipeline_parallelism_dtype=torch.bfloat16, - ) - - assert config.model.tensor_model_parallel_size == tensor_parallelism - assert config.model.pipeline_model_parallel_size == pipeline_parallelism - assert config.model.context_parallel_size == context_parallelism - - @pytest.mark.parametrize( - "global_batch_size,micro_batch_size", - [ - (8, 1), - (16, 2), - (32, 4), - (64, 8), - ], - ) - def test_pretrain_config_batch_size_combinations(self, global_batch_size, micro_batch_size): - """Test various batch size combinations.""" - config = pretrain_config(global_batch_size=global_batch_size, micro_batch_size=micro_batch_size) - - assert config.train.global_batch_size == global_batch_size - assert config.train.micro_batch_size == micro_batch_size - - @pytest.mark.parametrize("seq_length", [1024, 2048, 4096, 8192, 16384]) - def test_pretrain_config_sequence_lengths(self, seq_length): - """Test various sequence lengths.""" - config = pretrain_config(seq_length=seq_length) - - assert config.dataset.sequence_length == seq_length - - @pytest.mark.parametrize("precision", ["fp16_mixed", "bf16_mixed"]) - def test_precision_recipes(self, precision): - """Test precision configuration.""" - cfg = pretrain_config(precision_config=precision) - assert cfg.mixed_precision == precision diff --git a/tests/unit_tests/recipes/mamba/test_mamba2_1_3b.py b/tests/unit_tests/recipes/mamba/test_mamba2_1_3b.py deleted file mode 100644 index 368ea7b09f..0000000000 --- a/tests/unit_tests/recipes/mamba/test_mamba2_1_3b.py +++ /dev/null @@ -1,367 +0,0 @@ -# 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 -import tempfile -from unittest.mock import patch - -import pytest -import torch - -from megatron.bridge.models.mamba import MambaModelProvider1P3B -from megatron.bridge.recipes.mamba.mamba2_1_3b import model_config, pretrain_config -from megatron.bridge.training.comm_overlap import CommOverlapConfig -from megatron.bridge.training.config import ConfigContainer - - -@pytest.mark.unit -class TestModelConfig: - """Test cases for the model_config function.""" - - def test_model_config_default_parameters(self): - """Test model_config with default parameters.""" - config = model_config() - - assert isinstance(config, MambaModelProvider1P3B) - assert config.tensor_model_parallel_size == 1 - assert config.pipeline_model_parallel_size == 1 - assert config.pipeline_dtype is None - assert config.virtual_pipeline_model_parallel_size is None - assert config.context_parallel_size == 1 - assert config.sequence_parallel is False - - def test_model_config_custom_pipeline_parallelism(self): - """Test model_config with custom pipeline parallelism.""" - config = model_config(pipeline_parallelism=8, pipeline_parallelism_dtype=torch.float16) - - assert config.tensor_model_parallel_size == 1 # default - assert config.pipeline_model_parallel_size == 8 - assert config.pipeline_dtype is torch.float16 - - def test_model_config_with_pipeline_dtype(self): - """Test model_config with pipeline dtype specified.""" - config = model_config(pipeline_parallelism=2, pipeline_parallelism_dtype=torch.float16) - - assert config.pipeline_model_parallel_size == 2 - assert config.pipeline_dtype == torch.float16 - - def test_model_config_virtual_pipeline_parallelism(self): - """Test model_config with virtual pipeline parallelism.""" - config = model_config(virtual_pipeline_parallelism=4) - - assert config.virtual_pipeline_model_parallel_size == 4 - - def test_model_config_context_parallelism(self): - """Test model_config with custom context parallelism.""" - config = model_config(context_parallelism=8) - - assert config.context_parallel_size == 8 - - def test_model_config_all_custom_parameters(self): - """Test model_config with all parameters customized.""" - config = model_config( - tensor_parallelism=1, # Mamba 1.3B has only 1 attention head - pipeline_parallelism=4, - pipeline_parallelism_dtype=torch.bfloat16, - virtual_pipeline_parallelism=8, - context_parallelism=16, - sequence_parallelism=False, # Must be False for tensor_parallelism=1 - ) - - assert config.tensor_model_parallel_size == 1 - assert config.pipeline_model_parallel_size == 4 - assert config.pipeline_dtype == torch.bfloat16 - assert config.virtual_pipeline_model_parallel_size == 8 - assert config.context_parallel_size == 16 - assert config.sequence_parallel is False - - -@pytest.mark.unit -class TestPretrainConfig: - """Test cases for the pretrain_config function.""" - - def test_pretrain_config_default_parameters(self): - """Test pretrain_config with default parameters (mock mode).""" - config = pretrain_config() - - assert isinstance(config, ConfigContainer) - assert isinstance(config.model, MambaModelProvider1P3B) - - # Check training configuration - assert config.train.train_iters == 1_168_251 - assert config.train.global_batch_size == 8 - assert config.train.micro_batch_size == 1 - assert config.train.eval_interval == 100 - assert config.train.eval_iters == 32 - - # Check optimizer configuration - assert config.optimizer.optimizer == "adam" - assert config.optimizer.lr == 3e-4 - assert config.optimizer.min_lr == 3e-5 - assert config.optimizer.weight_decay == 0.1 - assert config.optimizer.bf16 is True - assert config.optimizer.fp16 is False - - # Check dataset configuration (should be in mock mode) - assert config.dataset.sequence_length == 4096 - assert config.dataset.split == "1,1,1" - assert config.dataset.blend is None - assert config.dataset.blend_per_split is None - - def test_pretrain_config_custom_training_parameters(self): - """Test pretrain_config with custom training parameters.""" - config = pretrain_config( - train_iters=10000, - global_batch_size=256, - micro_batch_size=2, - seq_length=4096, - lr=1e-4, - min_lr=1e-5, - lr_warmup_iters=1000, - ) - - assert config.train.train_iters == 10000 - assert config.train.global_batch_size == 256 - assert config.train.micro_batch_size == 2 - assert config.dataset.sequence_length == 4096 - assert config.optimizer.lr == 1e-4 - assert config.optimizer.min_lr == 1e-5 - assert config.scheduler.lr_warmup_iters == 1000 # Note: fixed in scheduler config - assert config.scheduler.lr_decay_iters is None # Will be set to train_iters during validation - - def test_pretrain_config_custom_model_parameters(self): - """Test pretrain_config with custom model parameters.""" - config = pretrain_config( - tensor_parallelism=1, # Mamba 1.3B has only 1 attention head - pipeline_parallelism=2, - context_parallelism=8, - sequence_parallelism=False, # Must be False for tensor_parallelism=1 - pipeline_parallelism_dtype=torch.bfloat16, - ) - - assert config.model.tensor_model_parallel_size == 1 - assert config.model.pipeline_model_parallel_size == 2 - assert config.model.context_parallel_size == 8 - assert config.model.sequence_parallel is False - assert config.model.pipeline_dtype == torch.bfloat16 - - def test_pretrain_config_with_custom_directory(self): - """Test pretrain_config with custom directory.""" - with tempfile.TemporaryDirectory() as temp_dir: - config = pretrain_config(dir=temp_dir, name="test_run") - - expected_run_dir = os.path.join(temp_dir, "test_run") - expected_checkpoint_dir = os.path.join(expected_run_dir, "checkpoints") - expected_tensorboard_dir = os.path.join(expected_run_dir, "tb_logs") - - assert config.checkpoint.save == expected_checkpoint_dir - assert config.logger.tensorboard_dir == expected_tensorboard_dir - - def test_pretrain_config_explicit_mock_mode(self): - """Test pretrain_config with explicit mock=True.""" - config = pretrain_config(mock=True) - - assert config.dataset.blend is None - assert config.dataset.blend_per_split is None - assert config.dataset.split == "1,1,1" - - def test_pretrain_config_with_data_paths(self): - """Test pretrain_config with data paths provided.""" - - data_paths = ["/path/to/data1", "/path/to/data2", "/path/to/data3"] - config = pretrain_config(data_paths=data_paths) - - # Check that non-mock mode is configured - assert config.dataset.split == "9999,8,2" - assert config.dataset.blend is not None - assert config.dataset.blend_per_split is None - - def test_pretrain_config_with_train_valid_test_paths(self): - """Test pretrain_config with separate train/valid/test paths.""" - - config = pretrain_config( - train_data_path=["/path/to/train1", "/path/to/train2", "/path/to/train3"], - valid_data_path=["/path/to/valid1", "/path/to/valid2", "/path/to/valid3"], - test_data_path=["/path/to/test1", "/path/to/test2", "/path/to/test3"], - ) - - # When blend_per_split is used, split should be None - assert config.dataset.split is None - assert config.dataset.blend is None - assert config.dataset.blend_per_split is not None - - def test_pretrain_config_prioritizes_blend(self): - """Test that blend takes priority over blend_per_split when both are provided.""" - config = pretrain_config( - train_data_path=["/path/to/train1", "/path/to/train2"], - valid_data_path=["/path/to/valid1", "/path/to/valid2"], - test_data_path=["/path/to/test1", "/path/to/test2"], - data_paths=["/path/to/data1", "/path/to/data2"], - ) - - # Should prioritize blend over blend_per_split - assert config.dataset.split == "9999,8,2" - assert config.dataset.blend is not None - assert config.dataset.blend_per_split is None - - @patch("megatron.bridge.recipes.utils.dataset_utils.get_blend_and_blend_per_split") - def test_pretrain_config_fallback_to_mock_when_no_weights(self, mock_get_blend): - """Test pretrain_config falls back to mock when no weights are returned.""" - # Mock function returns None for both weights - mock_get_blend.return_value = (None, None) - - config = pretrain_config(data_paths=["/some/path"]) - - # Should fall back to mock mode - assert config.dataset.blend is None - assert config.dataset.blend_per_split is None - assert config.dataset.split == "1,1,1" - - def test_pretrain_config_checkpoint_configuration(self): - """Test checkpoint configuration in pretrain_config.""" - config = pretrain_config() - - assert config.checkpoint.save_interval == 2000 - assert config.checkpoint.ckpt_format == "torch_dist" - assert config.checkpoint.fully_parallel_load is True - - def test_pretrain_config_ddp_configuration(self): - """Test distributed data parallel configuration.""" - config = pretrain_config() - - assert config.ddp.check_for_nan_in_grad is True - assert config.ddp.grad_reduce_in_fp32 is True - assert config.ddp.overlap_grad_reduce is True - assert config.ddp.overlap_param_gather is True - assert config.ddp.use_distributed_optimizer is True - - def test_pretrain_config_default_comm_overlap(self): - """Test default CommOverlapConfig setup.""" - config = pretrain_config() - - # Default setup should have comm overlap config - assert config.comm_overlap is None # Not set by default - - def test_pretrain_config_custom_comm_overlap(self): - """Test custom CommOverlapConfig.""" - custom_overlap = CommOverlapConfig( - tp_comm_overlap=True, - defer_embedding_wgrad_compute=True, - wgrad_deferral_limit=50, - data_parallel_size=1, # Add this to avoid None - ) - config = pretrain_config(comm_overlap_config=custom_overlap) - - # Should use the custom config - assert config.comm_overlap is not None - - def test_pretrain_config_scheduler_configuration(self): - """Test scheduler configuration.""" - config = pretrain_config(train_iters=50000) - - assert config.scheduler.start_weight_decay == 0.033 - assert config.scheduler.end_weight_decay == 0.033 - assert config.scheduler.weight_decay_incr_style == "constant" - assert config.scheduler.lr_decay_style == "cosine" - assert config.scheduler.lr_warmup_iters == 2000 - assert config.scheduler.lr_warmup_init == 0.0 - assert config.scheduler.lr_decay_iters is None # Will be set to train_iters during validation - assert config.scheduler.override_opt_param_scheduler is True - - def test_pretrain_config_tokenizer_configuration(self): - """Test tokenizer configuration.""" - config = pretrain_config() - - assert config.tokenizer.tokenizer_type == "HuggingFaceTokenizer" - assert config.tokenizer.tokenizer_model == "EleutherAI/gpt-neox-20b" - - def test_pretrain_config_rng_configuration(self): - """Test RNG configuration.""" - config = pretrain_config() - - assert config.rng.seed == 1234 - assert config.dataset.random_seed == 1234 - - def test_pretrain_config_dataset_configuration(self): - """Test dataset configuration details.""" - config = pretrain_config() - - assert config.dataset.reset_attention_mask is False - assert config.dataset.reset_position_ids is False - assert config.dataset.eod_mask_loss is False - assert config.dataset.num_dataset_builder_threads == 1 - assert config.dataset.data_sharding is True - assert config.dataset.dataloader_type == "single" - assert config.dataset.num_workers == 8 - - def test_pretrain_config_logger_configuration(self): - """Test logger configuration.""" - config = pretrain_config() - - assert config.logger.log_interval == 10 - assert "tb_logs" in config.logger.tensorboard_dir - - @pytest.mark.parametrize( - "tensor_parallelism,pipeline_parallelism,context_parallelism", - [ - (1, 1, 1), - (1, 4, 2), - (1, 2, 4), - (1, 2, 2), - (1, 4, 1), - ], - ) - def test_pretrain_config_parallelism_combinations( - self, tensor_parallelism, pipeline_parallelism, context_parallelism - ): - """Test various parallelism combinations.""" - config = pretrain_config( - tensor_parallelism=tensor_parallelism, - pipeline_parallelism=pipeline_parallelism, - context_parallelism=context_parallelism, - pipeline_parallelism_dtype=torch.bfloat16, - ) - - assert config.model.tensor_model_parallel_size == tensor_parallelism - assert config.model.pipeline_model_parallel_size == pipeline_parallelism - assert config.model.context_parallel_size == context_parallelism - - @pytest.mark.parametrize( - "global_batch_size,micro_batch_size", - [ - (8, 1), - (16, 2), - (32, 4), - (64, 8), - ], - ) - def test_pretrain_config_batch_size_combinations(self, global_batch_size, micro_batch_size): - """Test various batch size combinations.""" - config = pretrain_config(global_batch_size=global_batch_size, micro_batch_size=micro_batch_size) - - assert config.train.global_batch_size == global_batch_size - assert config.train.micro_batch_size == micro_batch_size - - @pytest.mark.parametrize("seq_length", [1024, 2048, 4096, 8192, 16384]) - def test_pretrain_config_sequence_lengths(self, seq_length): - """Test various sequence lengths.""" - config = pretrain_config(seq_length=seq_length) - - assert config.dataset.sequence_length == seq_length - - @pytest.mark.parametrize("precision", ["fp16_mixed", "bf16_mixed"]) - def test_precision_recipes(self, precision): - """Test precision configuration.""" - cfg = pretrain_config(precision_config=precision) - assert cfg.mixed_precision == precision diff --git a/tests/unit_tests/recipes/mamba/test_mamba2_2_7b.py b/tests/unit_tests/recipes/mamba/test_mamba2_2_7b.py deleted file mode 100644 index cfd911c3a0..0000000000 --- a/tests/unit_tests/recipes/mamba/test_mamba2_2_7b.py +++ /dev/null @@ -1,367 +0,0 @@ -# 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 -import tempfile -from unittest.mock import patch - -import pytest -import torch - -from megatron.bridge.models.mamba import MambaModelProvider2P7B -from megatron.bridge.recipes.mamba.mamba2_2_7b import model_config, pretrain_config -from megatron.bridge.training.comm_overlap import CommOverlapConfig -from megatron.bridge.training.config import ConfigContainer - - -@pytest.mark.unit -class TestModelConfig: - """Test cases for the model_config function.""" - - def test_model_config_default_parameters(self): - """Test model_config with default parameters.""" - config = model_config() - - assert isinstance(config, MambaModelProvider2P7B) - assert config.tensor_model_parallel_size == 1 - assert config.pipeline_model_parallel_size == 1 - assert config.pipeline_dtype is None - assert config.virtual_pipeline_model_parallel_size is None - assert config.context_parallel_size == 1 - assert config.sequence_parallel is False - - def test_model_config_custom_pipeline_parallelism(self): - """Test model_config with custom pipeline parallelism.""" - config = model_config(pipeline_parallelism=8, pipeline_parallelism_dtype=torch.float16) - - assert config.tensor_model_parallel_size == 1 # default - assert config.pipeline_model_parallel_size == 8 - assert config.pipeline_dtype is torch.float16 - - def test_model_config_with_pipeline_dtype(self): - """Test model_config with pipeline dtype specified.""" - config = model_config(pipeline_parallelism=2, pipeline_parallelism_dtype=torch.float16) - - assert config.pipeline_model_parallel_size == 2 - assert config.pipeline_dtype == torch.float16 - - def test_model_config_virtual_pipeline_parallelism(self): - """Test model_config with virtual pipeline parallelism.""" - config = model_config(virtual_pipeline_parallelism=4) - - assert config.virtual_pipeline_model_parallel_size == 4 - - def test_model_config_context_parallelism(self): - """Test model_config with custom context parallelism.""" - config = model_config(context_parallelism=8) - - assert config.context_parallel_size == 8 - - def test_model_config_all_custom_parameters(self): - """Test model_config with all parameters customized.""" - config = model_config( - tensor_parallelism=1, # Mamba 2.7B has only 1 attention head - pipeline_parallelism=4, - pipeline_parallelism_dtype=torch.bfloat16, - virtual_pipeline_parallelism=8, - context_parallelism=16, - sequence_parallelism=False, # Must be False for tensor_parallelism=1 - ) - - assert config.tensor_model_parallel_size == 1 - assert config.pipeline_model_parallel_size == 4 - assert config.pipeline_dtype == torch.bfloat16 - assert config.virtual_pipeline_model_parallel_size == 8 - assert config.context_parallel_size == 16 - assert config.sequence_parallel is False - - -@pytest.mark.unit -class TestPretrainConfig: - """Test cases for the pretrain_config function.""" - - def test_pretrain_config_default_parameters(self): - """Test pretrain_config with default parameters (mock mode).""" - config = pretrain_config() - - assert isinstance(config, ConfigContainer) - assert isinstance(config.model, MambaModelProvider2P7B) - - # Check training configuration - assert config.train.train_iters == 1_168_251 - assert config.train.global_batch_size == 8 - assert config.train.micro_batch_size == 1 - assert config.train.eval_interval == 100 - assert config.train.eval_iters == 32 - - # Check optimizer configuration - assert config.optimizer.optimizer == "adam" - assert config.optimizer.lr == 3e-4 - assert config.optimizer.min_lr == 3e-5 - assert config.optimizer.weight_decay == 0.1 - assert config.optimizer.bf16 is True - assert config.optimizer.fp16 is False - - # Check dataset configuration (should be in mock mode) - assert config.dataset.sequence_length == 4096 - assert config.dataset.split == "1,1,1" - assert config.dataset.blend is None - assert config.dataset.blend_per_split is None - - def test_pretrain_config_custom_training_parameters(self): - """Test pretrain_config with custom training parameters.""" - config = pretrain_config( - train_iters=10000, - global_batch_size=256, - micro_batch_size=2, - seq_length=4096, - lr=1e-4, - min_lr=1e-5, - lr_warmup_iters=1000, - ) - - assert config.train.train_iters == 10000 - assert config.train.global_batch_size == 256 - assert config.train.micro_batch_size == 2 - assert config.dataset.sequence_length == 4096 - assert config.optimizer.lr == 1e-4 - assert config.optimizer.min_lr == 1e-5 - assert config.scheduler.lr_warmup_iters == 1000 # Note: fixed in scheduler config - assert config.scheduler.lr_decay_iters is None # Will be set to train_iters during validation - - def test_pretrain_config_custom_model_parameters(self): - """Test pretrain_config with custom model parameters.""" - config = pretrain_config( - tensor_parallelism=1, # Mamba 2.7B has only 1 attention head - pipeline_parallelism=2, - context_parallelism=8, - sequence_parallelism=False, # Must be False for tensor_parallelism=1 - pipeline_parallelism_dtype=torch.bfloat16, - ) - - assert config.model.tensor_model_parallel_size == 1 - assert config.model.pipeline_model_parallel_size == 2 - assert config.model.context_parallel_size == 8 - assert config.model.sequence_parallel is False - assert config.model.pipeline_dtype == torch.bfloat16 - - def test_pretrain_config_with_custom_directory(self): - """Test pretrain_config with custom directory.""" - with tempfile.TemporaryDirectory() as temp_dir: - config = pretrain_config(dir=temp_dir, name="test_run") - - expected_run_dir = os.path.join(temp_dir, "test_run") - expected_checkpoint_dir = os.path.join(expected_run_dir, "checkpoints") - expected_tensorboard_dir = os.path.join(expected_run_dir, "tb_logs") - - assert config.checkpoint.save == expected_checkpoint_dir - assert config.logger.tensorboard_dir == expected_tensorboard_dir - - def test_pretrain_config_explicit_mock_mode(self): - """Test pretrain_config with explicit mock=True.""" - config = pretrain_config(mock=True) - - assert config.dataset.blend is None - assert config.dataset.blend_per_split is None - assert config.dataset.split == "1,1,1" - - def test_pretrain_config_with_data_paths(self): - """Test pretrain_config with data paths provided.""" - - data_paths = ["/path/to/data1", "/path/to/data2", "/path/to/data3"] - config = pretrain_config(data_paths=data_paths) - - # Check that non-mock mode is configured - assert config.dataset.split == "9999,8,2" - assert config.dataset.blend is not None - assert config.dataset.blend_per_split is None - - def test_pretrain_config_with_train_valid_test_paths(self): - """Test pretrain_config with separate train/valid/test paths.""" - - config = pretrain_config( - train_data_path=["/path/to/train1", "/path/to/train2", "/path/to/train3"], - valid_data_path=["/path/to/valid1", "/path/to/valid2", "/path/to/valid3"], - test_data_path=["/path/to/test1", "/path/to/test2", "/path/to/test3"], - ) - - # When blend_per_split is used, split should be None - assert config.dataset.split is None - assert config.dataset.blend is None - assert config.dataset.blend_per_split is not None - - def test_pretrain_config_prioritizes_blend(self): - """Test that blend takes priority over blend_per_split when both are provided.""" - config = pretrain_config( - train_data_path=["/path/to/train1", "/path/to/train2"], - valid_data_path=["/path/to/valid1", "/path/to/valid2"], - test_data_path=["/path/to/test1", "/path/to/test2"], - data_paths=["/path/to/data1", "/path/to/data2"], - ) - - # Should prioritize blend over blend_per_split - assert config.dataset.split == "9999,8,2" - assert config.dataset.blend is not None - assert config.dataset.blend_per_split is None - - @patch("megatron.bridge.recipes.utils.dataset_utils.get_blend_and_blend_per_split") - def test_pretrain_config_fallback_to_mock_when_no_weights(self, mock_get_blend): - """Test pretrain_config falls back to mock when no weights are returned.""" - # Mock function returns None for both weights - mock_get_blend.return_value = (None, None) - - config = pretrain_config(data_paths=["/some/path"]) - - # Should fall back to mock mode - assert config.dataset.blend is None - assert config.dataset.blend_per_split is None - assert config.dataset.split == "1,1,1" - - def test_pretrain_config_checkpoint_configuration(self): - """Test checkpoint configuration in pretrain_config.""" - config = pretrain_config() - - assert config.checkpoint.save_interval == 2000 - assert config.checkpoint.ckpt_format == "torch_dist" - assert config.checkpoint.fully_parallel_load is True - - def test_pretrain_config_ddp_configuration(self): - """Test distributed data parallel configuration.""" - config = pretrain_config() - - assert config.ddp.check_for_nan_in_grad is True - assert config.ddp.grad_reduce_in_fp32 is True - assert config.ddp.overlap_grad_reduce is True - assert config.ddp.overlap_param_gather is True - assert config.ddp.use_distributed_optimizer is True - - def test_pretrain_config_default_comm_overlap(self): - """Test default CommOverlapConfig setup.""" - config = pretrain_config() - - # Default setup should have comm overlap config - assert config.comm_overlap is None # Not set by default - - def test_pretrain_config_custom_comm_overlap(self): - """Test custom CommOverlapConfig.""" - custom_overlap = CommOverlapConfig( - tp_comm_overlap=True, - defer_embedding_wgrad_compute=True, - wgrad_deferral_limit=50, - data_parallel_size=1, # Add this to avoid None - ) - config = pretrain_config(comm_overlap_config=custom_overlap) - - # Should use the custom config - assert config.comm_overlap is not None - - def test_pretrain_config_scheduler_configuration(self): - """Test scheduler configuration.""" - config = pretrain_config(train_iters=50000) - - assert config.scheduler.start_weight_decay == 0.033 - assert config.scheduler.end_weight_decay == 0.033 - assert config.scheduler.weight_decay_incr_style == "constant" - assert config.scheduler.lr_decay_style == "cosine" - assert config.scheduler.lr_warmup_iters == 2000 - assert config.scheduler.lr_warmup_init == 0.0 - assert config.scheduler.lr_decay_iters is None # Will be set to train_iters during validation - assert config.scheduler.override_opt_param_scheduler is True - - def test_pretrain_config_tokenizer_configuration(self): - """Test tokenizer configuration.""" - config = pretrain_config() - - assert config.tokenizer.tokenizer_type == "HuggingFaceTokenizer" - assert config.tokenizer.tokenizer_model == "EleutherAI/gpt-neox-20b" - - def test_pretrain_config_rng_configuration(self): - """Test RNG configuration.""" - config = pretrain_config() - - assert config.rng.seed == 1234 - assert config.dataset.random_seed == 1234 - - def test_pretrain_config_dataset_configuration(self): - """Test dataset configuration details.""" - config = pretrain_config() - - assert config.dataset.reset_attention_mask is False - assert config.dataset.reset_position_ids is False - assert config.dataset.eod_mask_loss is False - assert config.dataset.num_dataset_builder_threads == 1 - assert config.dataset.data_sharding is True - assert config.dataset.dataloader_type == "single" - assert config.dataset.num_workers == 8 - - def test_pretrain_config_logger_configuration(self): - """Test logger configuration.""" - config = pretrain_config() - - assert config.logger.log_interval == 10 - assert "tb_logs" in config.logger.tensorboard_dir - - @pytest.mark.parametrize( - "tensor_parallelism,pipeline_parallelism,context_parallelism", - [ - (1, 1, 1), - (1, 4, 2), - (1, 2, 4), - (1, 2, 2), - (1, 4, 1), - ], - ) - def test_pretrain_config_parallelism_combinations( - self, tensor_parallelism, pipeline_parallelism, context_parallelism - ): - """Test various parallelism combinations.""" - config = pretrain_config( - tensor_parallelism=tensor_parallelism, - pipeline_parallelism=pipeline_parallelism, - context_parallelism=context_parallelism, - pipeline_parallelism_dtype=torch.bfloat16, - ) - - assert config.model.tensor_model_parallel_size == tensor_parallelism - assert config.model.pipeline_model_parallel_size == pipeline_parallelism - assert config.model.context_parallel_size == context_parallelism - - @pytest.mark.parametrize( - "global_batch_size,micro_batch_size", - [ - (8, 1), - (16, 2), - (32, 4), - (64, 8), - ], - ) - def test_pretrain_config_batch_size_combinations(self, global_batch_size, micro_batch_size): - """Test various batch size combinations.""" - config = pretrain_config(global_batch_size=global_batch_size, micro_batch_size=micro_batch_size) - - assert config.train.global_batch_size == global_batch_size - assert config.train.micro_batch_size == micro_batch_size - - @pytest.mark.parametrize("seq_length", [1024, 2048, 4096, 8192, 16384]) - def test_pretrain_config_sequence_lengths(self, seq_length): - """Test various sequence lengths.""" - config = pretrain_config(seq_length=seq_length) - - assert config.dataset.sequence_length == seq_length - - @pytest.mark.parametrize("precision", ["fp16_mixed", "bf16_mixed"]) - def test_precision_recipes(self, precision): - """Test precision configuration.""" - cfg = pretrain_config(precision_config=precision) - assert cfg.mixed_precision == precision diff --git a/tests/unit_tests/recipes/mamba/test_mamba2_370m.py b/tests/unit_tests/recipes/mamba/test_mamba2_370m.py deleted file mode 100644 index 9c239934b9..0000000000 --- a/tests/unit_tests/recipes/mamba/test_mamba2_370m.py +++ /dev/null @@ -1,366 +0,0 @@ -# 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 -import tempfile -from unittest.mock import patch - -import pytest -import torch - -from megatron.bridge.models.mamba import MambaModelProvider370M -from megatron.bridge.recipes.mamba.mamba2_370m import model_config, pretrain_config -from megatron.bridge.training.comm_overlap import CommOverlapConfig -from megatron.bridge.training.config import ConfigContainer - - -@pytest.mark.unit -class TestModelConfig: - """Test cases for the model_config function.""" - - def test_model_config_default_parameters(self): - """Test model_config with default parameters.""" - config = model_config() - - assert isinstance(config, MambaModelProvider370M) - assert config.tensor_model_parallel_size == 1 - assert config.pipeline_model_parallel_size == 1 - assert config.pipeline_dtype is None - assert config.virtual_pipeline_model_parallel_size is None - assert config.context_parallel_size == 1 - assert config.sequence_parallel is False - - def test_model_config_custom_pipeline_parallelism(self): - """Test model_config with custom pipeline parallelism.""" - config = model_config(pipeline_parallelism=8, pipeline_parallelism_dtype=torch.float16) - - assert config.tensor_model_parallel_size == 1 # default - assert config.pipeline_model_parallel_size == 8 - assert config.pipeline_dtype is torch.float16 - - def test_model_config_with_pipeline_dtype(self): - """Test model_config with pipeline dtype specified.""" - config = model_config(pipeline_parallelism=2, pipeline_parallelism_dtype=torch.float16) - - assert config.pipeline_model_parallel_size == 2 - assert config.pipeline_dtype == torch.float16 - - def test_model_config_virtual_pipeline_parallelism(self): - """Test model_config with virtual pipeline parallelism.""" - config = model_config(virtual_pipeline_parallelism=4) - - assert config.virtual_pipeline_model_parallel_size == 4 - - def test_model_config_context_parallelism(self): - """Test model_config with custom context parallelism.""" - config = model_config(context_parallelism=8) - - assert config.context_parallel_size == 8 - - def test_model_config_all_custom_parameters(self): - """Test model_config with all parameters customized.""" - config = model_config( - tensor_parallelism=1, # Mamba 370M has only 1 attention head - pipeline_parallelism=4, - pipeline_parallelism_dtype=torch.bfloat16, - virtual_pipeline_parallelism=8, - context_parallelism=16, - sequence_parallelism=False, # Cannot use sequence parallelism with tensor_parallelism=1 - ) - - assert config.tensor_model_parallel_size == 1 - assert config.pipeline_model_parallel_size == 4 - assert config.pipeline_dtype == torch.bfloat16 - assert config.virtual_pipeline_model_parallel_size == 8 - assert config.context_parallel_size == 16 - assert config.sequence_parallel is False - - -@pytest.mark.unit -class TestPretrainConfig: - """Test cases for the pretrain_config function.""" - - def test_pretrain_config_default_parameters(self): - """Test pretrain_config with default parameters (mock mode).""" - config = pretrain_config() - - assert isinstance(config, ConfigContainer) - assert isinstance(config.model, MambaModelProvider370M) - - # Check training configuration - assert config.train.train_iters == 1_168_251 - assert config.train.global_batch_size == 8 - assert config.train.micro_batch_size == 1 - assert config.train.eval_interval == 100 - assert config.train.eval_iters == 32 - - # Check optimizer configuration - assert config.optimizer.optimizer == "adam" - assert config.optimizer.lr == 3e-4 - assert config.optimizer.min_lr == 3e-5 - assert config.optimizer.weight_decay == 0.1 - assert config.optimizer.bf16 is True - assert config.optimizer.fp16 is False - - # Check dataset configuration (should be in mock mode) - assert config.dataset.sequence_length == 4096 - assert config.dataset.split == "1,1,1" - assert config.dataset.blend is None - assert config.dataset.blend_per_split is None - - def test_pretrain_config_custom_training_parameters(self): - """Test pretrain_config with custom training parameters.""" - config = pretrain_config( - train_iters=10000, - global_batch_size=256, - micro_batch_size=2, - seq_length=4096, - lr=1e-4, - min_lr=1e-5, - lr_warmup_iters=1000, - ) - - assert config.train.train_iters == 10000 - assert config.train.global_batch_size == 256 - assert config.train.micro_batch_size == 2 - assert config.dataset.sequence_length == 4096 - assert config.optimizer.lr == 1e-4 - assert config.optimizer.min_lr == 1e-5 - assert config.scheduler.lr_warmup_iters == 1000 - assert config.scheduler.lr_decay_iters is None # Will be set to train_iters during validation - - def test_pretrain_config_custom_model_parameters(self): - """Test pretrain_config with custom model parameters.""" - config = pretrain_config( - tensor_parallelism=1, # Mamba 370M has only 1 attention head - pipeline_parallelism=2, - context_parallelism=8, - sequence_parallelism=False, # Cannot use sequence parallelism with tensor_parallelism=1 - pipeline_parallelism_dtype=torch.bfloat16, - ) - - assert config.model.tensor_model_parallel_size == 1 - assert config.model.pipeline_model_parallel_size == 2 - assert config.model.context_parallel_size == 8 - assert config.model.sequence_parallel is False - assert config.model.pipeline_dtype == torch.bfloat16 - - def test_pretrain_config_with_custom_directory(self): - """Test pretrain_config with custom directory.""" - with tempfile.TemporaryDirectory() as temp_dir: - config = pretrain_config(dir=temp_dir, name="test_run") - - expected_run_dir = os.path.join(temp_dir, "test_run") - expected_checkpoint_dir = os.path.join(expected_run_dir, "checkpoints") - expected_tensorboard_dir = os.path.join(expected_run_dir, "tb_logs") - - assert config.checkpoint.save == expected_checkpoint_dir - assert config.logger.tensorboard_dir == expected_tensorboard_dir - - def test_pretrain_config_explicit_mock_mode(self): - """Test pretrain_config with explicit mock=True.""" - config = pretrain_config(mock=True) - - assert config.dataset.blend is None - assert config.dataset.blend_per_split is None - assert config.dataset.split == "1,1,1" - - def test_pretrain_config_with_data_paths(self): - """Test pretrain_config with data paths provided.""" - data_paths = ["/path/to/data1", "/path/to/data2", "/path/to/data3"] - config = pretrain_config(data_paths=data_paths) - - # Check that non-mock mode is configured - assert config.dataset.split == "9999,8,2" - assert config.dataset.blend is not None - assert config.dataset.blend_per_split is None - - def test_pretrain_config_with_train_valid_test_paths(self): - """Test pretrain_config with separate train/valid/test paths.""" - config = pretrain_config( - train_data_path=["/path/to/train1", "/path/to/train2", "/path/to/train3"], - valid_data_path=["/path/to/valid1", "/path/to/valid2", "/path/to/valid3"], - test_data_path=["/path/to/test1", "/path/to/test2", "/path/to/test3"], - ) - - # When blend_per_split is used, split should be None - assert config.dataset.split is None - assert config.dataset.blend is None - assert config.dataset.blend_per_split is not None - - def test_pretrain_config_prioritizes_blend(self): - """Test that blend takes priority over blend_per_split when both are provided.""" - config = pretrain_config( - train_data_path=["/path/to/train1", "/path/to/train2"], - valid_data_path=["/path/to/valid1", "/path/to/valid2"], - test_data_path=["/path/to/test1", "/path/to/test2"], - data_paths=["/path/to/data1", "/path/to/data2"], - ) - - # Should prioritize blend over blend_per_split - assert config.dataset.split == "9999,8,2" - assert config.dataset.blend is not None - assert config.dataset.blend_per_split is None - - @patch("megatron.bridge.recipes.utils.dataset_utils.get_blend_and_blend_per_split") - def test_pretrain_config_fallback_to_mock_when_no_weights(self, mock_get_blend): - """Test pretrain_config falls back to mock when no weights are returned.""" - # Mock function returns None for both weights - mock_get_blend.return_value = (None, None) - - config = pretrain_config(data_paths=["/some/path"]) - - # Should fall back to mock mode - assert config.dataset.blend is None - assert config.dataset.blend_per_split is None - assert config.dataset.split == "1,1,1" - - def test_pretrain_config_checkpoint_configuration(self): - """Test checkpoint configuration in pretrain_config.""" - config = pretrain_config() - - assert config.checkpoint.save_interval == 2000 - assert config.checkpoint.ckpt_format == "torch_dist" - assert config.checkpoint.fully_parallel_load is True - - def test_pretrain_config_ddp_configuration(self): - """Test distributed data parallel configuration.""" - config = pretrain_config() - - assert config.ddp.check_for_nan_in_grad is True - assert config.ddp.grad_reduce_in_fp32 is True - assert config.ddp.overlap_grad_reduce is True - assert config.ddp.overlap_param_gather is True - assert config.ddp.use_distributed_optimizer is True - - def test_pretrain_config_default_comm_overlap(self): - """Test default CommOverlapConfig setup.""" - config = pretrain_config() - - # Default setup should have comm overlap as None - assert config.comm_overlap is None - - def test_pretrain_config_custom_comm_overlap(self): - """Test custom CommOverlapConfig.""" - custom_overlap = CommOverlapConfig( - tp_comm_overlap=True, - defer_embedding_wgrad_compute=True, - wgrad_deferral_limit=50, - data_parallel_size=1, - ) - config = pretrain_config(comm_overlap_config=custom_overlap) - - # Should use the custom config - assert config.comm_overlap is not None - assert config.comm_overlap.tp_comm_overlap is True - - def test_pretrain_config_scheduler_configuration(self): - """Test scheduler configuration.""" - config = pretrain_config(train_iters=50000) - - assert config.scheduler.start_weight_decay == 0.033 - assert config.scheduler.end_weight_decay == 0.033 - assert config.scheduler.weight_decay_incr_style == "constant" - assert config.scheduler.lr_decay_style == "cosine" - assert config.scheduler.lr_warmup_iters == 2000 - assert config.scheduler.lr_warmup_init == 0.0 - assert config.scheduler.lr_decay_iters is None # Will be set to train_iters during validation - assert config.scheduler.override_opt_param_scheduler is True - - def test_pretrain_config_tokenizer_configuration(self): - """Test tokenizer configuration.""" - config = pretrain_config() - - assert config.tokenizer.tokenizer_type == "HuggingFaceTokenizer" - assert config.tokenizer.tokenizer_model == "EleutherAI/gpt-neox-20b" - - def test_pretrain_config_rng_configuration(self): - """Test RNG configuration.""" - config = pretrain_config() - - assert config.rng.seed == 1234 - assert config.dataset.random_seed == 1234 - - def test_pretrain_config_dataset_configuration(self): - """Test dataset configuration details.""" - config = pretrain_config() - - assert config.dataset.reset_attention_mask is False - assert config.dataset.reset_position_ids is False - assert config.dataset.eod_mask_loss is False - assert config.dataset.num_dataset_builder_threads == 1 - assert config.dataset.data_sharding is True - assert config.dataset.dataloader_type == "single" - assert config.dataset.num_workers == 8 - - def test_pretrain_config_logger_configuration(self): - """Test logger configuration.""" - config = pretrain_config() - - assert config.logger.log_interval == 10 - assert "tb_logs" in config.logger.tensorboard_dir - - @pytest.mark.parametrize( - "tensor_parallelism,pipeline_parallelism,context_parallelism", - [ - (1, 1, 1), - (1, 4, 2), - (1, 2, 4), - (1, 2, 2), - (1, 4, 1), - ], - ) - def test_pretrain_config_parallelism_combinations( - self, tensor_parallelism, pipeline_parallelism, context_parallelism - ): - """Test various parallelism combinations.""" - config = pretrain_config( - tensor_parallelism=tensor_parallelism, - pipeline_parallelism=pipeline_parallelism, - context_parallelism=context_parallelism, - pipeline_parallelism_dtype=torch.bfloat16, - ) - - assert config.model.tensor_model_parallel_size == tensor_parallelism - assert config.model.pipeline_model_parallel_size == pipeline_parallelism - assert config.model.context_parallel_size == context_parallelism - - @pytest.mark.parametrize( - "global_batch_size,micro_batch_size", - [ - (8, 1), - (16, 2), - (32, 4), - (64, 8), - ], - ) - def test_pretrain_config_batch_size_combinations(self, global_batch_size, micro_batch_size): - """Test various batch size combinations.""" - config = pretrain_config(global_batch_size=global_batch_size, micro_batch_size=micro_batch_size) - - assert config.train.global_batch_size == global_batch_size - assert config.train.micro_batch_size == micro_batch_size - - @pytest.mark.parametrize("seq_length", [1024, 2048, 4096, 8192, 16384]) - def test_pretrain_config_sequence_lengths(self, seq_length): - """Test various sequence lengths.""" - config = pretrain_config(seq_length=seq_length) - - assert config.dataset.sequence_length == seq_length - - @pytest.mark.parametrize("precision", ["fp16_mixed", "bf16_mixed"]) - def test_precision_recipes(self, precision): - """Test precision configuration.""" - cfg = pretrain_config(precision_config=precision) - assert cfg.mixed_precision == precision diff --git a/tests/unit_tests/recipes/mamba/test_mamba2_780m.py b/tests/unit_tests/recipes/mamba/test_mamba2_780m.py deleted file mode 100644 index 406ed6c40b..0000000000 --- a/tests/unit_tests/recipes/mamba/test_mamba2_780m.py +++ /dev/null @@ -1,367 +0,0 @@ -# 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 -import tempfile -from unittest.mock import patch - -import pytest -import torch - -from megatron.bridge.models.mamba import MambaModelProvider780M -from megatron.bridge.recipes.mamba.mamba2_780m import model_config, pretrain_config -from megatron.bridge.training.comm_overlap import CommOverlapConfig -from megatron.bridge.training.config import ConfigContainer - - -@pytest.mark.unit -class TestModelConfig: - """Test cases for the model_config function.""" - - def test_model_config_default_parameters(self): - """Test model_config with default parameters.""" - config = model_config() - - assert isinstance(config, MambaModelProvider780M) - assert config.tensor_model_parallel_size == 1 - assert config.pipeline_model_parallel_size == 1 - assert config.pipeline_dtype is None - assert config.virtual_pipeline_model_parallel_size is None - assert config.context_parallel_size == 1 - assert config.sequence_parallel is False - - def test_model_config_custom_pipeline_parallelism(self): - """Test model_config with custom pipeline parallelism.""" - config = model_config(pipeline_parallelism=8, pipeline_parallelism_dtype=torch.float16) - - assert config.tensor_model_parallel_size == 1 # default - assert config.pipeline_model_parallel_size == 8 - assert config.pipeline_dtype is torch.float16 - - def test_model_config_with_pipeline_dtype(self): - """Test model_config with pipeline dtype specified.""" - config = model_config(pipeline_parallelism=2, pipeline_parallelism_dtype=torch.float16) - - assert config.pipeline_model_parallel_size == 2 - assert config.pipeline_dtype == torch.float16 - - def test_model_config_virtual_pipeline_parallelism(self): - """Test model_config with virtual pipeline parallelism.""" - config = model_config(virtual_pipeline_parallelism=4) - - assert config.virtual_pipeline_model_parallel_size == 4 - - def test_model_config_context_parallelism(self): - """Test model_config with custom context parallelism.""" - config = model_config(context_parallelism=8) - - assert config.context_parallel_size == 8 - - def test_model_config_all_custom_parameters(self): - """Test model_config with all parameters customized.""" - config = model_config( - tensor_parallelism=1, # Mamba 780M has only 1 attention head - pipeline_parallelism=4, - pipeline_parallelism_dtype=torch.bfloat16, - virtual_pipeline_parallelism=8, - context_parallelism=16, - sequence_parallelism=False, # Must be False for tensor_parallelism=1 - ) - - assert config.tensor_model_parallel_size == 1 - assert config.pipeline_model_parallel_size == 4 - assert config.pipeline_dtype == torch.bfloat16 - assert config.virtual_pipeline_model_parallel_size == 8 - assert config.context_parallel_size == 16 - assert config.sequence_parallel is False - - -@pytest.mark.unit -class TestPretrainConfig: - """Test cases for the pretrain_config function.""" - - def test_pretrain_config_default_parameters(self): - """Test pretrain_config with default parameters (mock mode).""" - config = pretrain_config() - - assert isinstance(config, ConfigContainer) - assert isinstance(config.model, MambaModelProvider780M) - - # Check training configuration - assert config.train.train_iters == 1_168_251 - assert config.train.global_batch_size == 8 - assert config.train.micro_batch_size == 1 - assert config.train.eval_interval == 100 - assert config.train.eval_iters == 32 - - # Check optimizer configuration - assert config.optimizer.optimizer == "adam" - assert config.optimizer.lr == 3e-4 - assert config.optimizer.min_lr == 3e-5 - assert config.optimizer.weight_decay == 0.1 - assert config.optimizer.bf16 is True - assert config.optimizer.fp16 is False - - # Check dataset configuration (should be in mock mode) - assert config.dataset.sequence_length == 4096 - assert config.dataset.split == "1,1,1" - assert config.dataset.blend is None - assert config.dataset.blend_per_split is None - - def test_pretrain_config_custom_training_parameters(self): - """Test pretrain_config with custom training parameters.""" - config = pretrain_config( - train_iters=10000, - global_batch_size=256, - micro_batch_size=2, - seq_length=4096, - lr=1e-4, - min_lr=1e-5, - lr_warmup_iters=1000, - ) - - assert config.train.train_iters == 10000 - assert config.train.global_batch_size == 256 - assert config.train.micro_batch_size == 2 - assert config.dataset.sequence_length == 4096 - assert config.optimizer.lr == 1e-4 - assert config.optimizer.min_lr == 1e-5 - assert config.scheduler.lr_warmup_iters == 1000 # Note: fixed in scheduler config - assert config.scheduler.lr_decay_iters is None # Will be set to train_iters during validation - - def test_pretrain_config_custom_model_parameters(self): - """Test pretrain_config with custom model parameters.""" - config = pretrain_config( - tensor_parallelism=1, # Mamba 780M has only 1 attention head - pipeline_parallelism=2, - context_parallelism=8, - sequence_parallelism=False, # Must be False for tensor_parallelism=1 - pipeline_parallelism_dtype=torch.bfloat16, - ) - - assert config.model.tensor_model_parallel_size == 1 - assert config.model.pipeline_model_parallel_size == 2 - assert config.model.context_parallel_size == 8 - assert config.model.sequence_parallel is False - assert config.model.pipeline_dtype == torch.bfloat16 - - def test_pretrain_config_with_custom_directory(self): - """Test pretrain_config with custom directory.""" - with tempfile.TemporaryDirectory() as temp_dir: - config = pretrain_config(dir=temp_dir, name="test_run") - - expected_run_dir = os.path.join(temp_dir, "test_run") - expected_checkpoint_dir = os.path.join(expected_run_dir, "checkpoints") - expected_tensorboard_dir = os.path.join(expected_run_dir, "tb_logs") - - assert config.checkpoint.save == expected_checkpoint_dir - assert config.logger.tensorboard_dir == expected_tensorboard_dir - - def test_pretrain_config_explicit_mock_mode(self): - """Test pretrain_config with explicit mock=True.""" - config = pretrain_config(mock=True) - - assert config.dataset.blend is None - assert config.dataset.blend_per_split is None - assert config.dataset.split == "1,1,1" - - def test_pretrain_config_with_data_paths(self): - """Test pretrain_config with data paths provided.""" - - data_paths = ["/path/to/data1", "/path/to/data2", "/path/to/data3"] - config = pretrain_config(data_paths=data_paths) - - # Check that non-mock mode is configured - assert config.dataset.split == "9999,8,2" - assert config.dataset.blend is not None - assert config.dataset.blend_per_split is None - - def test_pretrain_config_with_train_valid_test_paths(self): - """Test pretrain_config with separate train/valid/test paths.""" - - config = pretrain_config( - train_data_path=["/path/to/train1", "/path/to/train2", "/path/to/train3"], - valid_data_path=["/path/to/valid1", "/path/to/valid2", "/path/to/valid3"], - test_data_path=["/path/to/test1", "/path/to/test2", "/path/to/test3"], - ) - - # When blend_per_split is used, split should be None - assert config.dataset.split is None - assert config.dataset.blend is None - assert config.dataset.blend_per_split is not None - - def test_pretrain_config_prioritizes_blend(self): - """Test that blend takes priority over blend_per_split when both are provided.""" - config = pretrain_config( - train_data_path=["/path/to/train1", "/path/to/train2"], - valid_data_path=["/path/to/valid1", "/path/to/valid2"], - test_data_path=["/path/to/test1", "/path/to/test2"], - data_paths=["/path/to/data1", "/path/to/data2"], - ) - - # Should prioritize blend over blend_per_split - assert config.dataset.split == "9999,8,2" - assert config.dataset.blend is not None - assert config.dataset.blend_per_split is None - - @patch("megatron.bridge.recipes.utils.dataset_utils.get_blend_and_blend_per_split") - def test_pretrain_config_fallback_to_mock_when_no_weights(self, mock_get_blend): - """Test pretrain_config falls back to mock when no weights are returned.""" - # Mock function returns None for both weights - mock_get_blend.return_value = (None, None) - - config = pretrain_config(data_paths=["/some/path"]) - - # Should fall back to mock mode - assert config.dataset.blend is None - assert config.dataset.blend_per_split is None - assert config.dataset.split == "1,1,1" - - def test_pretrain_config_checkpoint_configuration(self): - """Test checkpoint configuration in pretrain_config.""" - config = pretrain_config() - - assert config.checkpoint.save_interval == 2000 - assert config.checkpoint.ckpt_format == "torch_dist" - assert config.checkpoint.fully_parallel_load is True - - def test_pretrain_config_ddp_configuration(self): - """Test distributed data parallel configuration.""" - config = pretrain_config() - - assert config.ddp.check_for_nan_in_grad is True - assert config.ddp.grad_reduce_in_fp32 is True - assert config.ddp.overlap_grad_reduce is True - assert config.ddp.overlap_param_gather is True - assert config.ddp.use_distributed_optimizer is True - - def test_pretrain_config_default_comm_overlap(self): - """Test default CommOverlapConfig setup.""" - config = pretrain_config() - - # Default setup should have comm overlap config - assert config.comm_overlap is None # Not set by default - - def test_pretrain_config_custom_comm_overlap(self): - """Test custom CommOverlapConfig.""" - custom_overlap = CommOverlapConfig( - tp_comm_overlap=True, - defer_embedding_wgrad_compute=True, - wgrad_deferral_limit=50, - data_parallel_size=1, # Add this to avoid None - ) - config = pretrain_config(comm_overlap_config=custom_overlap) - - # Should use the custom config - assert config.comm_overlap is not None - - def test_pretrain_config_scheduler_configuration(self): - """Test scheduler configuration.""" - config = pretrain_config(train_iters=50000) - - assert config.scheduler.start_weight_decay == 0.033 - assert config.scheduler.end_weight_decay == 0.033 - assert config.scheduler.weight_decay_incr_style == "constant" - assert config.scheduler.lr_decay_style == "cosine" - assert config.scheduler.lr_warmup_iters == 2000 - assert config.scheduler.lr_warmup_init == 0.0 - assert config.scheduler.lr_decay_iters is None # Will be set to train_iters during validation - assert config.scheduler.override_opt_param_scheduler is True - - def test_pretrain_config_tokenizer_configuration(self): - """Test tokenizer configuration.""" - config = pretrain_config() - - assert config.tokenizer.tokenizer_type == "HuggingFaceTokenizer" - assert config.tokenizer.tokenizer_model == "EleutherAI/gpt-neox-20b" - - def test_pretrain_config_rng_configuration(self): - """Test RNG configuration.""" - config = pretrain_config() - - assert config.rng.seed == 1234 - assert config.dataset.random_seed == 1234 - - def test_pretrain_config_dataset_configuration(self): - """Test dataset configuration details.""" - config = pretrain_config() - - assert config.dataset.reset_attention_mask is False - assert config.dataset.reset_position_ids is False - assert config.dataset.eod_mask_loss is False - assert config.dataset.num_dataset_builder_threads == 1 - assert config.dataset.data_sharding is True - assert config.dataset.dataloader_type == "single" - assert config.dataset.num_workers == 8 - - def test_pretrain_config_logger_configuration(self): - """Test logger configuration.""" - config = pretrain_config() - - assert config.logger.log_interval == 10 - assert "tb_logs" in config.logger.tensorboard_dir - - @pytest.mark.parametrize( - "tensor_parallelism,pipeline_parallelism,context_parallelism", - [ - (1, 1, 1), - (1, 4, 2), - (1, 2, 4), - (1, 2, 2), - (1, 4, 1), - ], - ) - def test_pretrain_config_parallelism_combinations( - self, tensor_parallelism, pipeline_parallelism, context_parallelism - ): - """Test various parallelism combinations.""" - config = pretrain_config( - tensor_parallelism=tensor_parallelism, - pipeline_parallelism=pipeline_parallelism, - context_parallelism=context_parallelism, - pipeline_parallelism_dtype=torch.bfloat16, - ) - - assert config.model.tensor_model_parallel_size == tensor_parallelism - assert config.model.pipeline_model_parallel_size == pipeline_parallelism - assert config.model.context_parallel_size == context_parallelism - - @pytest.mark.parametrize( - "global_batch_size,micro_batch_size", - [ - (8, 1), - (16, 2), - (32, 4), - (64, 8), - ], - ) - def test_pretrain_config_batch_size_combinations(self, global_batch_size, micro_batch_size): - """Test various batch size combinations.""" - config = pretrain_config(global_batch_size=global_batch_size, micro_batch_size=micro_batch_size) - - assert config.train.global_batch_size == global_batch_size - assert config.train.micro_batch_size == micro_batch_size - - @pytest.mark.parametrize("seq_length", [1024, 2048, 4096, 8192, 16384]) - def test_pretrain_config_sequence_lengths(self, seq_length): - """Test various sequence lengths.""" - config = pretrain_config(seq_length=seq_length) - - assert config.dataset.sequence_length == seq_length - - @pytest.mark.parametrize("precision", ["fp16_mixed", "bf16_mixed"]) - def test_precision_recipes(self, precision): - """Test precision configuration.""" - cfg = pretrain_config(precision_config=precision) - assert cfg.mixed_precision == precision diff --git a/tests/unit_tests/recipes/mamba/test_mamba2_8b.py b/tests/unit_tests/recipes/mamba/test_mamba2_8b.py deleted file mode 100644 index ea97cf9acf..0000000000 --- a/tests/unit_tests/recipes/mamba/test_mamba2_8b.py +++ /dev/null @@ -1,383 +0,0 @@ -# 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 -import tempfile -from unittest.mock import patch - -import pytest -import torch - -from megatron.bridge.models.mamba import NVIDIAMambaModelProvider8B -from megatron.bridge.recipes.mamba.mamba2_8b import model_config, pretrain_config -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 ConfigContainer - - -@pytest.mark.unit -class TestModelConfig: - """Test cases for the model_config function.""" - - def test_model_config_default_parameters(self): - """Test model_config with default parameters.""" - config = model_config() - - assert isinstance(config, NVIDIAMambaModelProvider8B) - assert config.tensor_model_parallel_size == 8 - assert config.pipeline_model_parallel_size == 1 - assert config.pipeline_dtype is None - assert config.virtual_pipeline_model_parallel_size is None - assert config.context_parallel_size == 1 - assert config.sequence_parallel is False - - def test_model_config_custom_tensor_parallelism(self): - """Test model_config with custom tensor parallelism.""" - config = model_config(tensor_parallelism=16) # Mamba 8B has 32 attention heads - - assert config.tensor_model_parallel_size == 16 - assert config.pipeline_model_parallel_size == 1 # default - assert config.context_parallel_size == 1 # default - - def test_model_config_custom_pipeline_parallelism(self): - """Test model_config with custom pipeline parallelism.""" - config = model_config(pipeline_parallelism=8, pipeline_parallelism_dtype=torch.float16) - - assert config.tensor_model_parallel_size == 8 # default - assert config.pipeline_model_parallel_size == 8 - assert config.pipeline_dtype is torch.float16 - - def test_model_config_with_pipeline_dtype(self): - """Test model_config with pipeline dtype specified.""" - config = model_config(pipeline_parallelism=2, pipeline_parallelism_dtype=torch.float16) - - assert config.pipeline_model_parallel_size == 2 - assert config.pipeline_dtype == torch.float16 - - def test_model_config_virtual_pipeline_parallelism(self): - """Test model_config with virtual pipeline parallelism.""" - config = model_config(virtual_pipeline_parallelism=4) - - assert config.virtual_pipeline_model_parallel_size == 4 - - def test_model_config_context_parallelism(self): - """Test model_config with custom context parallelism.""" - config = model_config(context_parallelism=8) - - assert config.context_parallel_size == 8 - - def test_model_config_sequence_parallelism_enabled(self): - """Test model_config with sequence parallelism enabled.""" - config = model_config(sequence_parallelism=True, tensor_parallelism=16) # Mamba 8B has 32 attention heads - - assert config.sequence_parallel is True - - def test_model_config_all_custom_parameters(self): - """Test model_config with all parameters customized.""" - config = model_config( - tensor_parallelism=16, # Mamba 8B has 32 attention heads - pipeline_parallelism=4, - pipeline_parallelism_dtype=torch.bfloat16, - virtual_pipeline_parallelism=8, - context_parallelism=16, - sequence_parallelism=True, - ) - - assert config.tensor_model_parallel_size == 16 - assert config.pipeline_model_parallel_size == 4 - assert config.pipeline_dtype == torch.bfloat16 - assert config.virtual_pipeline_model_parallel_size == 8 - assert config.context_parallel_size == 16 - assert config.sequence_parallel is True - - -@pytest.mark.unit -class TestPretrainConfig: - """Test cases for the pretrain_config function.""" - - def test_pretrain_config_default_parameters(self): - """Test pretrain_config with default parameters (mock mode).""" - config = pretrain_config() - - assert isinstance(config, ConfigContainer) - assert isinstance(config.model, NVIDIAMambaModelProvider8B) - - # Check training configuration - assert config.train.train_iters == 1_168_251 - assert config.train.global_batch_size == 8 - assert config.train.micro_batch_size == 1 - assert config.train.eval_interval == 100 - assert config.train.eval_iters == 32 - - # Check optimizer configuration - assert config.optimizer.optimizer == "adam" - assert config.optimizer.lr == 3e-4 - assert config.optimizer.min_lr == 3e-5 - assert config.optimizer.weight_decay == 0.1 - assert config.optimizer.bf16 is True - assert config.optimizer.fp16 is False - - # Check dataset configuration (should be in mock mode) - assert config.dataset.sequence_length == 4096 - assert config.dataset.split == "1,1,1" - assert config.dataset.blend is None - assert config.dataset.blend_per_split is None - - def test_pretrain_config_custom_training_parameters(self): - """Test pretrain_config with custom training parameters.""" - config = pretrain_config( - train_iters=10000, - global_batch_size=256, - micro_batch_size=2, - seq_length=4096, - lr=1e-4, - min_lr=1e-5, - lr_warmup_iters=1000, - ) - - assert config.train.train_iters == 10000 - assert config.train.global_batch_size == 256 - assert config.train.micro_batch_size == 2 - assert config.dataset.sequence_length == 4096 - assert config.optimizer.lr == 1e-4 - assert config.optimizer.min_lr == 1e-5 - assert config.scheduler.lr_warmup_iters == 1000 # Note: fixed in scheduler config - assert config.scheduler.lr_decay_iters is None # Will be set to train_iters during validation - - def test_pretrain_config_custom_model_parameters(self): - """Test pretrain_config with custom model parameters.""" - config = pretrain_config( - tensor_parallelism=16, # Mamba 8B has 32 attention heads - pipeline_parallelism=2, - context_parallelism=8, - sequence_parallelism=True, - pipeline_parallelism_dtype=torch.bfloat16, - ) - - assert config.model.tensor_model_parallel_size == 16 - assert config.model.pipeline_model_parallel_size == 2 - assert config.model.context_parallel_size == 8 - assert config.model.sequence_parallel is True - assert config.model.pipeline_dtype == torch.bfloat16 - - def test_pretrain_config_with_custom_directory(self): - """Test pretrain_config with custom directory.""" - with tempfile.TemporaryDirectory() as temp_dir: - config = pretrain_config(dir=temp_dir, name="test_run") - - expected_run_dir = os.path.join(temp_dir, "test_run") - expected_checkpoint_dir = os.path.join(expected_run_dir, "checkpoints") - expected_tensorboard_dir = os.path.join(expected_run_dir, "tb_logs") - - assert config.checkpoint.save == expected_checkpoint_dir - assert config.logger.tensorboard_dir == expected_tensorboard_dir - - def test_pretrain_config_explicit_mock_mode(self): - """Test pretrain_config with explicit mock=True.""" - config = pretrain_config(mock=True) - - assert config.dataset.blend is None - assert config.dataset.blend_per_split is None - assert config.dataset.split == "1,1,1" - - def test_pretrain_config_with_data_paths(self): - """Test pretrain_config with data paths provided.""" - - data_paths = ["/path/to/data1", "/path/to/data2", "/path/to/data3"] - config = pretrain_config(data_paths=data_paths) - - # Check that non-mock mode is configured - assert config.dataset.split == "9999,8,2" - assert config.dataset.blend is not None - assert config.dataset.blend_per_split is None - - def test_pretrain_config_with_train_valid_test_paths(self): - """Test pretrain_config with separate train/valid/test paths.""" - - config = pretrain_config( - train_data_path=["/path/to/train1", "/path/to/train2", "/path/to/train3"], - valid_data_path=["/path/to/valid1", "/path/to/valid2", "/path/to/valid3"], - test_data_path=["/path/to/test1", "/path/to/test2", "/path/to/test3"], - ) - - # When blend_per_split is used, split should be None - assert config.dataset.split is None - assert config.dataset.blend is None - assert config.dataset.blend_per_split is not None - - def test_pretrain_config_prioritizes_blend(self): - """Test that blend takes priority over blend_per_split when both are provided.""" - config = pretrain_config( - train_data_path=["/path/to/train1", "/path/to/train2"], - valid_data_path=["/path/to/valid1", "/path/to/valid2"], - test_data_path=["/path/to/test1", "/path/to/test2"], - data_paths=["/path/to/data1", "/path/to/data2"], - ) - - # Should prioritize blend over blend_per_split - assert config.dataset.split == "9999,8,2" - assert config.dataset.blend is not None - assert config.dataset.blend_per_split is None - - @patch("megatron.bridge.recipes.utils.dataset_utils.get_blend_and_blend_per_split") - def test_pretrain_config_fallback_to_mock_when_no_weights(self, mock_get_blend): - """Test pretrain_config falls back to mock when no weights are returned.""" - # Mock function returns None for both weights - mock_get_blend.return_value = (None, None) - - config = pretrain_config(data_paths=["/some/path"]) - - # Should fall back to mock mode - assert config.dataset.blend is None - assert config.dataset.blend_per_split is None - assert config.dataset.split == "1,1,1" - - def test_pretrain_config_checkpoint_configuration(self): - """Test checkpoint configuration in pretrain_config.""" - config = pretrain_config() - - assert config.checkpoint.save_interval == 2000 - assert config.checkpoint.ckpt_format == "torch_dist" - assert config.checkpoint.fully_parallel_load is True - - def test_pretrain_config_ddp_configuration(self): - """Test distributed data parallel configuration.""" - config = pretrain_config() - - assert config.ddp.check_for_nan_in_grad is True - assert config.ddp.grad_reduce_in_fp32 is True - assert config.ddp.overlap_grad_reduce is True - assert config.ddp.overlap_param_gather is True - assert config.ddp.use_distributed_optimizer is True - - def test_pretrain_config_default_comm_overlap(self): - """Test default CommOverlapConfig setup.""" - config = pretrain_config() - - # Default setup should have comm overlap config - assert config.comm_overlap is None - - def test_pretrain_config_custom_comm_overlap(self): - """Test custom CommOverlapConfig.""" - custom_overlap = CommOverlapConfig( - tp_comm_overlap=True, - defer_embedding_wgrad_compute=True, - wgrad_deferral_limit=50, - data_parallel_size=1, # Add this to avoid None - ) - config = pretrain_config(comm_overlap_config=custom_overlap) - - # Should use the custom config - assert config.comm_overlap is not None - - def test_pretrain_config_scheduler_configuration(self): - """Test scheduler configuration.""" - config = pretrain_config(train_iters=50000) - - assert config.scheduler.start_weight_decay == 0.033 - assert config.scheduler.end_weight_decay == 0.033 - assert config.scheduler.weight_decay_incr_style == "constant" - assert config.scheduler.lr_decay_style == "cosine" - assert config.scheduler.lr_warmup_iters == 2000 - assert config.scheduler.lr_warmup_init == 0.0 - assert config.scheduler.lr_decay_iters is None # Will be set to train_iters during validation - assert config.scheduler.override_opt_param_scheduler is True - - def test_pretrain_config_tokenizer_configuration(self): - """Test tokenizer configuration.""" - config = pretrain_config() - - assert config.tokenizer.tokenizer_type == "NullTokenizer" - assert config.tokenizer.vocab_size == DEFAULT_NULL_TOKENIZER_VOCAB_SIZE - - def test_pretrain_config_rng_configuration(self): - """Test RNG configuration.""" - config = pretrain_config() - - assert config.rng.seed == 1234 - assert config.dataset.random_seed == 1234 - - def test_pretrain_config_dataset_configuration(self): - """Test dataset configuration details.""" - config = pretrain_config() - - assert config.dataset.reset_attention_mask is False - assert config.dataset.reset_position_ids is False - assert config.dataset.eod_mask_loss is False - assert config.dataset.num_dataset_builder_threads == 1 - assert config.dataset.data_sharding is True - assert config.dataset.dataloader_type == "single" - assert config.dataset.num_workers == 8 - - def test_pretrain_config_logger_configuration(self): - """Test logger configuration.""" - config = pretrain_config() - - assert config.logger.log_interval == 10 - assert "tb_logs" in config.logger.tensorboard_dir - - @pytest.mark.parametrize( - "tensor_parallelism,pipeline_parallelism,context_parallelism", - [ - (1, 1, 1), - (2, 1, 4), - (4, 2, 2), - (8, 2, 2), # Default tensor parallelism - (16, 2, 1), # Higher tensor parallelism for 8B model - (32, 1, 1), # Maximum tensor parallelism for 8B model - ], - ) - def test_pretrain_config_parallelism_combinations( - self, tensor_parallelism, pipeline_parallelism, context_parallelism - ): - """Test various parallelism combinations.""" - config = pretrain_config( - tensor_parallelism=tensor_parallelism, - pipeline_parallelism=pipeline_parallelism, - context_parallelism=context_parallelism, - pipeline_parallelism_dtype=torch.bfloat16, - ) - - assert config.model.tensor_model_parallel_size == tensor_parallelism - assert config.model.pipeline_model_parallel_size == pipeline_parallelism - assert config.model.context_parallel_size == context_parallelism - - @pytest.mark.parametrize( - "global_batch_size,micro_batch_size", - [ - (128, 1), - (512, 2), - (1024, 4), - (256, 8), - ], - ) - def test_pretrain_config_batch_size_combinations(self, global_batch_size, micro_batch_size): - """Test various batch size combinations.""" - config = pretrain_config(global_batch_size=global_batch_size, micro_batch_size=micro_batch_size) - - assert config.train.global_batch_size == global_batch_size - assert config.train.micro_batch_size == micro_batch_size - - @pytest.mark.parametrize("seq_length", [1024, 2048, 4096, 8192, 16384]) - def test_pretrain_config_sequence_lengths(self, seq_length): - """Test various sequence lengths.""" - config = pretrain_config(seq_length=seq_length) - - assert config.dataset.sequence_length == seq_length - - @pytest.mark.parametrize("precision", ["fp16_mixed", "bf16_mixed"]) - def test_precision_recipes(self, precision): - """Test precision configuration.""" - cfg = pretrain_config(precision_config=precision) - assert cfg.mixed_precision == precision diff --git a/tests/unit_tests/recipes/mamba/test_mamba2_hybrid_8b.py b/tests/unit_tests/recipes/mamba/test_mamba2_hybrid_8b.py deleted file mode 100644 index ed6d0fd0c2..0000000000 --- a/tests/unit_tests/recipes/mamba/test_mamba2_hybrid_8b.py +++ /dev/null @@ -1,380 +0,0 @@ -# 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 -import tempfile -from unittest.mock import patch - -import pytest -import torch - -from megatron.bridge.models.mamba import NVIDIAMambaHybridProvider8B -from megatron.bridge.recipes.mamba.mamba2_hybrid_8b import model_config, pretrain_config -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 ConfigContainer - - -@pytest.mark.unit -class TestModelConfig: - """Test cases for the model_config function.""" - - def test_model_config_default_parameters(self): - """Test model_config with default parameters.""" - config = model_config() - - assert isinstance(config, NVIDIAMambaHybridProvider8B) - assert config.tensor_model_parallel_size == 8 - assert config.pipeline_model_parallel_size == 1 - assert config.pipeline_dtype is None - assert config.virtual_pipeline_model_parallel_size is None - assert config.context_parallel_size == 1 - assert config.sequence_parallel is False - - def test_model_config_custom_tensor_parallelism(self): - """Test model_config with custom tensor parallelism.""" - config = model_config(tensor_parallelism=4) - - assert config.tensor_model_parallel_size == 4 - assert config.pipeline_model_parallel_size == 1 # default - assert config.context_parallel_size == 1 # default - - def test_model_config_custom_pipeline_parallelism(self): - """Test model_config with custom pipeline parallelism.""" - config = model_config(pipeline_parallelism=8, pipeline_parallelism_dtype=torch.float16) - - assert config.tensor_model_parallel_size == 8 # default - assert config.pipeline_model_parallel_size == 8 - assert config.pipeline_dtype is torch.float16 - - def test_model_config_with_pipeline_dtype(self): - """Test model_config with pipeline dtype specified.""" - config = model_config(pipeline_parallelism=2, pipeline_parallelism_dtype=torch.float16) - - assert config.pipeline_model_parallel_size == 2 - assert config.pipeline_dtype == torch.float16 - - def test_model_config_virtual_pipeline_parallelism(self): - """Test model_config with virtual pipeline parallelism.""" - config = model_config(virtual_pipeline_parallelism=4) - - assert config.virtual_pipeline_model_parallel_size == 4 - - def test_model_config_context_parallelism(self): - """Test model_config with custom context parallelism.""" - config = model_config(context_parallelism=8) - - assert config.context_parallel_size == 8 - - def test_model_config_sequence_parallelism_enabled(self): - """Test model_config with sequence parallelism enabled.""" - config = model_config(sequence_parallelism=True, tensor_parallelism=4) - - assert config.sequence_parallel is True - - def test_model_config_all_custom_parameters(self): - """Test model_config with all parameters customized.""" - config = model_config( - tensor_parallelism=4, - pipeline_parallelism=4, - pipeline_parallelism_dtype=torch.bfloat16, - virtual_pipeline_parallelism=8, - context_parallelism=16, - sequence_parallelism=True, - ) - - assert config.tensor_model_parallel_size == 4 - assert config.pipeline_model_parallel_size == 4 - assert config.pipeline_dtype == torch.bfloat16 - assert config.virtual_pipeline_model_parallel_size == 8 - assert config.context_parallel_size == 16 - assert config.sequence_parallel is True - - -@pytest.mark.unit -class TestPretrainConfig: - """Test cases for the pretrain_config function.""" - - def test_pretrain_config_default_parameters(self): - """Test pretrain_config with default parameters (mock mode).""" - config = pretrain_config() - - assert isinstance(config, ConfigContainer) - assert isinstance(config.model, NVIDIAMambaHybridProvider8B) - - # Check training configuration - assert config.train.train_iters == 1_168_251 - assert config.train.global_batch_size == 8 - assert config.train.micro_batch_size == 1 - assert config.train.eval_interval == 100 - assert config.train.eval_iters == 32 - - # Check optimizer configuration - assert config.optimizer.optimizer == "adam" - assert config.optimizer.lr == 3e-4 - assert config.optimizer.min_lr == 3e-5 - assert config.optimizer.weight_decay == 0.1 - assert config.optimizer.bf16 is True - assert config.optimizer.fp16 is False - - # Check dataset configuration (should be in mock mode) - assert config.dataset.sequence_length == 4096 - assert config.dataset.split == "1,1,1" - assert config.dataset.blend is None - assert config.dataset.blend_per_split is None - - def test_pretrain_config_custom_training_parameters(self): - """Test pretrain_config with custom training parameters.""" - config = pretrain_config( - train_iters=10000, - global_batch_size=256, - micro_batch_size=2, - seq_length=4096, - lr=1e-4, - min_lr=1e-5, - lr_warmup_iters=1000, - ) - - assert config.train.train_iters == 10000 - assert config.train.global_batch_size == 256 - assert config.train.micro_batch_size == 2 - assert config.dataset.sequence_length == 4096 - assert config.optimizer.lr == 1e-4 - assert config.optimizer.min_lr == 1e-5 - assert config.scheduler.lr_warmup_iters == 1000 # Note: fixed in scheduler config - assert config.scheduler.lr_decay_iters is None # Will be set to train_iters during validation - - def test_pretrain_config_custom_model_parameters(self): - """Test pretrain_config with custom model parameters.""" - config = pretrain_config( - tensor_parallelism=4, - pipeline_parallelism=2, - context_parallelism=8, - sequence_parallelism=True, - pipeline_parallelism_dtype=torch.bfloat16, - ) - - assert config.model.tensor_model_parallel_size == 4 - assert config.model.pipeline_model_parallel_size == 2 - assert config.model.context_parallel_size == 8 - assert config.model.sequence_parallel is True - assert config.model.pipeline_dtype == torch.bfloat16 - - def test_pretrain_config_with_custom_directory(self): - """Test pretrain_config with custom directory.""" - with tempfile.TemporaryDirectory() as temp_dir: - config = pretrain_config(dir=temp_dir, name="test_run") - - expected_run_dir = os.path.join(temp_dir, "test_run") - expected_checkpoint_dir = os.path.join(expected_run_dir, "checkpoints") - expected_tensorboard_dir = os.path.join(expected_run_dir, "tb_logs") - - assert config.checkpoint.save == expected_checkpoint_dir - assert config.logger.tensorboard_dir == expected_tensorboard_dir - - def test_pretrain_config_explicit_mock_mode(self): - """Test pretrain_config with explicit mock=True.""" - config = pretrain_config(mock=True) - - assert config.dataset.blend is None - assert config.dataset.blend_per_split is None - assert config.dataset.split == "1,1,1" - - def test_pretrain_config_with_data_paths(self): - """Test pretrain_config with data paths provided.""" - - data_paths = ["/path/to/data1", "/path/to/data2", "/path/to/data3"] - config = pretrain_config(data_paths=data_paths) - - # Check that non-mock mode is configured - assert config.dataset.split == "9999,8,2" - assert config.dataset.blend is not None - assert config.dataset.blend_per_split is None - - def test_pretrain_config_with_train_valid_test_paths(self): - """Test pretrain_config with separate train/valid/test paths.""" - - config = pretrain_config( - train_data_path=["/path/to/train1", "/path/to/train2", "/path/to/train3"], - valid_data_path=["/path/to/valid1", "/path/to/valid2", "/path/to/valid3"], - test_data_path=["/path/to/test1", "/path/to/test2", "/path/to/test3"], - ) - - # When blend_per_split is used, split should be None - assert config.dataset.split is None - assert config.dataset.blend is None - assert config.dataset.blend_per_split is not None - - def test_pretrain_config_prioritizes_blend(self): - """Test that blend takes priority over blend_per_split when both are provided.""" - config = pretrain_config( - train_data_path=["/path/to/train1", "/path/to/train2"], - valid_data_path=["/path/to/valid1", "/path/to/valid2"], - test_data_path=["/path/to/test1", "/path/to/test2"], - data_paths=["/path/to/data1", "/path/to/data2"], - ) - - # Should prioritize blend over blend_per_split - assert config.dataset.split == "9999,8,2" - assert config.dataset.blend is not None - assert config.dataset.blend_per_split is None - - @patch("megatron.bridge.recipes.utils.dataset_utils.get_blend_and_blend_per_split") - def test_pretrain_config_fallback_to_mock_when_no_weights(self, mock_get_blend): - """Test pretrain_config falls back to mock when no weights are returned.""" - # Mock function returns None for both weights - mock_get_blend.return_value = (None, None) - - config = pretrain_config(data_paths=["/some/path"]) - - # Should fall back to mock mode - assert config.dataset.blend is None - assert config.dataset.blend_per_split is None - assert config.dataset.split == "1,1,1" - - def test_pretrain_config_checkpoint_configuration(self): - """Test checkpoint configuration in pretrain_config.""" - config = pretrain_config() - - assert config.checkpoint.save_interval == 2000 - assert config.checkpoint.ckpt_format == "torch_dist" - assert config.checkpoint.fully_parallel_load is True - - def test_pretrain_config_ddp_configuration(self): - """Test distributed data parallel configuration.""" - config = pretrain_config() - - assert config.ddp.check_for_nan_in_grad is True - assert config.ddp.grad_reduce_in_fp32 is True - assert config.ddp.overlap_grad_reduce is True - assert config.ddp.overlap_param_gather is True - assert config.ddp.use_distributed_optimizer is True - - def test_pretrain_config_default_comm_overlap(self): - """Test default CommOverlapConfig setup.""" - config = pretrain_config() - - # Default setup should have comm overlap config - assert config.comm_overlap is None - - def test_pretrain_config_custom_comm_overlap(self): - """Test custom CommOverlapConfig.""" - custom_overlap = CommOverlapConfig( - tp_comm_overlap=True, - defer_embedding_wgrad_compute=True, - wgrad_deferral_limit=50, - data_parallel_size=1, # Add this to avoid None - ) - config = pretrain_config(comm_overlap_config=custom_overlap) - - # Should use the custom config - assert config.comm_overlap is not None - - def test_pretrain_config_scheduler_configuration(self): - """Test scheduler configuration.""" - config = pretrain_config(train_iters=50000) - - assert config.scheduler.start_weight_decay == 0.033 - assert config.scheduler.end_weight_decay == 0.033 - assert config.scheduler.weight_decay_incr_style == "constant" - assert config.scheduler.lr_decay_style == "cosine" - assert config.scheduler.lr_warmup_iters == 2000 - assert config.scheduler.lr_warmup_init == 0.0 - assert config.scheduler.lr_decay_iters is None # Will be set to train_iters during validation - assert config.scheduler.override_opt_param_scheduler is True - - def test_pretrain_config_tokenizer_configuration(self): - """Test tokenizer configuration.""" - config = pretrain_config() - - assert config.tokenizer.tokenizer_type == "NullTokenizer" - assert config.tokenizer.vocab_size == DEFAULT_NULL_TOKENIZER_VOCAB_SIZE - - def test_pretrain_config_rng_configuration(self): - """Test RNG configuration.""" - config = pretrain_config() - - assert config.rng.seed == 1234 - assert config.dataset.random_seed == 1234 - - def test_pretrain_config_dataset_configuration(self): - """Test dataset configuration details.""" - config = pretrain_config() - - assert config.dataset.reset_attention_mask is False - assert config.dataset.reset_position_ids is False - assert config.dataset.eod_mask_loss is False - assert config.dataset.num_dataset_builder_threads == 1 - assert config.dataset.data_sharding is True - assert config.dataset.dataloader_type == "single" - assert config.dataset.num_workers == 8 - - def test_pretrain_config_logger_configuration(self): - """Test logger configuration.""" - config = pretrain_config() - - assert config.logger.log_interval == 10 - assert "tb_logs" in config.logger.tensorboard_dir - - @pytest.mark.parametrize( - "tensor_parallelism,pipeline_parallelism,context_parallelism", - [ - (1, 1, 1), - (2, 1, 4), - (4, 2, 2), - ], - ) - def test_pretrain_config_parallelism_combinations( - self, tensor_parallelism, pipeline_parallelism, context_parallelism - ): - """Test various parallelism combinations.""" - config = pretrain_config( - tensor_parallelism=tensor_parallelism, - pipeline_parallelism=pipeline_parallelism, - context_parallelism=context_parallelism, - pipeline_parallelism_dtype=torch.bfloat16, - ) - - assert config.model.tensor_model_parallel_size == tensor_parallelism - assert config.model.pipeline_model_parallel_size == pipeline_parallelism - assert config.model.context_parallel_size == context_parallelism - - @pytest.mark.parametrize( - "global_batch_size,micro_batch_size", - [ - (128, 1), - (512, 2), - (1024, 4), - (256, 8), - ], - ) - def test_pretrain_config_batch_size_combinations(self, global_batch_size, micro_batch_size): - """Test various batch size combinations.""" - config = pretrain_config(global_batch_size=global_batch_size, micro_batch_size=micro_batch_size) - - assert config.train.global_batch_size == global_batch_size - assert config.train.micro_batch_size == micro_batch_size - - @pytest.mark.parametrize("seq_length", [1024, 2048, 4096, 8192, 16384]) - def test_pretrain_config_sequence_lengths(self, seq_length): - """Test various sequence lengths.""" - config = pretrain_config(seq_length=seq_length) - - assert config.dataset.sequence_length == seq_length - - @pytest.mark.parametrize("precision", ["fp16_mixed", "bf16_mixed"]) - def test_precision_recipes(self, precision): - """Test precision configuration.""" - cfg = pretrain_config(precision_config=precision) - assert cfg.mixed_precision == precision