lrm/lrm_sana/docs/plan.md

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