wangyu
b4be3983b7
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>
70 lines
2.9 KiB
Bash
Executable File
70 lines
2.9 KiB
Bash
Executable File
#!/bin/bash
|
|
# Full CompanionGuard-RL pipeline on 4x RTX 5090.
|
|
#
|
|
# Step 1: Generate data (calls LLM API, single process)
|
|
# Step 2: Annotate + split (calls LLM API, single process)
|
|
# Step 3: Train detector (4 GPU DDP, BF16)
|
|
# Step 4: Train intervention (4 GPU BC + 1 GPU PPO)
|
|
# Step 5: Evaluate (single GPU)
|
|
#
|
|
# Usage:
|
|
# export DASHSCOPE_API_KEY=your_key # for Qwen
|
|
# bash scripts/run_full_pipeline.sh
|
|
|
|
set -e
|
|
|
|
NUM_GPUS=4
|
|
echo "======================================================"
|
|
echo " CompanionGuard-RL Full Pipeline — 4x RTX 5090"
|
|
echo "======================================================"
|
|
|
|
# ── Step 1: Data generation ────────────────────────────────────────────
|
|
echo ""
|
|
echo "[1/5] Generating dataset..."
|
|
python scripts/generate_data.py --config configs/data_generation.yaml
|
|
|
|
# ── Step 2: LLM annotation + split ─────────────────────────────────────
|
|
echo ""
|
|
echo "[2/5] Annotating and splitting dataset..."
|
|
python scripts/annotate_data.py \
|
|
--input data/raw/generated.jsonl \
|
|
--output data/processed/annotated.jsonl \
|
|
--config configs/data_generation.yaml
|
|
|
|
# ── Step 3: Train detector ──────────────────────────────────────────────
|
|
echo ""
|
|
echo "[3/5] Training risk detector (4 GPU DDP, BF16)..."
|
|
accelerate launch \
|
|
--num_processes=${NUM_GPUS} \
|
|
--mixed_precision=bf16 \
|
|
--multi_gpu \
|
|
scripts/train_detector.py \
|
|
--config configs/detector_config.yaml
|
|
|
|
# ── Step 4: Train intervention policy ──────────────────────────────────
|
|
echo ""
|
|
echo "[4/5] Training intervention policy (BC: 4 GPU, PPO: 1 GPU)..."
|
|
accelerate launch \
|
|
--num_processes=${NUM_GPUS} \
|
|
--mixed_precision=bf16 \
|
|
--multi_gpu \
|
|
scripts/train_intervention.py \
|
|
--config configs/intervention_config.yaml \
|
|
--train-data data/processed/train.jsonl
|
|
|
|
# ── Step 5: Evaluate ────────────────────────────────────────────────────
|
|
echo ""
|
|
echo "[5/5] Evaluating..."
|
|
python scripts/evaluate.py \
|
|
--detector-ckpt checkpoints/detector/best.pt \
|
|
--agent-ckpt checkpoints/intervention/final.pt \
|
|
--test-data data/processed/test.jsonl \
|
|
--config configs/detector_config.yaml \
|
|
--intervention-config configs/intervention_config.yaml \
|
|
--output experiments/eval_results.json
|
|
|
|
echo ""
|
|
echo "======================================================"
|
|
echo " Pipeline complete. Results: experiments/eval_results.json"
|
|
echo "======================================================"
|