File size: 3,446 Bytes
ef8f3ad | 1 2 3 4 5 6 7 8 9 10 11 12 13 14 15 16 17 18 19 20 21 22 23 24 25 26 27 28 29 30 31 32 33 34 35 36 37 38 39 40 41 42 43 44 45 46 47 48 49 50 51 52 53 54 55 56 57 58 59 60 61 62 63 64 65 66 67 68 69 70 71 72 73 74 75 76 77 78 79 80 81 82 83 84 | # SANA LRM Implementation Plan
## Goal
Build a working latent-space reward model (LRM) under `lrm_sana` that follows the proven training path used in `flux` and aligns with the LRM objective from https://arxiv.org/abs/2502.01051.
## Scope
- Keep code inside `lrm_sana`.
- Keep all markdown docs inside `lrm_sana/docs`.
- Reuse the available dataset path/protocol from `flux`.
- Preserve flux-like run facilities: launcher profiles, logs structure, checkpoint behavior.
- Support these SANA checkpoints via model profiles:
- `Efficient-Large-Model/Sana_600M_512px_diffusers`
- `Efficient-Large-Model/Sana_1600M_512px_diffusers`
- `Efficient-Large-Model/Sana_Sprint_0.6B_1024px_diffusers`
- `Efficient-Large-Model/Sana_Sprint_1.6B_1024px_diffusers`
## Constraints
- Use `pickapic-anonymous/pickapic_v1` dataset flow from `flux`.
- Keep pseudo-preference CSV filtering support (`vqa_aes_clip_score_mp.csv`).
- Follow distributed DeepSpeed launch conventions from `flux/train_flux.sh`.
- Add variant-aware SANA text path so single/dual encoder variants can be handled safely.
## Phases
### Phase 1: Docs-First Baseline
- Create planning and architecture docs in `lrm_sana/docs`.
- Define implementation checkpoints and verification criteria before code rewrites.
### Phase 2: Project Scaffolding
- Mirror `flux` trainer structure into `lrm_sana/trainer`.
- Add top-level files: `lrm_sana/setup.py`, `lrm_sana/train_lrm_sana.sh`.
- Keep import paths local to `lrm_sana` package.
### Phase 3: Config and Registration Wiring
- Add SANA config groups and names:
- task: `step_sana`
- dataset: `step_sana`
- model: `step_sana_base`
- criterion: `step_clip_sana`
- Add `step_sana_base.yaml` with flux-like accelerator/dataset defaults.
### Phase 4: SANA Preference Model
- Implement `trainer/models/sana_preference_model.py`:
- SANA tokenizer/text encoder loading.
- SANA VAE latent encode + scaling.
- Scheduler timestep/noise mixing.
- SANA transformer forward.
- Text/image projection heads and `logit_scale`.
- Keep save/load behavior compatible with current trainer checkpoint flow.
### Phase 5: Dataset/Task/Criterion Integration
- Implement `step_sana_hf_dataset.py` from `step_flux_hf_dataset.py` baseline.
- Implement `step_sana_task.py` from `step_flux_task.py` baseline.
- Implement `step_clip_criterion_sana.py` from flux criterion baseline.
- Keep pairwise preference loss, distributed gather, and eval accuracy logic.
### Phase 6: Runtime Script and Profiles
- Adapt `train_lrm_sana.sh` from `train_flux.sh`:
- `RUN_PROFILE=main|quick`
- offline cache env defaults
- pseudo-preference CSV fallback
- DeepSpeed sharded launch
- Add model profile selection variable for the four SANA checkpoints.
### Phase 7: Verification
- Static checks:
- module imports
- hydra config composition
- dataclass/config registration validity
- Runtime checks:
- quick smoke run (`max_steps=1`)
- output/log/checkpoint tree validation
## Deliverables
- Working SANA LRM code under `lrm_sana`.
- Documentation set in `lrm_sana/docs`.
- Flux-like launcher for H200 node testing.
## Success Criteria
- Training script launches and completes a quick run end-to-end.
- Dataloader, model forward, criterion loss, and eval metric path work without interface mismatch.
- Logs/checkpoints/config snapshots are written to expected locations.
- Model profile switch works across the four requested SANA checkpoints.
|