lrm/lrm_sana/docs/migration_notes.md

Migration Notes: Flux -> SANA LRM

Reused Components

• Training loop orchestration from flux/trainer/scripts/train.py.
• Accelerator stack (base, debug, deepspeed).
• Optimizer and LR scheduler modules.
• Pairwise loss structure and evaluation flow.
• Dataset filtering logic including pseudo-preference support.

SANA-Specific Adaptations

• Replace Flux model wrapper with SANA latent-transformer wrapper.
• Replace Flux tokenizer/text path with variant-aware SANA path.
• Keep pairwise reward objective unchanged; adapt only feature extraction.
• Add model profile support for four requested checkpoints.

Naming Changes

• step_flux -> step_sana
• step_flux_base -> step_sana_base
• step_clip_flux -> step_clip_sana
• FluxPreferenceModel -> SanaPreferenceModel

Config Baseline Choice

• Dataset source follows flux working setting:
  • pickapic-anonymous/pickapic_v1
• Pseudo preference CSV path follows flux launcher fallback behavior.
• Logging/output path remains flux-style to simplify operational parity.

Runtime Notes

• Quick profile should force minimal run:
  • test split for train/valid/test
  • image size downscaled
  • max_steps=1
• Main profile should use production defaults with DeepSpeed sharded mode.

Known Risks During Migration

• Mismatch between tokenizer outputs and model forward signature.
• Caption embedding/channel shape mismatch in SANA transformer conditioning.
• Missing subfolders for some checkpoints (tokenizer_2 or text_encoder_2).
• Diffusers/transformers version compatibility for SANA classes.

Mitigation Strategy

• Use profile metadata for each checkpoint (explicit, not implicit only).
• Fail fast with clear validation errors at model init time.
• Keep quick-run smoke checks as gate before longer runs.
• Preserve fallback behavior from flux launcher for data and cache paths.