[JAX] Fix transformers CVEs in the Primus MaxDiffusion image

[olehtika]

Motivation

rocm/jax-training:maxtext-v26.*, the base of docker/primus_maxdiffusion.ubuntu.amd.Dockerfile, pins transformers 4.57.3, which carries CVE-2026-4372, CVE-2026-5241 and CVE-2026-9856. transformers 5.10.0 is the first release that fixes all three.

Upgrading is not a one-liner: v5 removed every Flax implementation, including the FlaxCLIPTextModel and FlaxT5EncoderModel that FLUX's text encoders import. On a plain upgrade FLUX dies at import, and WAN dies on a bare transformers.models.clip.modeling_flax_clip import that v5 deleted outright.

Technical Details

Rather than look for Flax replacements that no longer exist, this runs FLUX's PyTorch CLIP-L and T5-XXL encoders under JAX through Torchax, which maxdiffusion already uses for LTX2's Gemma3 and WAN's UMT5 encoders. FLUX's text encoders are frozen and only run during embedding precompute — the trainer deletes them before the training loop — so tracing them costs nothing on the training path.

The image gains:

1. transformers>=5.10.0, installed after requirements-maxdiffusion.txt. That file pins transformers==4.57.3, so installing any earlier means pip puts the vulnerable version straight back; the first build of this change shipped an image that looked patched and had 4.57.3 in it. A version assertion at the end of the build now fails loudly rather than let that recur.
2. torchax==0.0.13. Pinned because it dispatches against private torch overloads: 0.0.13 imports cleanly on the base's torch 2.12.0 but fails at import on torch >= 2.14, which dropped the aten.prod Dimname overload it references. Installed --no-deps as a guard rather than a fix — 0.0.13 declares no unconditional dependencies, so nothing is being suppressed today, but a later release resolving its own torch or jax over the base's ROCm builds would be a bad day.
3. Three patches:
   • maxdiffusion-flux-transformers5.patch — the Torchax text-encoder wrappers, FLUX configs repointed at the text_encoder/text_encoder_2 subfolders of the official black-forest-labs/FLUX.1-dev repo instead of third-party Flax mirrors, and a CPU-only parity test.
   • maxdiffusion-pin-transformers5-compat.patch — the remaining v5 and hub-1.x fallout: guarded Flax imports on the WAN and pipeline_flax_utils paths, and use_auth_token renamed to token.
   • primus-flux-torch-text-encoders.patch — Primus's configs/models/maxdiffusion/flux_dev.yaml names ariG23498/clip-vit-large-patch14-text-flax and ariG23498/t5-v1-1-xxl-flax, which hold no PyTorch weights and so cannot be read once the encoders run under Torchax. Repointed at the official FLUX.1-dev subfolders: same model, first-party copy, and the same weights maxdiffusion's own base_flux_dev.yml names.

The two maxdiffusion patches are rebased onto 68e069659f0af80694559e29939a17f879fe7f6a, the commit Primus pins at third_party/maxdiffusion (verified at dc3f4ba, the submodule commit develop pins) and the one the base's tree is checked out at. An ARG MAXDIFFUSION_COMMIT plus git rev-parse guard fails the build if that pin moves, rather than applying the patches to a tree they were never rebased onto.

The patched code works on both transformers 4.57.3 and 5.x — dtype= replaces the torch_dtype= kwarg that 4.56 deprecated and v5 removed, and attn_implementation="eager" is accepted by both — so this does not strand anyone still on the old pin. The new *_subfolder config keys are read through a helper that tolerates their absence, since pyconfig raises ValueError rather than AttributeError for a key it does not hold, and a config written before these keys existed should fall back to the repository root rather than crash.

The maxdiffusion and Primus changes are submitted to those projects separately, as AI-Hypercomputer/maxdiffusion#483 and AMD-AGI/Primus#1159. Each patch and its guard come out of this image once the corresponding pin carries the fix, leaving only the transformers upgrade.

Relationship to AI-Hypercomputer/maxdiffusion#482

Worth stating, because it becomes a prerequisite the moment this image's pin moves. AI-Hypercomputer/maxdiffusion#482 fixes three independent FLUX regressions on maxdiffusion main, introduced by 64450df4 (2026-07-16, the Flux.2-klein onboarding): the timestep/guidance embeddings get projected twice and fail at init on a broadcast error, the legacy BFL checkpoint loader still maps to a norm_out submodule that was renamed, and norm_out silently inherits the opposite shift/scale convention, which inflates the loss about 8x without failing.

None of that affects this PR. The patches here are rebased onto 68e0696 (2026-03-03), which is an ancestor of the Klein commit, so the regressions do not exist in the tree they apply to — and #482 touches embeddings_flax.py, flux/util.py and transformer_flux_flax.py, none of which these patches touch, so there is no overlap to reconcile either. What it does mean is that a future ARG MAXDIFFUSION_COMMIT bump has to land on a tree carrying both #482 and the Torchax text encoders; the Torchax change alone would build an image whose FLUX training is broken.

Relationship to #247

#247 moves this image's base from maxtext-v26.6 to maxtext-v26.7. The tags maxtext-v26.7 and maxtext-v26.7-jax0.11.0-te2.17.0 are the same manifest, and that is the base everything below was tested on, so this change and #247 agree on the target base. This PR deliberately leaves ARG BASE_DOCKER alone to avoid conflicting with #247; happy to rebase on top of it, or fold this into that PR, whichever the maintainers prefer.

Test Plan

All three MI355X MaxDiffusion configs, run through this repo's Primus path on 8x MI355X (gfx950), 20 steps each at bf16:

madengine run --tags maxdiffusion_MI355X_flux_dev-pretrain
madengine run --tags maxdiffusion_MI355X_wan2.1_1.3b-pretrain
madengine run --tags maxdiffusion_MI355X_wan2.1_14b-pretrain

built from this Dockerfile with BASE_DOCKER=rocm/jax-training:maxtext-v26.7, against Primus at the pinned dc3f4ba.

Test Result

All three succeed, rc=0, with step times and losses matching a pre-upgrade baseline run of the same configs:

Config	Status	s/step (steady)	TFLOP/s/GPU	Loss @ step 18
maxdiffusion_MI355X_flux_dev-pretrain	SUCCESS	1.548	583.3	572
maxdiffusion_MI355X_wan2.1_1.3b-pretrain	SUCCESS	4.437	920.9	1.540
maxdiffusion_MI355X_wan2.1_14b-pretrain	SUCCESS	26.868	789.3	1.551

The built image carries transformers 5.17.0, huggingface_hub 1.31.0 and torchax 0.0.13; all three patches apply cleanly; the commit guard passes against the base's tree; and the version assertion reports transformers 5.17.0 OK. FLUX loads its text encoders from the FLUX.1-dev text_encoder/text_encoder_2 subfolders, confirming the Torchax path is the one exercised rather than a silent fallback.

One note for reviewers: testing was on the v26.7 base, matching #247. I have not verified the v26.6 base that develop currently names — the commit guard should still pass there, since Primus's pin is unchanged, and v26.6's torch is only a problem for the torchax pin if it falls outside the range 0.0.13 is known to import on, which is 2.8.0 through 2.12.0 (it breaks at 2.14). v26.7 carries 2.12.0.

Submission Checklist

• Look over the contributing guidelines at https://github.com/ROCm/TheRock/blob/main/GOVERNANCE.md#pull-requests.