Skip to content

Relax triton pin from ==3.5.0 to >=3.4.0 - #52

Closed
cennn wants to merge 1 commit into
mainfrom
fix/relax-triton-pin
Closed

Relax triton pin from ==3.5.0 to >=3.4.0#52
cennn wants to merge 1 commit into
mainfrom
fix/relax-triton-pin

Conversation

@cennn

@cennn cennn commented Aug 4, 2026

Copy link
Copy Markdown
Collaborator

Summary

  • Change triton==3.5.0 to triton>=3.4.0 in requirements.txt
  • The hard pin shadows the NGC-bundled pytorch-triton (3.4.0) when installing in PyTorch NGC containers, causing different Triton backend_hash values and numerically different compiled kernels — breaking compile-mode bit-exact reproducibility across images
  • Since pyproject.toml declares no triton runtime dependency, this file is for development use only; relaxing the constraint lets the environment's existing triton be reused

Context

When MAGI-2's Dockerfile runs pip install -r requirements.txt before installing MagiCompiler, triton==3.5.0 overwrites the NGC base image's pytorch-triton 3.4.0. This produces a different backend_hash in Inductor-generated Triton kernels, leading to max ~5e-4 numerical differences in VAE encode that cascade to max ~1.78 divergence after DiT processing (chaotic amplification). The issue is absent in eager mode and only manifests in torch.compile mode.

Test plan

  • Verify MagiCompiler imports and functions correctly with triton 3.4.0
  • Rebuild MAGI-2 image without triton override, confirm compile-mode parity with Athena

The hard pin on triton==3.5.0 causes issues in NGC-based Docker images
where pytorch-triton 3.4.0 is bundled with PyTorch. Installing
triton==3.5.0 shadows the bundled version, leading to different Triton
compiler backends (backend_hash) which produces numerically different
compiled kernels — breaking bit-exact reproducibility between images
that use different triton versions.

Since MagiCompiler's pyproject.toml declares no triton runtime
dependency, this requirements.txt is for development only. Relaxing
to >=3.4.0 allows the environment's existing triton to be reused.
@cennn

cennn commented Aug 4, 2026

Copy link
Copy Markdown
Collaborator Author

Closing: the fix belongs in downstream Dockerfiles (Athena & MAGI-2) that strip the triton pin before installing, rather than relaxing it in MagiCompiler itself.

@cennn cennn closed this Aug 4, 2026
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.

1 participant