[CI] [Training] Initial e2e small training test (#504)

2025-06-11 13:53:36 -07:00
parent 0f2bbe71ac
commit c0026ca5ba
8 changed files with 274 additions and 100 deletions
@@ -22,7 +22,7 @@ torchrun --nnodes 1 --nproc_per_node $NUM_GPUS\
    --sp_size 4 \
    --tp_size 4 \
    --hsdp_replicate_dim 1 \
-    --hsdp_shards 4 \
+    --hsdp_shard_dim 4 \
    --num_gpus $NUM_GPUS \
    --train_sp_batch_size 1\
    --dataloader_num_workers 10\