[dev] ci: Add DSv3 proxy

tests/functional_tests/test_cases/bert/bert_release/model_config.yaml

@@ -45,6 +45,7 @@ MODEL_ARGS:
   --log-params-norm: true
   --log-validation-ppl-to-tensorboard: true
   --wandb-project: megatron-core-release-runs
+  --wandb-entity: adlr
   --wandb-exp-name: ${WANDB_EXPERIMENT}
   --attention-backend: unfused
   --exit-interval: 20000

tests/functional_tests/test_cases/gpt/gpt3_15b_8t_release/model_config.yaml

@@ -84,6 +84,7 @@ MODEL_ARGS:
   --log-interval: 100
   --tensorboard-dir: ${TENSORBOARD_PATH}
   --wandb-project: megatron-core-release-runs
+  --wandb-entity: adlr
   --wandb-exp-name: ${WANDB_EXPERIMENT}
   # Add mixed precision args
   --bf16: true

tests/functional_tests/test_cases/gpt/gpt3_15b_8t_release_sm/model_config.yaml

@@ -84,6 +84,7 @@ MODEL_ARGS:
   --log-interval: 100
   --tensorboard-dir: ${TENSORBOARD_PATH}
   --wandb-project: megatron-core-release-runs
+  --wandb-entity: adlr
   --wandb-exp-name: ${WANDB_EXPERIMENT}
   # Add mixed precision args
   --bf16: true

tests/functional_tests/test_cases/mixtral/deepseekv3_proxy_flex_tp1pp4emp16etc1cp1_gb_200_release_sm/model_config.yml

@@ -0,0 +1,169 @@
+# The proxy model is used for local code quality check.
+# The proxy model should contain all the necessary components and settings but fewer parameters.
+ENV_VARS:
+  TORCH_NCCL_AVOID_RECORD_STREAMS: 0
+  NVTE_ALLOW_NONDETERMINISTIC_ALGO: 1
+  PYTORCH_CUDA_ALLOC_CONF: expandable_segments:True
+  NCCL_NVLS_ENABLE: 0
+  NVTE_FUSED_ATTN: 1
+  NVTE_NORM_FWD_USE_CUDNN: 1
+  NVTE_NORM_BWD_USE_CUDNN: 1
+  PYTHONWARNINGS: ignore
+  NCCL_DEBUG: VERSION
+  NON_DETERMINSTIC_RESULTS: 1
+  NVSHMEM_IB_ENABLE_IBGDA: 0
+  CUDA_DEVICE_MAX_CONNECTIONS: 1
+  NUM_OF_HYBRID_EP_RANKS_PER_NVLINK_DOMAIN: 16
+  USE_MNNVL: 1
+TEST_TYPE: "release"
+MODEL_ARGS:
+  # Distributed args
+  --distributed-timeout-minutes: 60
+  --tensor-model-parallel-size: 2
+  --pipeline-model-parallel-size: 4
+  --pipeline-model-parallel-layout: Et*2\\|\\(tt\\|\\)*5t\\|tmL # Et*2|(tt|)*5t|tmL
+  --expert-model-parallel-size: 16
+  --context-parallel-size: 1
+  --expert-tensor-parallel-size: 1
+  --use-distributed-optimizer: true
+  --overlap-grad-reduce: true
+  --overlap-param-gather: true
+
+  # Training args
+  --use-mcore-models: true
+  --sequence-parallel: true
+  --use-flash-attn: true
+  --disable-bias-linear: true
+  --micro-batch-size: 1
+  --global-batch-size: 512
+  --train-samples: 24414062
+  --exit-duration-in-mins: 220
+  --no-check-for-nan-in-loss-and-grad: true
+  --cross-entropy-loss-fusion: true
+  --cross-entropy-fusion-impl: te
+  --manual-gc: true
+  --manual-gc-interval: 10
+
+  # Transformer Engine args
+  --transformer-impl: transformer_engine
+
+  # Data args
+  --seq-length: 4096
+  --data-cache-path: ${DATA_CACHE_PATH}
+  --tokenizer-type: GPTSentencePieceTokenizer
+  --tokenizer-model: ${DATA_PATH}/utils/nemotron_2_256k.model
+  --data-path: $DATA_BLEND
+  --split: 99,1,0
+  --no-mmap-bin-files: true
+  --no-create-attention-mask-in-dataloader: true
+  --num-workers: 6
+
+  # Add network size args
+  --num-layers: 14 # original 61 layers
+  --hidden-size: 7168
+  --ffn-hidden-size: 18432
+  --num-attention-heads: 128
+  --kv-channels: 128
+  --max-position-embeddings: 4096
+  --position-embedding-type: rope
+  --rotary-base: 10000
+  --make-vocab-size-divisible-by: 3232
+  --normalization: RMSNorm
+  --norm-epsilon: 1e-6
+  --swiglu: true
+  --untie-embeddings-and-output-weights: true
+  --multi-latent-attention: true
+  --mtp-num-layers: 1
+  --mtp-loss-scaling-factor: 0.1
+
+  # Add regularization args
+  --attention-dropout: 0.0
+  --hidden-dropout: 0.0
+  --clip-grad: 1.0
+  --weight-decay: 0.1
+  --qk-layernorm: true
+
+  # Add learning rate args
+  --lr-decay-samples: 24413696
+  --lr-warmup-samples: 1536000
+  --lr-warmup-init: 1e-7
+  --lr: 1e-5
+  --min-lr: 1e-6
+  --lr-decay-style: cosine
+  --adam-beta1: 0.9
+  --adam-beta2: 0.95
+
+  # Add MoE args
+  --num-experts: 64 # local 4 + 1 shared, EP16
+  --moe-layer-freq: ([0]*3+[1]*11)
+  --moe-ffn-hidden-size: 2048
+  --moe-shared-expert-intermediate-size: 2048
+  --moe-router-load-balancing-type: seq_aux_loss
+  --moe-router-topk: 8
+  --moe-token-dispatcher-type: flex
+  --moe-flex-dispatcher-backend: hybridep
+  --moe-router-pre-softmax: true
+  --moe-grouped-gemm: true
+  --moe-aux-loss-coeff: 1e-4
+  --moe-router-group-topk: 4
+  --moe-router-num-groups: 8
+  --moe-router-topk-scaling-factor: 2.5
+  --moe-router-score-function: sigmoid
+  --moe-router-enable-expert-bias: true
+  --moe-router-bias-update-rate: 1e-3
+  --moe-router-dtype: fp32
+  --moe-permute-fusion: true
+
+  # Add MLA args
+  --q-lora-rank: 1536
+  --kv-lora-rank: 512
+  --qk-head-dim: 128
+  --qk-pos-emb-head-dim: 64
+  --v-head-dim: 128
+  --rotary-scaling-factor: 40
+  --mscale: 1.0
+  --mscale-all-dim: 1.0
+
+  # Add validation args
+  --eval-iters: 32
+  --eval-interval: 200
+
+  # Add checkpointing args
+  --auto-detect-ckpt-format:
+    true
+    # Add checkpointing args
+  --save: ${CHECKPOINT_SAVE_PATH}
+  --load: ${CHECKPOINT_LOAD_PATH}
+  --save-interval: 500
+  --save-retain-interval: 10000
+  --dist-ckpt-strictness: log_all
+
+  # Add initialization args
+  --init-method-std: 0.02
+
+  # Add logging args
+  --log-timers-to-tensorboard: true
+  --log-memory-to-tensorboard: true
+  --log-num-zeros-in-grad: true
+  --log-params-norm: true
+  --log-validation-ppl-to-tensorboard: true
+  --log-throughput: true
+  --log-interval: 1
+  --logging-level: 40
+  --tensorboard-dir: ${TENSORBOARD_PATH}
+  --wandb-project: megatron-core-release-runs
+  --wandb-entity: adlr
+  --wandb-exp-name: ${WANDB_EXPERIMENT}
+  --wandb-save-dir: ${WANDB_SAVE_PATH}
+
+  # Add mixed precision args
+  --bf16: true
+
+  # enable experimental
+  --enable-experimental: true
+  --exit-interval: 9536
+METRICS:
+  - "iteration-time"
+  - "lm loss"
+  - "mem-allocated-bytes"
+  - "mem-max-allocated-bytes"

tests/functional_tests/test_cases/mixtral/deepseekv3_proxy_flex_tp1pp4emp16etc1cp1_release_sm/model_config.yml

@@ -13,7 +13,7 @@ ENV_VARS:
   NON_DETERMINSTIC_RESULTS: 1
   NVSHMEM_IB_ENABLE_IBGDA: 0
   CUDA_DEVICE_MAX_CONNECTIONS: 1
-TEST_TYPE: 'release'
+TEST_TYPE: "release"
 MODEL_ARGS:
   # Distributed args
   --distributed-timeout-minutes: 60
@@ -150,6 +150,7 @@ MODEL_ARGS:
   --logging-level: 40
   --tensorboard-dir: ${TENSORBOARD_PATH}
   --wandb-project: megatron-core-release-runs
+  --wandb-entity: adlr
   --wandb-exp-name: ${WANDB_EXPERIMENT}
   --wandb-save-dir: ${WANDB_SAVE_PATH}
 
@@ -160,7 +161,7 @@ MODEL_ARGS:
   --enable-experimental: true
   --exit-interval: 9536
 METRICS:
-  - 'iteration-time'
-  - 'lm loss'
-  - 'mem-allocated-bytes'
-  - 'mem-max-allocated-bytes'
+  - "iteration-time"
+  - "lm loss"
+  - "mem-allocated-bytes"
+  - "mem-max-allocated-bytes"

tests/functional_tests/test_cases/mixtral/deepseekv3_proxy_flex_tp1pp4emp16etp1cp1_gb_200_release/model_config.yaml

@@ -0,0 +1,168 @@
+# The proxy model is used for local code quality check.
+# The proxy model should contain all the necessary components and settings but fewer parameters.
+ENV_VARS:
+  TORCH_NCCL_AVOID_RECORD_STREAMS: 0
+  NVTE_ALLOW_NONDETERMINISTIC_ALGO: 1
+  PYTORCH_CUDA_ALLOC_CONF: expandable_segments:True
+  NCCL_NVLS_ENABLE: 0
+  NVTE_FUSED_ATTN: 1
+  NVTE_NORM_FWD_USE_CUDNN: 1
+  NVTE_NORM_BWD_USE_CUDNN: 1
+  PYTHONWARNINGS: ignore
+  NCCL_DEBUG: VERSION
+  NON_DETERMINSTIC_RESULTS: 1
+  NVSHMEM_IB_ENABLE_IBGDA: 0
+  CUDA_DEVICE_MAX_CONNECTIONS: 1
+  NUM_OF_HYBRID_EP_RANKS_PER_NVLINK_DOMAIN: 16
+  USE_MNNVL: 1
+TEST_TYPE: "release"
+MODEL_ARGS:
+  # Distributed args
+  --distributed-timeout-minutes: 60
+  --tensor-model-parallel-size: 2
+  --pipeline-model-parallel-size: 4
+  --pipeline-model-parallel-layout: Et*2\\|\\(tt\\|\\)*5t\\|tmL # Et*2|(tt|)*5t|tmL
+  --expert-model-parallel-size: 16
+  --context-parallel-size: 1
+  --expert-tensor-parallel-size: 1
+  --use-distributed-optimizer: true
+  --overlap-grad-reduce: true
+  --overlap-param-gather: true
+
+  # Training args
+  --use-mcore-models: true
+  --sequence-parallel: true
+  --use-flash-attn: true
+  --disable-bias-linear: true
+  --micro-batch-size: 1
+  --global-batch-size: 512
+  --train-samples: 24414062
+  --exit-duration-in-mins: 220
+  --no-check-for-nan-in-loss-and-grad: true
+  --cross-entropy-loss-fusion: true
+  --cross-entropy-fusion-impl: te
+  --manual-gc: true
+  --manual-gc-interval: 10
+
+  # Transformer Engine args
+  --transformer-impl: transformer_engine
+
+  # Data args
+  --seq-length: 4096
+  --data-cache-path: ${DATA_CACHE_PATH}
+  --tokenizer-type: GPTSentencePieceTokenizer
+  --tokenizer-model: ${DATA_PATH}/utils/nemotron_2_256k.model
+  --data-path: $DATA_BLEND
+  --split: 99,1,0
+  --no-mmap-bin-files: true
+  --no-create-attention-mask-in-dataloader: true
+  --num-workers: 6
+
+  # Add network size args
+  --num-layers: 14 # original 61 layers
+  --hidden-size: 7168
+  --ffn-hidden-size: 18432
+  --num-attention-heads: 128
+  --kv-channels: 128
+  --max-position-embeddings: 4096
+  --position-embedding-type: rope
+  --rotary-base: 10000
+  --make-vocab-size-divisible-by: 3232
+  --normalization: RMSNorm
+  --norm-epsilon: 1e-6
+  --swiglu: true
+  --untie-embeddings-and-output-weights: true
+  --multi-latent-attention: true
+  --mtp-num-layers: 1
+  --mtp-loss-scaling-factor: 0.1
+
+  # Add regularization args
+  --attention-dropout: 0.0
+  --hidden-dropout: 0.0
+  --clip-grad: 1.0
+  --weight-decay: 0.1
+  --qk-layernorm: true
+
+  # Add learning rate args
+  --lr-decay-samples: 24413696
+  --lr-warmup-samples: 1536000
+  --lr-warmup-init: 1e-7
+  --lr: 1e-5
+  --min-lr: 1e-6
+  --lr-decay-style: cosine
+  --adam-beta1: 0.9
+  --adam-beta2: 0.95
+
+  # Add MoE args
+  --num-experts: 64 # local 4 + 1 shared, EP16
+  --moe-layer-freq: ([0]*3+[1]*11)
+  --moe-ffn-hidden-size: 2048
+  --moe-shared-expert-intermediate-size: 2048
+  --moe-router-load-balancing-type: seq_aux_loss
+  --moe-router-topk: 8
+  --moe-token-dispatcher-type: flex
+  --moe-flex-dispatcher-backend: hybridep
+  --moe-router-pre-softmax: true
+  --moe-grouped-gemm: true
+  --moe-aux-loss-coeff: 1e-4
+  --moe-router-group-topk: 4
+  --moe-router-num-groups: 8
+  --moe-router-topk-scaling-factor: 2.5
+  --moe-router-score-function: sigmoid
+  --moe-router-enable-expert-bias: true
+  --moe-router-bias-update-rate: 1e-3
+  --moe-router-dtype: fp32
+  --moe-permute-fusion: true
+
+  # Add MLA args
+  --q-lora-rank: 1536
+  --kv-lora-rank: 512
+  --qk-head-dim: 128
+  --qk-pos-emb-head-dim: 64
+  --v-head-dim: 128
+  --rotary-scaling-factor: 40
+  --mscale: 1.0
+  --mscale-all-dim: 1.0
+
+  # Add validation args
+  --eval-iters: 32
+  --eval-interval: 200
+
+  # Add checkpointing args
+  --auto-detect-ckpt-format:
+    true
+    # Add checkpointing args
+  --save: ${CHECKPOINT_SAVE_PATH}
+  --load: ${CHECKPOINT_LOAD_PATH}
+  --save-interval: 500
+  --save-retain-interval: 10000
+  --dist-ckpt-strictness: log_all
+
+  # Add initialization args
+  --init-method-std: 0.02
+
+  # Add logging args
+  --log-timers-to-tensorboard: true
+  --log-memory-to-tensorboard: true
+  --log-num-zeros-in-grad: true
+  --log-params-norm: true
+  --log-validation-ppl-to-tensorboard: true
+  --log-throughput: true
+  --log-interval: 1
+  --logging-level: 40
+  --tensorboard-dir: ${TENSORBOARD_PATH}
+  --wandb-project: megatron-core-release-runs
+  --wandb-entity: adlr
+  --wandb-exp-name: ${WANDB_EXPERIMENT}
+  --wandb-save-dir: ${WANDB_SAVE_PATH}
+
+  # Add mixed precision args
+  --bf16: true
+
+  # enable experimental
+  --enable-experimental: true
+METRICS:
+  - "iteration-time"
+  - "lm loss"
+  - "mem-allocated-bytes"
+  - "mem-max-allocated-bytes"

tests/functional_tests/test_cases/mixtral/deepseekv3_proxy_flex_tp1pp4emp16etp1cp1_release/model_config.yaml

@@ -150,6 +150,7 @@ MODEL_ARGS:
   --logging-level: 40
   --tensorboard-dir: ${TENSORBOARD_PATH}
   --wandb-project: megatron-core-release-runs
+  --wandb-entity: adlr
   --wandb-exp-name: ${WANDB_EXPERIMENT}
   --wandb-save-dir: ${WANDB_SAVE_PATH}