cp: fix(mistral3): preserve medium VLM checkpoint layout (2758) into r0.5.0 - #2762
Merged
Conversation
Signed-off-by: Alexandros Koumparoulis <akoumparouli@nvidia.com> Signed-off-by: NeMo Bot <nemo-bot@nvidia.com>
Contributor
Author
|
/ok to test feecf1a |
akoumpa
approved these changes
Jun 25, 2026
akoumpa
added a commit
that referenced
this pull request
Aug 30, 2026
* build: bump FA3/FA4 pin to pick up CuTe compile-key fixes The pinned flash-attention ref 002cce0a1 (2026-07-03) predates both CuTe compile-key fixes, so every forward/backward with a tensor max_seqlen -- i.e. every varlen/packed-sequence step -- rebuilt the compile key and recompiled the kernel. Reported against a Qwen3.5-VL packed-sequence finetune on B300. Bump to ce088ab9ce0f (main), which adds over the old pin: - #2507 [CuTe, Bwd] fix backward compile key churn (max_seqlen tensor) - #2762 [CuTe, Fwd] stabilize tensor max_seqlen compile key - #2745 [CuTe] fix forward dynamic-shape correctness - #2819 [CuTe] speed up scalar SM100 mask compilation #2798 raised the cute requirement to nvidia-cutlass-dsl>=4.6.2, so pin the FA4 CUTLASS DSL install to 4.6.2 to match. Kept exact rather than floating so the image stays reproducible. Note this ref also feeds the FA3 (Hopper) wheel, which is built by default, so the bump affects the default x86 image and not just INSTALL_FA4=true builds. Signed-off-by: Alexandros Koumparoulis <akoumparouli@nvidia.com> * build: pin full flash-attention commit Signed-off-by: Alexandros Koumparoulis <akoumparouli@nvidia.com> * build: bump cutlass-dsl to 4.6.2, ffpa-attn to 0.2.3, quack-kernels to 0.6.4 flash_attn.cute (FA4) at the pinned FLASH_ATTN_REF builds against nvidia-cutlass-dsl 4.6.2. quack-kernels and ffpa-attn both import the same `cutlass` package, so all three have to agree on one version. - quack-kernels 0.6.1 -> 0.6.4 (pins cutlass-dsl 4.6.2; 0.6.1 pinned 4.6.0) - ffpa-attn 0.2.2 -> 0.2.3 (first release pinning cutlass-dsl 4.6.2 and quack-kernels 0.6.4; 0.2.2 hard-pinned 4.6.0/0.6.1) - ffpa extra's explicit cutlass-dsl pin 4.6.0 -> 4.6.2 Because ffpa-attn 0.2.3 already declares the versions FA4 needs, the resolution is conflict-free -- no override-dependencies entries required. Corrects the apache-tvm-ffi note as well: FA4 declares apache-tvm-ffi>=0.1.12 but never imports tvm-ffi, so the <=0.1.11 cap that keeps tilelang working is not mutually exclusive with FA4. Both lock files regenerated with uv 0.8.22 (the version CI pins); the diff is confined to the six bumped packages. Signed-off-by: Alexandros Koumparoulis <akoumparouli@nvidia.com> * build: enable FA4 (INSTALL_FA4=true) by default Flips the default in both stages that declare the ARG: wheel_builder (gates building the flash_attn/cute wheel) and automodel_final (gates the nvidia-cutlass-dsl[cu13] install and the flash_attn/cute symlink). Docker does not inherit ARGs across stages, so flipping only the first would build the wheel while skipping the DSL install and the symlink -- a half-installed FA4. Both FA4 wheels are installed --no-deps, so flash-attn-4's declared apache-tvm-ffi>=0.1.12 bound is never resolved and the <=0.1.11 cap that keeps the tilelang kernels working still holds. Signed-off-by: Alexandros Koumparoulis <akoumparouli@nvidia.com> * build: correct the apache-tvm-ffi comment The note claimed flash-attn-4 never imports tvm-ffi. That is wrong: flash_attn/cute/cache_utils.py has a bare top-level `import tvm_ffi`, reached eagerly through flash_attn.cute.__init__ -> interface.py, so importing flash_attn.cute hard-requires it. The cap itself is still right, but for a different reason. FA4 is installed --no-deps, so its declared >=0.1.12 floor is never resolved, and the tvm_ffi that is present comes via quack-kernels (>=0.1.6,<0.2) pinned here to 0.1.11. FA4 therefore runs against an ffi older than it declares; the two symbols it references, tvm_ffi.Function and tvm_ffi.__version__, both exist in 0.1.11. Note this is unexercised: nothing selects attn_implementation="flash_attention_4" until the FA4 backend lands, so flash_attn.cute is never imported by any test here. Worth knowing given docker/Dockerfile sets FLASH_ATTENTION_CUTE_DSL_CACHE_ENABLED=1, which turns on the AOT cache path in that same cache_utils.py. Verified against flash-attention ce088ab9, the pinned FLASH_ATTN_REF. Comment only; no dependency or lock change. Signed-off-by: Alexandros Koumparoulis <akoumparouli@nvidia.com> --------- Signed-off-by: Alexandros Koumparoulis <akoumparouli@nvidia.com>
This file contains hidden or bidirectional Unicode text that may be interpreted or compiled differently than what appears below. To review, open the file in an editor that reveals hidden Unicode characters.
Learn more about bidirectional Unicode characters
Sign up for free
to join this conversation on GitHub.
Already have an account?
Sign in to comment
2 participants
Add this suggestion to a batch that can be applied as a single commit.This suggestion is invalid because no changes were made to the code.Suggestions cannot be applied while the pull request is closed.Suggestions cannot be applied while viewing a subset of changes.Only one suggestion per line can be applied in a batch.Add this suggestion to a batch that can be applied as a single commit.Applying suggestions on deleted lines is not supported.You must change the existing code in this line in order to create a valid suggestion.Outdated suggestions cannot be applied.This suggestion has been applied or marked resolved.Suggestions cannot be applied from pending reviews.Suggestions cannot be applied on multi-line comments.Suggestions cannot be applied while the pull request is queued to merge.Suggestion cannot be applied right now. Please check back later.
beep boop [馃]: Hi @akoumpa 馃憢,