1배경 / 목적
B200 스텝 프로파일(전날 분석)에서 LingBot-VA 학습은 에피소드 1개 = micro-step 1개(batch 1)라 GPU당 토큰 예산 ≈68k의 10–15%만 쓰고, 유효 batch 128은 accum 16으로만 만들어 8 GPU에서 micro-step마다 FSDP all-gather/reduce-scatter가 반복(util 50–59%, 34 s/step). 공식 코드에는 B개 에피소드를 하나의 FlexAttention 시퀀스로 packing하는 구조(seq_ids 기반 block mask)가 있으나 길이가 같은 에피소드만 가능하다(default_collate가 F가 다른 텐서를 stack 못 함 → batch 2 시도 시 stack expects each tensor to be equal size).
사용자 지침 3가지: (1) bucket마다 최대 F가 다르면 F 값마다 재컴파일(4–6 s) → 패딩 F를 16/32 프레임 배수로 올림해 shape 수를 묶을 것. (2) 패딩 프레임도 linear/FFN GEMM은 그대로 타므로 bucketing으로 낭비 10% 미만을 맞출 것(향후 human video prefix 길이도 같은 bucket 기준). (3) batch 1×4 vs batch 4 loss 일치 검증에 길이가 다른 4개(패딩 있는 batch)도 추가해 loss 마스크를 확인할 것.
2작업 내용
2.1 packing 모듈 (lingbot-va)
wan_va/dataset/packing.py (신규)
episode_latent_frames(dataset) # 메타에서 F = (n_video−1)//4+1 (latent 로딩 없음). bucket key 함수 — human prefix 길이 추가 지점
LengthBucketBatchSampler(lengths, batch_size, rank, world_size, seed, pool_factor=50)
# 에폭마다 seed+epoch 셔플 → rank round-robin shard(꼬리 패딩, 모든 rank 동일 batch 수) → batch×50 pool을 길이 정렬 → batch 절단 → batch 순서 셔플
make_pad_collate(pad_multiple) # F를 round_up(max F, 8 latent = 32 비디오 프레임)로 zero-pad, valid_frames (B,) 반환
wan_va/train.py
batch_size>1 → batch_sampler=LengthBucketBatchSampler, collate_fn=pad_collate, epoch reset 시 set_epoch
_prepare_input_dict: latent_dict/action_dict['valid_frames'] 전달
compute_loss: frame_valid = arange(F) < valid_frames[:,None] → latent/action loss를 유효 프레임 평균으로(패딩 프레임 제외)
wan_va/modules/model.py FlexAttnFunc.init_mask(..., valid_frames=None)
latent_seq_id / action_seq_id (B,F,H,W)에서 frame ≥ valid_frames[b] → −1 → self/cross attention mask에서 제외(기존 padded_length 메커니즘과 동일)
configs/va_furniture_all_e128_pack_cfg.py furniture_all_e128_pack
8 GPU × batch 4 × accum 4 = 유효 128, lr 3e-5, warmup 200, pad_frames_multiple=8 (LINGBOT_PAD_MULTIPLE), bucket_pool_factor=50
2.2 검증 스크립트 · 프로파일 훅
scripts/lingbot/verify_packing.py(1 GPU, 학습 없음): sampler 통계 → 같은 길이 4개 batch → 다른 길이 4개 batch에서 packed forward/loss/backward(grad finite) 후 입력을 슬라이스한 단일 forward 4회와 예측·loss 비교. LINGBOT_PROFILE=n(train.py) 훅으로 micro-step별 토큰 수·fwd+bwd 시간·peak 메모리 출력.
주의(검증 중 발견): FlexAttention block mask가 클래스 속성이라 backward 시점에 다시 읽힘. packed forward와 backward 사이에 다른 shape의 forward를 끼우면 shape 불일치 에러 → 검증 순서를 backward 먼저로 조정. 실제 학습 루프는 해당 없음.
2.3 resume 시 lr 스케줄 이어가기 (부수 수정)
wan_va/train.py (resume_from = .../checkpoint_step_N)
step0 = N
lr_scheduler = LambdaLR(opt, lambda s: warmup_constant_lambda(s + step0, warmup_steps) * min(1, (s+1)/resume_warmup_steps))
# resume_warmup_steps 기본 50 (env LINGBOT_RESUME_WARMUP) — Adam moment가 초기화되므로 짧은 re-warmup만
단위 검증: 1500 재개 → s=0 6e-7, s=49 3e-5 도달 / warmup 도중(100/200) 재개 → 3e-7 (=3e-5×0.5×0.02)
3결과 (수치)
3.1 검증 (로그인 GPU 1장, all_low)
| 항목 | 결과 |
|---|
| sampler (rank 0/8) | 40 batch/epoch, padding 낭비 2.3%, distinct padded F 19 (F 103–362, 8 배수) |
| 같은 길이 4개 (F=120) | packed vs pooled loss rel diff latent 1e-4 / action 2e-5, grad 835/835 finite, peak 52 GB |
| 다른 길이 4개 (119/177/285/239 → 288) | rel diff latent 3e-5 / action 2.7e-4, grad 전부 finite(fully-masked 패딩 row NaN 없음), peak 73 GB(optimizer 제외) |
예측값 자체는 packed vs single이 bf16 누적 순서 차이 수준(rel err 1–5e-2)으로만 다름 → loss 일치가 의미 있는 기준.
3.2 1 GPU packed 프로파일 (smoke cfg, batch 4, LINGBOT_PROFILE=8)
| micro-step | 토큰 | fwd+bwd | 토큰당 |
|---|
| batch 1 (전날 실측) | 4.7k–10.7k | 0.85–1.9 s | 0.18 ms |
| batch 4, F=128 | 20.5k | 3.24 s | 0.158 ms |
| batch 4, F=232 | 35.5k | 5.57 s | 0.157 ms |
| batch 4, F=280 | 42.4k | 6.41 s | 0.151 ms |
| batch 4, F=288 (새 shape 컴파일 포함) | 43.5k | 9.38 s | 0.216 ms |
첫 micro-step 105 s(JIT). peak 메모리 98–125 GB(optimizer 포함, FSDP 없음). 에피소드당 1.3–1.4 s로 batch 1과 거의 같음 → 1 GPU에서는 packing 이득이 사실상 없음.
3.3 8 GPU 실측 (112276, 1500부터 packing)
| 구성 | micro-step/step | s/step | GPU util | 남은 1500 step |
|---|
| e128 unpacked (111882, batch 1 × accum 16) | 16 | 34 s | 50–59% | ≈14 h |
| e128 packed (112276, batch 4 × accum 4) | 4 | 19–29 s, 중앙값 23 s (step 1502–1513, 컴파일 일부 포함) | 53–56% (10분 window, JIT 포함) | ≈9.5 h |
첫 step 152 s(JIT). fb 92–99 GB/GPU, CPU 10%. util 수치는 비슷하지만 스텝 시간이 32% 줄었으므로 통신 대기가 계산으로 전환된 것.
핵심 발견: packing 이득 ≈1.5×는 예상대로 토큰당 계산(15%)이 아니라 FSDP micro-step 동기화 16→4회 감소에서 나옴. 8 GPU 이상에서만 의미 있는 최적화.
3.4 resume warmup 재시작 (버그)
| e128 구간 | job | 상태 |
|---|
| 0 → 761 | 111322 | warmup 0–200, full lr 200–761 |
| 750 → 1004 (2회 재개) | 111606 / 111700 | 750–950 warmup 재시작 |
| 1000 → 1732 | 111882 | 1000–1200 warmup 재시작 |
| 1500 → (packed) | 112276 | 1500–1700 warmup 재시작(구 코드로 시작, 다음 requeue부터 수정 적용) |
영향: 1732 step 중 600 step(35%)이 warmup → step-1000 평가(4가구 0/10)는 실질 full-lr ≈600 step(=77k 샘플)으로, e40 5000 step(200k 샘플)과의 비교가 아직 공정하지 않음. e40 baseline은 warmup 10이라 영향 미미.