|
| 1 | +# DiffusionGemma |
| 2 | + |
| 3 | +Megatron Bridge supports the public `google/diffusiongemma-26B-A4B-it` |
| 4 | +checkpoint for block-diffusion SFT. |
| 5 | + |
| 6 | +## What is implemented |
| 7 | + |
| 8 | +- HF ↔ Megatron conversion and exact checkpoint tensor coverage. |
| 9 | +- Text and image-conditioned encoder/decoder inference. |
| 10 | +- Uniform-state corruption with one diffusion time per example. |
| 11 | +- `base-sft`, `reweighted-ce`, `loo-ce`, and `reweighted-loo-ce` objectives. |
| 12 | +- Detached two-pass self-conditioning. |
| 13 | +- Optional clean-target encoder next-token loss. It runs a separate encoder |
| 14 | + pass so the denoising decoder never sees the target tokens as context. |
| 15 | +- Tensor parallelism with sequence parallelism and expert parallelism. |
| 16 | +- Distributed optimizer training and torch-dist checkpoint save/resume. |
| 17 | + |
| 18 | +The reference SFT objective does not add a router balancing penalty. The |
| 19 | +provider exposes `moe_aux_loss_coeff` for experiments, but its default is zero. |
| 20 | + |
| 21 | +## Recipe |
| 22 | + |
| 23 | +Use `diffusion_gemma_26b_sft_config()` from |
| 24 | +`megatron.bridge.recipes.diffusion_gemma`. The production recipe uses TP=1, |
| 25 | +EP=8, one pipeline stage, BF16 mixed precision, a distributed optimizer, and |
| 26 | +full uniform activation recomputation. The recomputation setting keeps |
| 27 | +large-image prompts with long text and prior-target context within H200 |
| 28 | +memory; callers can override it for shorter inputs. |
| 29 | +Pass `train_path` and optional `validation_path`/`test_path` to train on |
| 30 | +existing JSONL split files; the adapter never creates a new split. Without |
| 31 | +`train_path`, the recipe uses synthetic smoke-test data. Each JSONL row has the |
| 32 | +form: |
| 33 | + |
| 34 | +```json |
| 35 | +{"messages":[{"role":"user","content":[{"type":"image","image":"images/example.png"},{"type":"text","text":"Describe the image."}]}],"completion":"A blue square above a red circle."} |
| 36 | +``` |
| 37 | + |
| 38 | +Image paths are local and relative to the JSONL file. The caller is responsible |
| 39 | +for supplying only instructions/context available at inference time and for |
| 40 | +curating completion labels. Multi-block completions are split without |
| 41 | +truncation; only prior clean target blocks are appended to the encoder context. |
| 42 | + |
| 43 | +```python |
| 44 | +from megatron.bridge.recipes.diffusion_gemma import diffusion_gemma_26b_sft_config |
| 45 | +from megatron.bridge.training.diffusion_gemma_step import forward_step |
| 46 | +from megatron.bridge.training.pretrain import pretrain |
| 47 | + |
| 48 | +cfg = diffusion_gemma_26b_sft_config( |
| 49 | + train_path="data/train.jsonl", validation_path="data/validation.jsonl", |
| 50 | + test_path="data/test.jsonl", |
| 51 | +) |
| 52 | +pretrain(cfg, forward_step) |
| 53 | +``` |
| 54 | + |
| 55 | +Run this entry point with `uv run python -m torch.distributed.run |
| 56 | +--nproc_per_node=8 train.py`. Initial weights are loaded through the HF bridge |
| 57 | +hook; `checkpoint.load` points to the resume directory, not the HF model ID. |
| 58 | + |
| 59 | +The dataset must emit the following fields: |
| 60 | + |
| 61 | +```text |
| 62 | +input_ids, attention_mask, position_ids, mm_token_type_ids, |
| 63 | +canvas_ids, canvas_mask, |
| 64 | +pixel_values, image_position_ids # optional for image examples |
| 65 | +``` |
| 66 | + |
| 67 | +Prompts are left-padded; canvases are right-padded. When sequence parallelism |
| 68 | +is enabled, the prompt length must be divisible by TP. |
| 69 | + |
| 70 | +## Validation evidence |
| 71 | + |
| 72 | +The tiny model has HF loss/gradient parity at roughly 1e-3 relative gradient |
| 73 | +error and overfits a fixed batch. TP=2 and EP=2 checks pass, including a |
| 74 | +prompt crossing the toy sliding-window boundary. A real 26B BF16 |
| 75 | +run on eight H200 GPUs completed three unfrozen optimizer steps with EP=8 and a |
| 76 | +distributed optimizer; no updates were skipped or non-finite. After warmup, |
| 77 | +the measured steps were about 4.6–6.1 seconds and peak rank-0 allocated memory |
| 78 | +was about 74.6 GiB. |
| 79 | +With the clean-target encoder objective enabled, the three-step smoke also |
| 80 | +passed: warmed-up steps measured 6.1–8.4 seconds and peak rank-0 allocated |
| 81 | +memory was about 83.6 GiB. These are small synthetic microbatches, not a |
| 82 | +prediction of production throughput. No 26B optimizer checkpoint |
| 83 | +was written during either smoke run. |
| 84 | +On a synthetic long-context multimodal example with a large image, a long text |
| 85 | +prompt, and a multi-block completion, the recipe completed three EP=8 steps |
| 86 | +with recomputation enabled. Peak rank-0 allocation was 123.8 GiB and warmed-up |
| 87 | +steps took about 3.4 seconds. The same example without recomputation exceeded |
| 88 | +the 139.8 GiB H200 capacity during attention. |
| 89 | + |
| 90 | +The two-GPU `pretrain()` checkpoint test saved at step three and resumed in |
| 91 | +a fresh process. Steps four through six reproduced the uninterrupted losses, |
| 92 | +gradient norms, and consumed-sample positions exactly. |
| 93 | + |
| 94 | +## Known limitations |
| 95 | + |
| 96 | +- Context parallelism and pipeline parallelism are not implemented by the |
| 97 | + diffusion training step. |
| 98 | +- Context lengths beyond the validated sliding-window boundary still need |
| 99 | + production-shaped memory and quality evaluation. |
| 100 | +- `MockDiffusionGemmaDatasetConfig` is for smoke tests only. Pass |
| 101 | + `train_path` to use `DiffusionGemmaDatasetConfig` with curated JSONL input files. |
| 102 | +- A successful optimizer smoke test does not demonstrate task quality or |
| 103 | + generalization. Those require held-out task evaluation on a saved split. |
0 commit comments