lrm/lrm_sana/docs/checklist.md

SANA LRM Implementation Checklist

A. Documentation First

• Create docs/plan.md.
• Create docs/architecture.md.
• Create docs/checklist.md.
• Create docs/migration_notes.md.

B. Code Scaffold

• Ensure lrm_sana/setup.py exists.
• Ensure lrm_sana/train_lrm_sana.sh exists.
• Ensure full lrm_sana/trainer/* structure exists.

C. Config Registration

• Add SANA config registration in trainer/models/__init__.py.
• Add SANA config registration in trainer/datasets/__init__.py.
• Add SANA config registration in trainer/tasks/__init__.py.
• Add SANA config registration in trainer/criterions/__init__.py.
• Add SANA trainer config dataclass in trainer/configs/step_sana_configs.py.
• Add trainer/conf/step_sana_base.yaml.

D. Model Integration

• Implement trainer/models/sana_preference_model.py.
• Implement variant-aware checkpoint selection support.
• Confirm latent encode/noise/transformer/pooling path.
• Confirm text/image projection + logit_scale logic.

E. Dataset/Task/Criterion

• Implement trainer/datasets/step_sana_hf_dataset.py.
• Implement trainer/tasks/step_sana_task.py.
• Implement trainer/criterions/step_clip_criterion_sana.py.
• Preserve pseudo-preference filtering and timestep behavior.

F. Runtime Launcher

• Adapt train_lrm_sana.sh from flux launcher.
• Keep RUN_PROFILE=main|quick flow.
• Keep offline cache and DeepSpeed launch behavior.
• Add model profile switch for 4 SANA checkpoints.

G. Validation

• 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.