[Draft] 355_benchmark

examples/megatron/configs/MI355X/deepseek_v3-BF16-pretrain.yaml

@@ -40,7 +40,7 @@ modules:
 
       # parallel
       tensor_model_parallel_size: ${PRIMUS_TP:1}
-      pipeline_model_parallel_size: ${PRIMUS_PP:1}
+      pipeline_model_parallel_size: ${PRIMUS_PP:8}
       expert_model_parallel_size: ${PRIMUS_EP:8}
       overlap_grad_reduce: true
       overlap_param_gather: true
@@ -71,6 +71,24 @@ modules:
       ckpt_format: torch
       eval_iters: 0
 
+      # Turbo
+      enable_primus_turbo: true
+      use_turbo_attention: false
+      use_turbo_grouped_mlp: false
+
+      # deepep
+      use_turbo_deepep: true
+      moe_shared_expert_overlap: false
+      moe_router_dtype: fp32
+
+      # 64 or 80 for ep8, 32 for ep16-64 is best practice
+      turbo_deepep_num_cu: 64
+      turbo_deepep_use_comm_stream: false
+
+      # sync-free moe support stage 1-2, 0 means not use sync-free moe
+      # stage 2 is recommended for better performance
+      turbo_sync_free_moe_stage: 1
+
       # Cross entropy flags
       # cross_entropy_fusion_impl: "te"
       # cross_entropy_loss_fusion: true

examples/megatron/prepare.py

@@ -264,7 +264,9 @@ def build_megatron_helper(primus_path: Path, patch_args: Path, backend_path: str
 
     emerging_optimizers_path = primus_path / "third_party/Emerging-Optimizers"
     log_info(f"Building Emerging Optimizers in {emerging_optimizers_path}")
-    ret = subprocess.run(["pip", "install", "-e", str(emerging_optimizers_path)], check=True)
+    ret = subprocess.run(
+        ["pip", "install", "--no-build-isolation", "-e", str(emerging_optimizers_path)], check=True
+    )
     if ret.returncode != 0:
         log_error_and_exit("Building Emerging Optimizers failed.")
 

examples/run_local_pretrain.sh

@@ -113,7 +113,7 @@ echo "ENV_ARGS: ${ENV_ARGS[*]}"
 HOSTNAME=$(hostname)
 ARGS=("$@")
 
-VOLUME_ARGS=(-v "$PRIMUS_PATH":"$PRIMUS_PATH" -v "$DATA_PATH":"$DATA_PATH")
+VOLUME_ARGS=(-v "$PRIMUS_PATH":"$PRIMUS_PATH" -v "$DATA_PATH":"$DATA_PATH" -v "/shared_aig/c4:/shared_aig/c4")
 if [[ -f "$PATH_TO_BNXT_TAR_PACKAGE" ]]; then
     VOLUME_ARGS+=(-v "$PATH_TO_BNXT_TAR_PACKAGE":"$PATH_TO_BNXT_TAR_PACKAGE")
 fi
@@ -134,10 +134,10 @@ export CLEAN_DOCKER_CONTAINER=${CLEAN_DOCKER_CONTAINER:-0}
 
 # ------------------ Optional Container Cleanup ------------------
 docker_podman_proxy() {
-    if command -v podman &>/dev/null; then
-        podman "$@"
-    elif command -v docker &>/dev/null; then
+    if command -v docker &>/dev/null; then
         docker "$@"
+    elif command -v podman &>/dev/null; then
+        podman "$@"
     else
         echo "Neither Docker nor Podman found!" >&2
         return 1
@@ -164,6 +164,13 @@ else
     echo "Node-${NODE_RANK}: Launching training container."
 fi
 
+if ! docker_podman_proxy image inspect "$DOCKER_IMAGE" &>/dev/null; then
+    echo "Node-${NODE_RANK}: Image not found locally, pulling $DOCKER_IMAGE..."
+    docker_podman_proxy pull "$DOCKER_IMAGE"
+else
+    echo "Node-${NODE_RANK}: Image $DOCKER_IMAGE already exists, skipping pull."
+fi
+
 # ------------------ Launch Training Container ------------------
 docker_podman_proxy run --rm \
     --env MASTER_ADDR \

examples/run_pretrain.sh

@@ -197,8 +197,10 @@ if [ "$USING_AINIC" == "1" ]; then
     export NCCL_IB_GID_INDEX=1
     # export NCCL_IB_ROCE_VERSION_NUM=2
     export NCCL_MAX_P2P_CHANNELS=56
-    export NCCL_IB_TC=104
-    export NCCL_IB_FIFO_TC=192
+    # export NCCL_IB_TC=104
+    # export NCCL_IB_FIFO_TC=192
+    export NCCL_IB_TC=41
+    export NCCL_IB_FIFO_TC=185
     export NET_OPTIONAL_RECV_COMPLETION=1
     export NCCL_IB_USE_INLINE=1
     export RCCL_GDR_FLUSH_GPU_MEM_NO_RELAXED_ORDERING=0

examples/run_slurm_pretrain.sh

@@ -38,10 +38,14 @@ export LOG_DIR=${LOG_DIR:-"./output"}
 LOG_FILE="${LOG_DIR}/log_slurm_pretrain.txt"
 mkdir -p "$LOG_DIR"
 
+    #  --nodelist="uswslocpm2m-106-[273,297,310,319,687,732,836,892]" \
 srun -N "${NNODES}" \
      --exclusive \
      --export ALL \
      --ntasks-per-node=1 \
+     --time="${SLURM_TIME:-07:00:00}" \
+     --nodelist="${SLURM_NODELIST:-}" \
+     --partition="${SLURM_PARTITION:-amd-aig}" \
      --cpus-per-task="${CPUS_PER_TASK:-128}" \
      bash -c "
           readarray -t node_array < <(scontrol show hostnames \"\$SLURM_JOB_NODELIST\")

prepare_c4_data.sh

@@ -0,0 +1,144 @@
+#!/bin/bash
+###############################################################################
+# Prepare C4 English dataset for Megatron training with DeepSeek V3
+#
+# This script:
+#   1. Downloads C4-en data from HuggingFace (configurable amount)
+#       GIT_LFS_SKIP_SMUDGE=1 git clone https://huggingface.co/datasets/allenai/c4
+#       cd c4
+#       git lfs pull --include "en/*"
+#   2. Converts to JSONL format
+#   3. Tokenizes into Megatron .bin/.idx format using DeepSeekV3Tokenizer
+#
+# Usage:
+#   bash prepare_c4_data.sh [--num_shards N] [--data_dir /path/to/data]
+#
+# By default downloads 1 shard (~350MB compressed, ~3M documents) for testing.
+# Full C4-en has 1024 shards. Adjust --num_shards for more data.
+###############################################################################
+
+set -e
+
+# ======================== Configuration ========================
+NUM_SHARDS=${NUM_SHARDS:-200}           # Number of C4 shards to download (1-1024)
+DATA_DIR=${DATA_DIR:-"/shared/c4"}
+PRIMUS_PATH=${PRIMUS_PATH:-"/shared/john/Primus"}
+TOKENIZER_TYPE="DeepSeekV3Tokenizer"
+TOKENIZER_MODEL="deepseek-ai/DeepSeek-V3"
+WORKERS=${WORKERS:-$(nproc)}          # Number of preprocessing workers
+HF_TOKEN=${HF_TOKEN:-"your_hf_token"}             # Set your HuggingFace token
+
+# Parse arguments
+while [[ $# -gt 0 ]]; do
+    case $1 in
+        --num_shards) NUM_SHARDS="$2"; shift 2;;
+        --data_dir)   DATA_DIR="$2";   shift 2;;
+        --workers)    WORKERS="$2";    shift 2;;
+        *) echo "Unknown option: $1"; exit 1;;
+    esac
+done
+
+# ======================== Paths ========================
+export RAW_DIR="${DATA_DIR}/en"       # Pre-downloaded shards live here
+export JSONL_DIR="${DATA_DIR}/jsonl"
+export TOKENIZED_DIR="${DATA_DIR}/tokenized"
+export TRAIN_OUTPUT_PREFIX="${TOKENIZED_DIR}/c4_en_train"
+export NUM_SHARDS
+
+mkdir -p "$RAW_DIR" "$JSONL_DIR" "$TOKENIZED_DIR"
+
+echo "============================================"
+echo "C4 English Data Preparation"
+echo "============================================"
+echo "NUM_SHARDS:     ${NUM_SHARDS} (out of 1024 total)"
+echo "DATA_DIR:       ${DATA_DIR}"
+echo "PRIMUS_PATH:    ${PRIMUS_PATH}"
+echo "TOKENIZER:      ${TOKENIZER_TYPE} / ${TOKENIZER_MODEL}"
+echo "WORKERS:        ${WORKERS}"
+echo "============================================"
+
+# ======================== Step 1: Merge shards into JSONL ========================
+echo ""
+echo ">>> Step 1: Merging C4 English shards into JSONL (${NUM_SHARDS} shards)..."
+echo "    (Download skipped — using pre-downloaded shards in ${RAW_DIR})"
+
+JSONL_FILE="${JSONL_DIR}/c4_en_train.jsonl"
+
+if [ -f "${JSONL_FILE}" ]; then
+    echo "JSONL file already exists: ${JSONL_FILE}"
+    echo "Skipping merge. Delete it to re-merge."
+else
+    # Verify shards exist
+    MISSING=0
+    for i in $(seq 0 $((NUM_SHARDS - 1))); do
+        SHARD_NAME=$(printf "c4-train.%05d-of-01024.json.gz" "$i")
+        if [ ! -f "${RAW_DIR}/${SHARD_NAME}" ]; then
+            echo "  WARNING: Missing shard ${SHARD_NAME}"
+            MISSING=$((MISSING + 1))
+        fi
+    done
+    if [ "$MISSING" -gt 0 ]; then
+        echo "ERROR: ${MISSING} shard(s) missing in ${RAW_DIR}. Cannot proceed."
+        exit 1
+    fi
+
+    echo "Decompressing and merging shards into JSONL ..."
+    for i in $(seq 0 $((NUM_SHARDS - 1))); do
+        SHARD_NAME=$(printf "c4-train.%05d-of-01024.json.gz" "$i")
+        SHARD_PATH="${RAW_DIR}/${SHARD_NAME}"
+        echo "  [${i}/${NUM_SHARDS}] Decompressing ${SHARD_NAME} ..."
+        zcat "${SHARD_PATH}" >> "${JSONL_FILE}"
+    done
+
+    DOC_COUNT=$(wc -l < "${JSONL_FILE}")
+    echo "Done! Total documents: ${DOC_COUNT}"
+    echo "Saved to: ${JSONL_FILE}"
+fi
+
+echo ">>> Step 1 complete."
+
+# ======================== Step 2: Tokenize ========================
+echo ""
+echo ">>> Step 2: Tokenizing with ${TOKENIZER_TYPE}..."
+
+JSONL_FILE="${JSONL_DIR}/c4_en_train.jsonl"
+
+if [ -f "${TRAIN_OUTPUT_PREFIX}_text_document.bin" ] && [ -f "${TRAIN_OUTPUT_PREFIX}_text_document.idx" ]; then
+    echo "Tokenized files already exist:"
+    echo "  ${TRAIN_OUTPUT_PREFIX}_text_document.bin"
+    echo "  ${TRAIN_OUTPUT_PREFIX}_text_document.idx"
+    echo "Skipping tokenization. Delete them to re-tokenize."
+else
+    # Need to set up Python path for Megatron imports
+    export PYTHONPATH="${PRIMUS_PATH}/third_party/Megatron-LM:${PRIMUS_PATH}:${PYTHONPATH:-}"
+
+    python3 "${PRIMUS_PATH}/examples/megatron/preprocess_data.py" \
+        --input "${JSONL_FILE}" \
+        --tokenizer-type "${TOKENIZER_TYPE}" \
+        --tokenizer-model "${TOKENIZER_MODEL}" \
+        --output-prefix "${TRAIN_OUTPUT_PREFIX}" \
+        --workers "${WORKERS}" \
+        --append-eod \
+        --partitions 1
+
+    echo ">>> Step 2 complete."
+fi
+
+# ======================== Summary ========================
+echo ""
+echo "============================================"
+echo "Data preparation complete!"
+echo "============================================"
+echo ""
+echo "Tokenized data files:"
+ls -lh "${TOKENIZED_DIR}/"
+echo ""
+echo "To use this data for training, set in run_dsv3.sh:"
+echo ""
+echo "  1. Change:  --mock_data True  →  --mock_data False"
+echo "  2. Add env:  export PRIMUS_TOKENIZED_DATA_PATH=${TRAIN_OUTPUT_PREFIX}_text_document"
+echo ""
+echo "Or pass directly via environment variable before running:"
+echo "  export PRIMUS_TOKENIZED_DATA_PATH=${TRAIN_OUTPUT_PREFIX}_text_document"
+echo ""
+echo "============================================"

primus/backends/megatron/patches/training_log/print_rank_last_patches.py

@@ -193,12 +193,26 @@ def inject(
                 local_rank = torch.cuda.current_device()
                 r_total, r_used, r_free = get_rocm_smi_mem_info(local_rank)
                 r_ratio = r_used / r_total
+
+                # get the max rocm_mem_usage
+                usage_tensor = torch.tensor([r_used], device="cuda", dtype=torch.float32)
+                world_size = torch.distributed.get_world_size()
+                gathered_usage = [torch.zeros_like(usage_tensor) for _ in range(world_size)]
+                torch.distributed.all_gather(gathered_usage, usage_tensor)
+
+                rocm_mem_usages = [t.item() for t in gathered_usage]
+                max_usage = max(rocm_mem_usages)
+                max_rank = rocm_mem_usages.index(max_usage)
+
                 rocm_mem_str = (
                     f" | rocm mem usage/free/total/usage_ratio: "
                     f"{r_used / 1024 ** 3:.2f}GB/"
                     f"{r_free / 1024 ** 3:.2f}GB/"
                     f"{r_total / 1024 ** 3:.2f}GB/"
                     f"{r_ratio * 100:.2f}%"
+                    f" | rank-{max_rank} rocm max mem usage/usage_ratio: "
+                    f"{max_usage / 1024 ** 3:.2f}GB/"
+                    f"{max_usage / r_total * 100:.2f}%"
                 )
                 # Cache for reuse on non-sampled iterations
                 self._last_rocm_mem_str = rocm_mem_str

start_training_dsv2_lite.sh

@@ -0,0 +1,51 @@
+#!/bin/bash
+
+export HF_TOKEN="your_hf_token"  # make it your own hf token
+export WANDB_API_KEY="your_wandb_api_key"  # make it your own wandb api key
+export DOCKER_IMAGE="docker.io/tasimage/primus:pr-563-ainic"
+#export SLURM_TREE_WIDTH=128 
+
+export NNODES=128
+export SLURM_TIME=07:00:00
+export SLURM_PARTITION=amd-aig
+
+# export NCCL_DEBUG=INFO
+export USING_AINIC=1
+export NCCL_IB_HCA="ionic_0:1,ionic_2:1,ionic_3:1,ionic_4:1,ionic_5:1,ionic_7:1,ionic_8:1,ionic_9:1"
+export GLOO_SOCKET_IFNAME=ens9np0 
+export NCCL_SOCKET_IFNAME=ens9np0 
+export CLEAN_DOCKER_CONTAINER=1
+
+export MBS=12
+export GBS=$((96 * NNODES))
+export PROFILE=False
+export TURBO_GROUPED_MLP=False
+export TURBO_DEEPEEP=True
+export LEGACY_GG=True
+export PRIMUS_DETERMINISTIC=0
+
+# export EXP=examples/megatron/configs/MI355X/llama3.1_8B-BF16-pretrain.yaml
+export EXP=examples/megatron/configs/MI355X/deepseek_v2_lite-BF16-pretrain.yaml
+export PRIMUS_TEAM=amd
+export PRIMUS_USER=tas
+export PRIMUS_EXP_NAME=dsv2_lite-pretrain-mbs_$MBS-gbs_$GBS-turbogg_$TURBO_GROUPED_MLP-turbodeepep_$TURBO_DEEPEEP-legacygg_$LEGACY_GG-profile_$PROFILE
+
+mkdir -p output/$PRIMUS_TEAM/$PRIMUS_USER/$PRIMUS_EXP_NAME
+bash ./examples/run_slurm_pretrain.sh \
+  --train_iters 10 \
+  --disable_wandb True \
+  --disable_tensorboard True \
+  --micro_batch_size $MBS \
+  --global_batch_size $GBS \
+  --seq_length 4096 \
+  --max_position_embeddings 4096 \
+  --use_turbo_grouped_mlp $TURBO_GROUPED_MLP \
+  --use_turbo_deepep $TURBO_DEEPEEP \
+  --moe_use_legacy_grouped_gemm $LEGACY_GG \
+  --cross_entropy_fusion_impl "te" \
+  --cross_entropy_loss_fusion True \
+  --profile $PROFILE \
+  --use_pytorch_profiler $PROFILE \
+  --profile_step_end 7 \
+  --profile_step_start 6 \
+  2>&1 | tee output/$PRIMUS_TEAM/$PRIMUS_USER/$PRIMUS_EXP_NAME/log.txt