Skip to content

Target repeated MOSS transformer blocks for regional compilation - #3289

Merged
bghira merged 1 commit into
mainfrom
fix/moss-regional-compilation
Oct 20, 2026
Merged

bghira merged 1 commit into
mainfrom
fix/moss-regional-compilation

Conversation

@bghira

@bghira bghira commented Oct 20, 2026

Copy link
Copy Markdown
Owner

Accelerate's generic regional compiler also wraps MOSS auxiliary modules with different fixed parameter shapes, including the shared block and output-head modulation layers. These specializations can exhaust Dynamo's shared wrapper cache before image-size variation is considered.

Add an opt-in model-foundation context that uses Diffusers' existing compile_repeated_blocks() and _repeated_blocks declaration. MOSS enables this policy, compiling transformer blocks in place while suppressing Accelerate's automatic auxiliary-region compilation during preparation. Compiler settings are restored even if preparation fails. Training and inference preparation use the same hook; other model families retain their current policy.

Validation:

  • .venv/bin/python -m unittest -v tests.test_moss_image tests.test_trainer: 153 tests pass, including output/gradient parity, parameter identity, compilation scope, plugin restoration, and training/inference wiring.
  • Two H100s, tiny MOSS fixture (width 32, ten blocks): actual trainer preparation and three BF16 LoRA optimizer steps pass with both DDP and FSDP2, using ten transformer blocks, gradient checkpointing, dynamic image/caption sizes, and Inductor with fullgraph enabled. Auxiliary modulation layers remain eager.
  • The two-H100 DDP check also passes with recompile_limit=8 and fail_on_recompile_limit_hit=True, so hitting the limit cannot silently fall back to eager execution.
  • Full suite: .venv/bin/python -m unittest -v -f completed 7,052 tests successfully (292 skipped).
  • All pre-commit hooks and git diff --check pass.

Closes #3286.

@bghira
bghira merged commit 97d8fe2 into main Oct 20, 2026
2 checks passed
@bghira
bghira deleted the fix/moss-regional-compilation branch October 20, 2026 15:51
Sign up for free to join this conversation on GitHub. Already have an account? Sign in to comment

Labels

None yet

Projects

None yet

Development

Successfully merging this pull request may close these issues.

MOSS regional compilation exhausts Dynamo’s recompilation limit

1 participant