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>
47 lines
1.1 KiB
YAML
47 lines
1.1 KiB
YAML
model:
|
||
name: "hfl/chinese-macbert-large"
|
||
hidden_size: 1024
|
||
num_heads: 8
|
||
dropout: 0.1
|
||
use_lora: false
|
||
|
||
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
|
||
max_history_turns: 5
|
||
num_workers: 4 # DataLoader worker processes per GPU
|
||
|
||
training:
|
||
epochs: 10
|
||
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 # 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
|
||
primary: 1.0
|
||
fine: 1.0
|
||
|
||
evaluation:
|
||
binary_threshold: 0.5
|
||
fine_threshold: 0.4
|
||
|
||
logging:
|
||
project: "CompanionGuard-RL"
|
||
run_name: "detector-macbert-4gpu"
|
||
use_wandb: true
|
||
|
||
output:
|
||
checkpoint_dir: "checkpoints/detector"
|