VGGRPO Phase C — GRPO 루프 배선 + 1-step smoke (block

TL;DR

Stage 2(forward) + Phase B(reward)를 실제 GRPO 루프에 배선해 end-to-end 1-step smoke 통과. 어댑터가 FF diffusers-shaped BaseAdapter를 우회하므로 커스텀 드라이버로 구동.
14B video 정책 backward가 1 B200을 초과 → micro-batch1·CFG-uncond-detach로도 OOM → block_wise full-recompute SAC로 fit (cosmos block_wise 버그를 monkeypatch로 우회).
SMOKE PASS: rewards 비상수, advantage std=0.999, log_prob.grad_fn=True, ratio_init=1.0, grad_norm=1.2e-4 → 정책경사가 LoRA로 올바르게 흐름.

PASS

1-step smoke

1.000000

ratio init

0.999

advantage std

143.9M

LoRA params

1통합 전략 — 커스텀 드라이버 (Path b)

우리 CosmosV2WAdapter는 FF BaseAdapter.__init__(Arguments+Accelerator, diffusers DiffusionPipeline 가정, accelerator.prepare, FSDP)을 의도적으로 우회했고, cosmos 파이프라인은 diffusers가 아니다 → FF GRPOTrainer에 그대로 못 꽂힌다.

결정: FF 트레이너와 싸우지 않고 GRPO를 직접 배선

우리 inference/forward + 우리 reward + hand-rolled group-advantage (r-mean)/(std+eps) + clipped surrogate ratio=exp(lp-old_lp). FF 소스 무수정. 어댑터에 get_trainable_parameters()(LoRA) 추가. cosmos를 FF의 diffusers-shape에 억지로 맞추는 것보다 안전.

2메모리 3단계 디버깅 핵심

14B video 정책 forward+backward(latent [1,16,24,54,96], 토큰 24×54×96)가 1 B200(178GB) 초과. 원인·해결을 단계적으로 규명:

시도	결과	원인
batch K=2 · CFG grad	OOM 178GB	배치×CFG×14B 활성화 과다
micro-batch 1 (순차 grad 누적)	OOM 178GB	단일 forward도 초과
CFG uncond detach (값 동일·grad 드롭)	OOM 178GB	cond 단독도 초과
block_wise SAC (full recompute)	PASS ✅	블록 입력만 저장 (~155GB)

cosmos block_wise는 버그. SACConfig.get_context_fn()이 None 반환 → ptd_checkpoint_wrapper가 context_fn() 무조건 호출 → TypeError: 'NoneType' object is not callable (rollout forward부터 깸).

해결: pipeline.py에서 get_context_fn을 모든 op PREFER_RECOMPUTE인 full-recompute 정책으로 monkeypatch(빌드 전), CosmosV2WPseudoPipelineConfig(sac_mode="block_wise")로 활성화. 기본 predict2_14b_720(mm-style)은 MLP addmm 출력 저장 → OOM.

3결과 — SMOKE PASS

tmp/smoke_grpo5.log (14B, sac_mode=block_wise, K=2, num_steps=8, W=4):

rewards: [0.9578, 0.9602] std=1.66e-3 (비-상수 ✅) advantage: [-0.707, +0.707] std=0.999 (group-normalize ✅) WIRING: log_prob.requires_grad=True grad_fn=True ratio_init=1.000000 (grad→LoRA, rollout≡train ✅) ti=0: ratio_mean=1.000000 loss=-1.79e-5 grad_norm=1.20e-4 (LoRA grad + optimizer step ✅) SMOKE PASS: reward_std>0=T adv_std>0=T wiring=T grad>0=T ratio_init≈1=T

정책경사가 LoRA로 올바르게 흐른다. ratio=1.0 정확은 rollout과 train forward가 동일 분포임을 재확인(무편향). loss·grad_norm이 작은 건 init에서 ratio≈1이라 정상 — 정책이 아직 안 발산. rollout trajectory: all_latents (10,16,24,54,96), log_probs (8,).

4다음 (Next)

전체 학습 드라이버: smoke는 1-iteration. epoch 루프 + 전체 train_timesteps + checkpoint 저장 + wandb + 실제 prompt/cond-image 데이터셋 → scripts/cosmos_ff/train_grpo.py + .sh(sbm/sbmr 제출).

속도/스케일(Phase D): block_wise는 full recompute라 정책 스텝이 느림(매 스텝 14B 재계산). train_timesteps 전체(≤35)×K×inner_epochs 순차 → 1 iteration이 길다. 완화: K/그룹, train 스텝 subsample, multi-GPU FSDP+context-parallel(cosmos CP 지원), reward async. uncond detach는 grad 근사(정확도 영향 모니터).