feat(deepseek-v4): let inference.sh launch across nodes

examples/models/deepseek_v4/inference.sh always ran
torch.distributed.run --nproc_per_node=$((PP * EP)), i.e. all ranks on one
node. With the DeepSeek-V4-Pro defaults (PP=4, EP=8) that is 32 processes on
a single node, which a 4-GPU GB200 NVL72 node cannot run although the header
says the defaults are for GB200.

Add NNODES / NODE_RANK / MASTER_ADDR / MASTER_PORT (and GPUS_PER_NODE) knobs.
When NNODES > 1 the launcher passes --nnodes/--node_rank/--master_addr/
--master_port to torch.distributed.run; PP*EP must be divisible by NNODES.
With the variables unset the behaviour is unchanged. Document the Slurm form
(one task per node) in the README.

Verified on 8 x GB200 (TP1/PP4/EP8, real DeepSeek-V4-Pro weights from an
imported Megatron checkpoint): coherent generation, 12 min 49 s including
the 3 TB load.

Signed-off-by: linmuchuiyang <6169320+linmuchuiyang@users.noreply.github.com>

Author: linmuchuiyang

examples/models/deepseek_v4/README.md

@@ -222,6 +222,11 @@ required storage and validate checkpoint resume explicitly.
   [`inference.sh`](inference.sh) example uses legacy full-prefix generation and
   is not a verified KV-cache path. HF-native inference is the target after a
   verified Megatron-to-HF export.
+- [`inference.sh`](inference.sh) needs `PP * EP` GPUs (32 for the Pro
+  defaults `PP=4`, `EP=8`). On 4-GPU nodes such as GB200 NVL72 spread them over
+  several nodes with one task per node:
+  `srun --ntasks-per-node=1 bash -c 'NNODES=$SLURM_JOB_NUM_NODES NODE_RANK=$SLURM_NODEID MASTER_ADDR=$(scontrol show hostnames "$SLURM_JOB_NODELIST" | head -n1) ./inference.sh'`.
+  With the variables unset the script launches all ranks on the current node.
 - Fused mHC requires Blackwell (`sm_100`). Hopper uses the unfused path.
 - CPU import and export require enough RAM for the full BF16 model plus
   conversion workspace and remain unverified in the current card.

examples/models/deepseek_v4/inference.sh

@@ -22,6 +22,11 @@
 #   EP: expert-parallel size (default: 4 for Flash, 8 for Pro)
 #   PP: pipeline-parallel size (default: 1 for Flash, 4 for Pro)
 #   PROMPT: prompt string (default: "Explain hyper-connections in transformer models.")
+#   NNODES / NODE_RANK / MASTER_ADDR / MASTER_PORT: multi-node launch (defaults 1 / 0 / 127.0.0.1 / 29500).
+#       The PP*EP ranks are spread over NNODES nodes; GPUS_PER_NODE defaults to PP*EP/NNODES.
+#       Under Slurm run one task per node, for example:
+#         srun --ntasks-per-node=1 bash -c 'NNODES=$SLURM_JOB_NUM_NODES NODE_RANK=$SLURM_NODEID \
+#           MASTER_ADDR=$(scontrol show hostnames "$SLURM_JOB_NODELIST" | head -n1) ./inference.sh'
 #   NUM_OF_HYBRID_EP_RANKS_PER_NVLINK_DOMAIN: optional HybridEP topology override
 #   NVLINK_DOMAIN_SIZE: optional physical NVLink-domain size override
 #   USE_MNNVL: optional HybridEP fabric override (0 or 1)
@@ -51,9 +56,26 @@ if [[ -z "${PP:-}" ]]; then
 fi
 TP=1
 
+# Multi-node launch: PP*EP ranks over NNODES nodes (default: everything on this node).
+NNODES=${NNODES:-1}
+NODE_RANK=${NODE_RANK:-0}
+MASTER_ADDR=${MASTER_ADDR:-127.0.0.1}
+MASTER_PORT=${MASTER_PORT:-29500}
+WORLD=$((PP * EP))
+if (( WORLD % NNODES != 0 )); then
+    echo "ERROR: PP*EP=${WORLD} is not divisible by NNODES=${NNODES}" >&2
+    exit 1
+fi
+GPUS_PER_NODE=${GPUS_PER_NODE:-$((WORLD / NNODES))}
+TORCHRUN_ARGS=(--nproc_per_node="${GPUS_PER_NODE}")
+if (( NNODES > 1 )); then
+    TORCHRUN_ARGS+=(--nnodes="${NNODES}" --node_rank="${NODE_RANK}"
+                    --master_addr="${MASTER_ADDR}" --master_port="${MASTER_PORT}")
+fi
+
 # DSv4 hybrid attention does not support inference contexts yet, so use full-prefix generation.
 # Inference directly from the HF checkpoint (Bridge dequantises in-flight).
-uv run python -m torch.distributed.run --nproc_per_node=$((PP * EP)) \
+uv run python -m torch.distributed.run "${TORCHRUN_ARGS[@]}" \
     examples/conversion/hf_to_megatron_generate_text.py \
     --hf_model_path "${HF_MODEL_ID}" \
     --prompt "${PROMPT}" \
@@ -65,7 +87,7 @@ uv run python -m torch.distributed.run --nproc_per_node=$((PP * EP)) \
 # Inference from a previously-imported Megatron checkpoint (faster cold start).
 MEGATRON_DIR="${WORKSPACE}/models/${MODEL_VARIANT}"
 if [[ -d "${MEGATRON_DIR}/iter_0000000" ]]; then
-    uv run python -m torch.distributed.run --nproc_per_node=$((PP * EP)) \
+    uv run python -m torch.distributed.run "${TORCHRUN_ARGS[@]}" \
         examples/conversion/hf_to_megatron_generate_text.py \
         --hf_model_path "${HF_MODEL_ID}" \
         --megatron_model_path "${MEGATRON_DIR}" \