Skip to content

Fixes for several issues: mask compile, cache offload, buffers, FSDP2 eval, trainer teardown, MoE logits, tokenizers, processors - #48971

Closed
Ryukijano wants to merge 16 commits into
huggingface:mainfrom
Ryukijano:fix/mask-offload-fsdp2-buffers
Closed

Ryukijano wants to merge 16 commits into
huggingface:mainfrom
Ryukijano:fix/mask-offload-fsdp2-buffers

Conversation

@Ryukijano

Copy link
Copy Markdown
Contributor

What this PR does

This PR bundles a set of small, independent fixes for several open issues. Each fix is a separate signed-off commit with its own regression coverage, so the PR can be split by subsystem if preferred.

Mask handling under torch.compile (#48924, #48925)

  • create_causal_mask / bidirectional mask helpers now distinguish ordinary torch.compile (where is_tracing(None) is true but a None padding mask still permits skipping mask construction) from export/JIT/CUDA-graph capture, which genuinely requires static behavior.
  • modeling_opt no longer forwards the fabricated all-ones position-ID mask to create_causal_mask, preserving the None fast path.
  • Verified: sdpa_mask returns None under compile; compiled OPT forward matches eager (max diff ~1.2e-7).

Hybrid cache offloading of linear-attention states (#48947)

  • Cache.update synchronizes the offload prefetch stream only for KV updates. Hybrid layers also hold conv_states/recurrent_states that are updated through update_conv_state/update_recurrent_state and read directly by model code, so offloading them races and produces CUDA/CPU device mismatches (reproduced in torch.cat during decode).
  • LinearAttentionAnd*Layer.offload()/.prefetch() now manage only the dynamic KV state; missing implementations were added to the sliding hybrid layers, and Cache.is_linear excludes all hybrid classes.
  • Regression test added in tests/utils/test_cache_utils.py.

Non-persistent buffer materialization (#48892)

  • _move_missing_keys_from_meta_to_device materializes non-persistent buffers with torch.zeros_like instead of torch.empty_like (garbage memory), and warns once for non-empty buffers that remain all-zero after initialization, i.e. not covered by _init_weights. Rotary inv_freq buffers do not false-positive.

FSDP2 evaluation before training (#48841)

  • evaluation_loop now creates the optimizer and calls accelerator.prepare(model, optimizer) when FSDP2 is active and the model has not been prepared yet, matching the training path so FSDP2 rewrites optimizer parameter references consistently. (Supersedes the approach in the closed Fix FSDP2 evaluation before training #48858.)

Trainer teardown / process-group lifecycle (#48874)

  • New Trainer.end() (idempotent) calls accelerator.end_training(), which finishes trackers and destroys the distributed process group. close() is an alias; Trainer is now a context manager (with Trainer(...) as trainer:).
  • An atexit fallback destroys the process group at interpreter exit only when accelerate owns it (a user-initialized torch.distributed group without PartialState is left alone). The group is deliberately not destroyed in _finalize_training since evaluate/predict/save_model/push_to_hub still need it.
  • Verified with single-process gloo repros; the NCCL destructor warning itself isn't observable on this platform.

_no_placement_params in balanced memory (#48756)

  • get_balanced_memory accepts no_placement_params and subtracts excluded-parameter sizes from the total, the holding no-split module, and the holding leaf module — so a ~100 GB never-placed table (e.g. Qwen4-Exp's n-gram embedding) no longer inflates every device's budget. Adapted to this branch's newer biggest_leaf buffer logic; regression test added in test_modeling_utils.py.

base_model_prefix corrections (#48862)

  • rename_source_key now requires prefix + "." when stripping (prevents matching model2. for prefix model) and escapes the prefix in the regex.
  • Corrected prefixes to match actual submodule attribute names: CLIPSeg → clip, PPOCRv5/v6 det → model, TimesFM → decoder, UVDoc → backbone, VideoPrism → vision_model/text_model/video_model.

MoE router logits / aux loss (#48886, #48889)

  • cohere2_moe gains the output_router_logits config flag + forward plumbing its code already honoured.
  • afmoe/zaya (propagated) return explicit aux_loss=None: these models balance experts with per-expert bias (expert_bias, e_score_correction_bias, balancing_biases), so no aux loss is computed — router_logits are still returned for inspection/external balancing. Wiring a Switch-style loss would also crash zaya (num_experts+1 router logits incl. skip expert).
  • MoeCausalLMOutputWithPast docstring updated accordingly. This intentionally matches the approach on origin/moe-router-logits-exposed rather than adding losses these recipes don't use.

Gemma4 audio eager mask (#48818)

  • Gemma4AudioAttention treated the 4-D float additive mask (0 = keep, finfo.min = masked) as a boolean keep-mask and inverted it; _convert_4d_mask_to_blocked_5d also padded float masks with False (= 0.0 = keep). Float masks are now normalized to boolean keep-masks (mask == 0). Eager vs SDPA max diff: 0.358 → 0.0.

Image/video processor size dict mutation (#48917)

  • Processors mutating size["min_pixels"]/["max_pixels"] now copy the dict first (size = dict(size) if isinstance(size, dict) else size), so class-level defaults and caller-passed dicts are no longer mutated. Applied to qwen2_vl (image ×2 + video), minimax_m3_vl (incl. the [672,672] legacy path), ernie4_5_vl_moe video, and all modular-derived processors (video_llama_3, paddleocr_vl, hunyuan_vl, glm_image, cohere_compass). Unlike Fix image processor class-level size mutation and min_pixels handling #48916 this guards non-dict size inputs too (int/list still valid per ImagesKwargs).

Empty SentencePiece precompiled_charsmap (#48942)

  • Protobuf returns b"" for unset bytes fields, so b"" — not None — reached every if _spm_precompiled_charsmap is not None: guard and normalizers.Precompiled(b"") crashed. Guards switched to truthiness in 14 tokenizers (t5, xlnet, xlm_roberta, xglm, videoprism, rembert, reformer, pegasus, nllb, mbart50, lasr [modular+generated], camembert, albert). convert_slow_tokenizer.py already used truthiness.

Luke/MLuke entity tokens (#48225)

  • entity_token_1/entity_token_2 IDs are resolved by token name via convert_tokens_to_ids instead of positional extra_special_tokens_ids[0]/[1], which broke whenever other extra special tokens occupied those positions. Applied to both mluke and luke (same bug).

Chat CLI requests dependency (#48283)

  • transformers.cli.chat no longer hard-requires requests: streaming now uses huggingface_hub.utils.httpx (already a transitive dep) with a context-managed httpx.stream, preserving indefinite streaming and error handling. transformers env/version verified working with requests blocked.

Test robustness

Notes on overlap with existing PRs

Several issues have open competing PRs (#48916/#48921, #48871, #48875, #48758, #48958, #48928#48940, #48932, #48922, #48318/#48412, #48389, #48530/#48543/#48587/#48944, #48226/#48230/#48235/#48408). Where they exist this implementation generally differs (e.g. guarded dict copy vs unconditional, finfo.min-aware mask normalization, escaped-prefix regex fix, upstream-aligned MoE approach, ownership-gated atexit teardown). Happy to split this into per-issue PRs or drop commits where a competing approach is preferred — tell me which direction you want.

Testing

  • Targeted regression tests pass for each commit (cache offload test, balanced-memory test, compiled OPT forward, tokenizer/processor suites: e.g. Gemma4 141 passed + 3280 subtests, tokenizers 213 + 509 passed).
  • ruff check and ruff format --check clean on all 68 changed files.

…ts does not cover them [huggingfaceGH-48892]

Signed-off-by: Gyanateet <gyanateet@gmail.com>
Signed-off-by: Gyanateet <gyanateet@gmail.com>
Signed-off-by: Gyanateet <gyanateet@gmail.com>
…OutputWithPast docs [huggingfaceGH-48886] [huggingfaceGH-48889]

Signed-off-by: Gyanateet <gyanateet@gmail.com>
…process group [huggingfaceGH-48874]

Signed-off-by: Gyanateet <gyanateet@gmail.com>
Copilot AI lite review requested due to automatic review settings September 21, 2026 00:09

Copilot AI left a comment

Copy link
Copy Markdown
Contributor

Choose a reason for hiding this comment

The reason will be displayed to describe this comment to others. Learn more.

Copilot was unable to review this pull request because the user who requested the review has reached their quota limit.

@github-actions

Copy link
Copy Markdown
Contributor

Thank you for your contribution 🤗!

CI Security Gate — automatic approval blocked

This PR was not automatically approved for CI because the security gate failed.

Possible reasons:

  • The PR touches 50 or more files — only PRs with fewer than 50 changed files are automatically approved
  • A changed file is outside the allowed directories (src/, tests/, docs/, utils/), has a disallowed extension (only .py, .txt, .md permitted outside tests/ and docs/), or is not .md/.yml inside docs/ — this covers files the PR deletes or renames, not only the ones it edits
  • A new high-severity security issue was detected in the changed Python files (Bandit check)
  • The PR touches a path this repository protects from untrusted PRs, such as the file that decides who reviews it — a maintainer must make that change in a separate PR

See the workflow run for the exact violations.

A maintainer can review and manually approve CI if a finding is a false positive.

@github-actions

Copy link
Copy Markdown
Contributor

[For maintainers] Suggested jobs to run (before merge)

run-slow: afmoe, albert, bark, camembert, clipseg, cohere2_moe, cohere_compass, ernie4_5_vl_moe, gemma4, glm_image, hunyuan_vl, lasr, luke, mbart50, mellum, minimax_m3_vl

@Ryukijano
Ryukijano deleted the fix/mask-offload-fsdp2-buffers branch September 21, 2026 11:47
Sign up for free to join this conversation on GitHub. Already have an account? Sign in to comment

Projects

None yet

Development

Successfully merging this pull request may close these issues.

3 participants