add mlperf llama3.1 8b pretrain recipes

Author: mollys

src/megatron/bridge/perf_recipes/llama/__init__.py

@@ -70,6 +70,9 @@
     llama3_70b_sft_32gpu_gb200_fp8mx_config,
 )
 from megatron.bridge.perf_recipes.llama.gb200.llama31 import (
+    llama31_8b_pretrain_8gpu_gb200_nvfp4_config,
+    llama31_8b_pretrain_72gpu_gb200_nvfp4_config,
+    llama31_8b_pretrain_512gpu_gb200_fp8cs_config,
     llama31_405b_pretrain_128gpu_gb200_bf16_config,
     llama31_405b_pretrain_128gpu_gb200_fp8cs_config,
     llama31_405b_pretrain_128gpu_gb200_fp8mx_config,
@@ -102,6 +105,9 @@
     llama3_70b_sft_32gpu_gb300_fp8mx_config,
 )
 from megatron.bridge.perf_recipes.llama.gb300.llama31 import (
+    llama31_8b_pretrain_8gpu_gb300_nvfp4_config,
+    llama31_8b_pretrain_72gpu_gb300_nvfp4_config,
+    llama31_8b_pretrain_512gpu_gb300_fp8cs_config,
     llama31_405b_pretrain_128gpu_gb300_bf16_config,
     llama31_405b_pretrain_128gpu_gb300_fp8cs_config,
     llama31_405b_pretrain_128gpu_gb300_fp8mx_config,

src/megatron/bridge/perf_recipes/llama/common.py

@@ -14,6 +14,8 @@
 # ruff: noqa: F401
 """Common helpers for llama performance recipes."""
 
+import torch
+
 from megatron.bridge.perf_recipes._common import (
     _benchmark_common,
     _enable_overlap_param_gather_with_optimizer_step,
@@ -56,3 +58,189 @@ def _llama_benchmark_common(cfg: ConfigContainer) -> None:
     cfg.model.moe_token_dispatcher_type = "alltoall"
     if cuda_graph_impl == "none" and te_rng_disabled:
         cfg.rng.te_rng_tracker = cfg.model.use_te_rng_tracker = False
+
+
+def _llama31_8b_common(cfg: ConfigContainer) -> None:
+    """Set common configurations for all Llama3.1 8B configs."""
+    cfg.optimizer.adam_beta2 = 0.95
+    cfg.optimizer.adam_eps = 1e-05
+    cfg.optimizer.bf16 = True
+    cfg.optimizer.fp8_recipe = 'tensorwise'
+    cfg.optimizer.overlap_param_gather = True
+    cfg.optimizer.params_dtype = torch.bfloat16
+    cfg.optimizer.use_distributed_optimizer = True
+    cfg.optimizer.weight_decay = 0.1
+    cfg.optimizer.bucket_size = 768000000
+    cfg.ddp.data_parallel_sharding_strategy = 'optim_grads_params'
+    cfg.ddp.delay_wgrad_compute = False
+    cfg.ddp.overlap_grad_reduce = True
+    cfg.ddp.overlap_param_gather = True
+    cfg.ddp.use_distributed_optimizer = True
+    cfg.ddp.bucket_size = 768000000
+    cfg.ddp.grad_reduce_in_fp32 = False
+    cfg.ddp.average_in_collective = False
+    cfg.model.add_bias_linear = False
+    cfg.model.apply_rope_fusion = True
+    cfg.model.attention_dropout = 0.0
+    cfg.model.gradient_accumulation_fusion = True
+    cfg.model.hidden_dropout = 0.0
+    cfg.model.hidden_size = 4096
+    cfg.model.kv_channels = 128
+    cfg.model.masked_softmax_fusion = True
+    cfg.model.microbatch_group_size_per_vp_stage = 1
+    cfg.model.normalization = 'RMSNorm'
+    cfg.model.num_attention_heads = 32
+    cfg.model.num_layers = 32
+    cfg.model.num_layers_at_end_in_bf16 = 0
+    cfg.model.num_layers_at_start_in_bf16 = 0
+    cfg.model.num_query_groups = 8
+    cfg.model.overlap_p2p_comm = True
+    cfg.model.params_dtype = torch.bfloat16
+    cfg.model.persist_layer_norm = True
+    cfg.model.pipeline_dtype = torch.bfloat16
+    cfg.model.recompute_modules = ['core_attn']
+    cfg.model.tp_only_amax_red = True
+    cfg.model.use_te_rng_tracker = True
+    cfg.model.wgrad_deferral_limit = 50
+    cfg.model.batch_p2p_comm = False
+    cfg.model.fused_single_qkv_rope = True
+    cfg.train.check_optimizer_step_success = False
+    cfg.train.decrease_batch_size_if_needed = False
+    cfg.train.empty_unused_memory_level = 0
+    cfg.train.exit_signal_handler = False
+    cfg.train.exit_signal_handler_for_dataloader = False
+    cfg.train.exit_signal_handler_for_training = False
+    cfg.train.iterations_to_skip = []
+    cfg.train.manual_gc = True
+    cfg.train.manual_gc_eval = False
+    cfg.train.manual_gc_interval = 500
+    cfg.train.skip_sync_grad_norm_across_mp = True
+    cfg.scheduler.end_weight_decay = 0.1
+    cfg.scheduler.lr_decay_style = 'cosine'
+    cfg.scheduler.lr_warmup_init = 0.0
+    cfg.scheduler.lr_warmup_samples = 0
+    cfg.scheduler.lr_wsd_decay_style = 'exponential'
+    cfg.scheduler.override_opt_param_scheduler = False
+    cfg.scheduler.start_weight_decay = 0.1
+    cfg.scheduler.use_checkpoint_opt_param_scheduler = False
+    cfg.scheduler.weight_decay_incr_style = 'constant'
+    cfg.dataset.add_extra_token_to_sequence = True
+    cfg.dataset.allow_ambiguous_pad_tokens = False
+    cfg.dataset.create_attention_mask = False
+    cfg.dataset.data_parallel_size = 1
+    cfg.dataset.data_sharding = True
+    cfg.dataset.dataloader_type = 'single'
+    cfg.dataset.defer_npy_index_mmap = True
+    cfg.dataset.drop_last = True
+    cfg.dataset.drop_last_partial_validation_sequence = True
+    cfg.dataset.eod_mask_loss = False
+    cfg.dataset.fast_cache_load = True
+    cfg.dataset.hybrid_context_parallel = False
+    cfg.dataset.inter_document_masking = False
+    cfg.dataset.mid_level_dataset_surplus = 0.005
+    cfg.dataset.mmap_bin_files = True
+    cfg.dataset.mock = False
+    cfg.dataset.num_dataset_builder_threads = 1
+    cfg.dataset.num_workers = 8
+    cfg.dataset.persistent_workers = True
+    cfg.dataset.pin_memory = True
+    cfg.dataset.reset_attention_mask = False
+    cfg.dataset.reset_position_ids = False
+    cfg.dataset.seq_length = 8192
+    cfg.dataset.sequence_length = 8192
+    cfg.dataset.sequence_parallel_size = 0
+    cfg.checkpoint.also_save_hf_checkpoint = False
+    cfg.checkpoint.async_ckpt_cpu_priority = 10
+    cfg.checkpoint.async_ckpt_io_priority = 3
+    cfg.checkpoint.async_ckpt_use_cpu_shm = False
+    cfg.checkpoint.async_save = False
+    cfg.checkpoint.async_strategy = 'mcore'
+    cfg.checkpoint.async_write_results_mp_mode = 'fork'
+    cfg.checkpoint.auto_detect_ckpt_format = False
+    cfg.checkpoint.ckpt_assume_constant_structure = False
+    cfg.checkpoint.ckpt_convert_update_legacy_dist_opt_format = False
+    cfg.checkpoint.ckpt_format = 'torch_dist'
+    cfg.checkpoint.ckpt_fully_parallel_load_exchange_algo = 'broadcast'
+    cfg.checkpoint.ckpt_fully_parallel_load_process_group = 'dp'
+    cfg.checkpoint.ckpt_fully_parallel_save_process_group = 'dp'
+    cfg.checkpoint.ckpt_load_validate_sharding_integrity = True
+    cfg.checkpoint.dist_ckpt_optim_fully_reshardable = False
+    cfg.checkpoint.dist_ckpt_save_pre_mcore_014 = False
+    cfg.checkpoint.dist_ckpt_strictness = 'log_all'
+    cfg.checkpoint.dist_ckpt_workers = 1
+    cfg.checkpoint.distrib_optim_fully_reshardable_mem_efficient = False
+    cfg.checkpoint.exit_on_missing_checkpoint = False
+    cfg.checkpoint.finetune = False
+    cfg.checkpoint.fully_parallel_load = True
+    cfg.checkpoint.fully_parallel_save = True
+    cfg.checkpoint.hf_distributed_save = False
+    cfg.checkpoint.hf_save_every_n_ranks = 1
+    cfg.checkpoint.hf_trust_remote_code = False
+    cfg.checkpoint.load_optim = False
+    cfg.checkpoint.load_rng = False
+    cfg.checkpoint.most_recent_k = -1
+    cfg.checkpoint.non_persistent_local_ckpt_algo = 'fully_parallel'
+    cfg.checkpoint.replication = False
+    cfg.checkpoint.replication_factor = 2
+    cfg.checkpoint.save_optim = True
+    cfg.checkpoint.save_rng = True
+    cfg.checkpoint.save_tokenizer_assets = True
+    cfg.checkpoint.storage_writers_per_rank = 1
+    cfg.checkpoint.strict_fsdp_dtensor_load = False
+    cfg.checkpoint.use_checkpoint_args = False
+    cfg.checkpoint.use_mp_args_from_checkpoint_args = False
+    cfg.checkpoint.use_persistent_ckpt_worker = True
+    cfg.checkpoint.use_tokenizer_model_from_checkpoint_args = True
+    cfg.checkpoint.verify_integrity = False
+    cfg.logger.barrier_with_L1_time = True
+    cfg.logger.filter_warnings = True
+    cfg.logger.log_device_memory_used = False
+    cfg.logger.log_energy = False
+    cfg.logger.log_interval = 1
+    cfg.logger.log_l2_norm_grad_to_tensorboard = False
+    cfg.logger.log_loss_scale_to_tensorboard = True
+    cfg.logger.log_max_attention_logit = False
+    cfg.logger.log_memory_to_tensorboard = False
+    cfg.logger.log_num_zeros_in_grad = False
+    cfg.logger.log_params_norm = False
+    cfg.logger.log_progress = False
+    cfg.logger.log_runtime_to_tensorboard = False
+    cfg.logger.log_throughput = False
+    cfg.logger.log_throughput_to_tensorboard = False
+    cfg.logger.log_timers_to_tensorboard = False
+    cfg.logger.log_validation_ppl_to_tensorboard = False
+    cfg.logger.log_world_size_to_tensorboard = False
+    cfg.logger.logging_level = 20
+    cfg.logger.mlflow_log_artifacts = True
+    cfg.logger.moe_routing_trace_capture_hidden_states = False
+    cfg.logger.moe_routing_trace_capture_logits = False
+    cfg.logger.moe_routing_trace_dump_weights = False
+    cfg.logger.runtime_time_unit = 'hours'
+    cfg.logger.set_level_for_all_loggers = False
+    cfg.logger.skip_train_metrics_log = False
+    cfg.logger.tensorboard_log_interval = 1
+    cfg.logger.tensorboard_queue_size = 1000
+    cfg.logger.throughput_window_size = 100
+    cfg.logger.timing_log_level = 0
+    cfg.logger.timing_log_option = 'minmax'
+    cfg.tokenizer.force_system_message = False
+    cfg.tokenizer.hf_tokenizer_kwargs = {'use_fast': True}
+    cfg.tokenizer.make_vocab_size_divisible_by = 1
+    cfg.tokenizer.null_tokenizer_pad_id = -1
+    cfg.tokenizer.pad_vocab_size = False
+    cfg.tokenizer.rank = 0
+    cfg.tokenizer.sp_tokenizer_kwargs = {}
+    cfg.tokenizer.tensor_model_parallel_size = 1
+    cfg.tokenizer.tiktoken_num_special_tokens = 1000
+    cfg.tokenizer.tokenizer_hf_no_include_special_tokens = False
+    cfg.tokenizer.tokenizer_hf_no_use_fast = False
+    cfg.tokenizer.tokenizer_sentencepiece_ignore_extra_whitespaces = True
+    cfg.tokenizer.tokenizer_sentencepiece_legacy = False
+    cfg.tokenizer.tokenizer_type = "NullTokenizer"
+    cfg.tokenizer.trust_remote_code = False
+    cfg.tokenizer.vocab_extra_ids = 0
+    cfg.rng.data_parallel_random_init: False
+    cfg.rng.inference_rng_tracker = False
+    cfg.rng.te_rng_tracker = True
+    cfg.ddp.align_param_gather = False
+    cfg.optimizer.overlap_param_gather_with_optimizer_step = False