Planning scale up

[yair-schiff]

tl;dr: Train large (7-8B) E2D2 model.
Goal: Beat LLaDA on downstream tasks, be fully open-source, super fast inference!

Design decisions

• E2D2, block parameterization with block size $S=32$
• From scratch or fine-tune? TODO
• Dataset: TODO
• Hyperparams: TODO
  • Can also rely on this as resource: Scaling laws for dLLMs

Open questions

• Where can we run this? Options include:
  • Use Jax and run on Marin cc @eric-czech
  • Empire AI comes through (unlikely)
  • Lambda labs
• What is the total budget?

Major TODOs:

• Good profiling to convince ourselves current implementations are utilizing hardware appropriately (e.g., also make sure flex attention and compile are working as expected and providing boost) ((exp): Benchmark MFU as compute / nodes scale #7)
• Determine dataset mix
• Use remaining Lambda credits to do hyperparam sweep
• Download and process data to permanent file storage location
• Integrate an eval harness ((feat): Add evaluation harness #5)

Other TODOs

• Create our own backbone_modeling.py file (as opposed to rely on modeling_qwen3.py from transformers)
  • If using Jax, we should support a jax-compatible version of backbone and denoiser files. ((feat): Create backbone modeling file #66 )
• Clean up our custom generate method? Currently it has a lot of assumptions and perhaps too much custom logic (Clean up our custom generate method (rm assumptions and custom logic) #67)
• Need to sort out resumption: currently seeing loss spike ((todo): Validate that resuming from checkpoint works correctly #4).
• There is slight difference in loss / PPL logged to wandb and what we get when ((bug): Fix discrepancy between val loss logged during training and obtained from running likelihood eval with checkpoints #65).
• Confirm that push to hub works well and integrates with eval harness(es)
• Upgrade packages and pin software versions ((exp): Compare CUDA versions #58)
• General code cleanup: Perhaps we don't need to be so "general" / support so many baseline denoisers and backbones