Add multi-GPU DDP for Omni LoRA training

This commit is contained in:
Xin Wang
2026-08-23 22:41:39 +08:00
parent cb4abfdf82
commit 7cdb2d90c2
2 changed files with 402 additions and 92 deletions

View File

@@ -970,6 +970,56 @@ manifest contains only selected training anchors, so `train_omni_turn_lora.py`
uses the sparse dataset automatically and must be run without
`--online-current-tail`.
### Multi-GPU Thinker-LoRA training
Fully cached AV training supports native single-node DDP through `torchrun`.
Each rank holds one complete 3B Thinker replica, reads a disjoint shard of the
selected cache anchors, and synchronizes only the trainable LoRA and turn-head
gradients. Model checkpoint loading is deliberately staggered by rank so four
processes do not read all checkpoint shards simultaneously. During gradient
accumulation, DDP synchronization occurs only on the final microbatch. Only
rank zero writes the adapter, head, summary, and progress log.
For four GPUs, use per-GPU batch 2 and accumulation 2 for global batch 16:
```bash
CUDA_VISIBLE_DEVICES=0,1,2,3 \
OMP_NUM_THREADS=2 \
MKL_NUM_THREADS=2 \
OPENBLAS_NUM_THREADS=2 \
PYTORCH_ALLOC_CONF=expandable_segments:True \
torchrun --standalone --nproc_per_node=4 train_omni_turn_lora.py \
--model-id /data1/yfzx/wangx/Models/Qwen2.5-Omni-3B \
--cache-dir /data1/yfzx/wangx/Datasets/easycom/omni_lora_cache_sparse_8s/Session_7 \
--cache-dir /data1/yfzx/wangx/Datasets/easycom/omni_lora_cache_sparse_8s/Session_11 \
--cache-dir /data1/yfzx/wangx/Datasets/easycom/omni_lora_cache_sparse_8s/Session_10 \
--cache-dir /data1/yfzx/wangx/Datasets/easycom/omni_lora_cache_sparse_8s/Session_6 \
--cache-dir /data1/yfzx/wangx/Datasets/easycom/omni_lora_cache_sparse_8s/Session_9 \
--cache-dir /data1/yfzx/wangx/Datasets/easycom/omni_lora_cache_sparse_8s/Session_4 \
--cache-dir /data1/yfzx/wangx/Datasets/easycom/omni_lora_cache_sparse_8s/Session_8 \
--cache-dir /data1/yfzx/wangx/Datasets/easycom/omni_lora_cache_sparse_8s/Session_5 \
--cache-dir /data1/yfzx/wangx/Datasets/easycom/omni_lora_cache_sparse_8s/Session_12 \
--output-dir checkpoints/omni_turn_lora_easycom_sparse_8s_ddp4 \
--max-steps 10170 \
--batch-size 2 \
--gradient-accumulation-steps 2 \
--num-workers 0 \
--block-cache-size 32 \
--learning-rate 2e-5 \
--head-learning-rate 1e-4 \
--warmup-steps 500 \
--dtype bfloat16 \
--flash-attention-2 \
--log-every 20
```
`10170` is only the current three-epoch estimate for roughly 54,000 selected
anchors. Recompute it from the completed manifests as
`ceil(num_examples / 16) * 3`. Do not pass `--device`: under `torchrun`, each
process uses its `LOCAL_RANK`. DDP intentionally rejects
`--online-current-tail`; the online producer-consumer path remains a separate
benchmark.
For a speed-only alternative, `train_omni_turn_lora.py` supports a two-GPU
producer-consumer benchmark. The producer GPU runs the frozen AV towers and
places current-tail embeddings in a bounded CPU queue while the consumer GPU