feat: multi-GPU support for 4x RTX 5090 (PCIe DDP, BF16)

Hardware analysis:
  4x RTX 5090 32GB without NVLink is fully sufficient.
  PCIe 5.0 all-reduce overhead <1% of step time for MacBERT-large (340M params).
  BF16 mixed precision gives ~2x throughput vs FP32 on 5090.

Module B (Detector) — full 4-GPU DDP via Accelerate:
  - DistributedSampler with per-epoch shuffling (correct DDP data split)
  - BF16 autocast via accelerator.mixed_precision
  - Gradient accumulation handled by accelerator.accumulate()
  - Only rank-0 saves checkpoints and logs to wandb
  - accelerator.gather_for_metrics() for correct multi-GPU validation
  - per_gpu_batch_size=32, effective_batch = 32×4 = 128

Module C (Intervention) — hybrid parallel strategy:
  - Stage 1 (BC warm-up): all 4 GPUs via Accelerate DDP
    TensorDataset broadcast from rank-0 to all processes
  - Stage 2 (PPO): GPU-0 only — env-agent loop is inherently sequential
  - Detector preprocessing: distributed across all 4 GPUs via shard split
    + all_gather_object to collect results on rank-0

Configs updated:
  detector_config.yaml:    per_gpu_batch_size=32, gradient_accumulation_steps=1,
                           mixed_precision=bf16, num_workers=4
  intervention_config.yaml: BC per_gpu_batch_size=256, PPO batch_size=256

Launch scripts added:
  scripts/run_detector.sh         — single command: 4-GPU detector training
  scripts/run_intervention.sh     — single command: hybrid BC+PPO training
  scripts/run_full_pipeline.sh    — end-to-end pipeline steps 1-5

Co-Authored-By: Claude Sonnet 4.6 <noreply@anthropic.com>
This commit is contained in:
2026-05-09 17:56:13 +08:00
parent 4a0e71fb23
commit b4be3983b7
7 changed files with 637 additions and 184 deletions

View File

@@ -7,36 +7,39 @@ model:
data:
train_path: "data/processed/train.jsonl"
val_path: "data/processed/val.jsonl"
test_path: "data/processed/test.jsonl"
max_persona_len: 128
max_context_len: 512
max_response_len: 256
val_path: "data/processed/val.jsonl"
test_path: "data/processed/test.jsonl"
max_persona_len: 128
max_context_len: 512
max_response_len: 256
max_history_turns: 5
num_workers: 4 # DataLoader worker processes per GPU
training:
epochs: 10
batch_size: 16
per_gpu_batch_size: 32 # 4 GPUs × 32 = 128 effective batch per step
gradient_accumulation_steps: 1 # effective_batch = per_gpu × n_gpu × accum
lr: 2e-5
warmup_steps: 200
weight_decay: 0.01
gradient_clip: 1.0
eval_steps: 200
save_steps: 500
eval_steps: 200 # global steps between validation runs
mixed_precision: "bf16" # RTX 5090 has native BF16; use "fp16" for older GPUs
seed: 42
loss_weights:
binary: 1.0
level: 1.0
binary: 1.0
level: 1.0
primary: 1.0
fine: 1.0
fine: 1.0
evaluation:
binary_threshold: 0.5
fine_threshold: 0.4
fine_threshold: 0.4
logging:
project: "CompanionGuard-RL"
run_name: "detector-macbert"
project: "CompanionGuard-RL"
run_name: "detector-macbert-4gpu"
use_wandb: true
output: