File size: 1,816 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 | # SANA LRM Implementation Checklist
## A. Documentation First
- [x] Create `docs/plan.md`.
- [x] Create `docs/architecture.md`.
- [x] Create `docs/checklist.md`.
- [x] Create `docs/migration_notes.md`.
## B. Code Scaffold
- [x] Ensure `lrm_sana/setup.py` exists.
- [x] Ensure `lrm_sana/train_lrm_sana.sh` exists.
- [x] Ensure full `lrm_sana/trainer/*` structure exists.
## C. Config Registration
- [x] Add SANA config registration in `trainer/models/__init__.py`.
- [x] Add SANA config registration in `trainer/datasets/__init__.py`.
- [x] Add SANA config registration in `trainer/tasks/__init__.py`.
- [x] Add SANA config registration in `trainer/criterions/__init__.py`.
- [x] Add SANA trainer config dataclass in `trainer/configs/step_sana_configs.py`.
- [x] Add `trainer/conf/step_sana_base.yaml`.
## D. Model Integration
- [x] Implement `trainer/models/sana_preference_model.py`.
- [x] Implement variant-aware checkpoint selection support.
- [x] Confirm latent encode/noise/transformer/pooling path.
- [x] Confirm text/image projection + `logit_scale` logic.
## E. Dataset/Task/Criterion
- [x] Implement `trainer/datasets/step_sana_hf_dataset.py`.
- [x] Implement `trainer/tasks/step_sana_task.py`.
- [x] Implement `trainer/criterions/step_clip_criterion_sana.py`.
- [x] Preserve pseudo-preference filtering and timestep behavior.
## F. Runtime Launcher
- [x] Adapt `train_lrm_sana.sh` from flux launcher.
- [x] Keep `RUN_PROFILE=main|quick` flow.
- [x] Keep offline cache and DeepSpeed launch behavior.
- [x] Add model profile switch for 4 SANA checkpoints.
## G. Validation
- [x] Run static import checks.
- [ ] Run hydra config composition check.
- [ ] Run quick smoke training (`max_steps=1`).
- [ ] Verify logs/checkpoints/config snapshot output.
- [ ] Verify each SANA profile initializes cleanly.
|