MPL2 License#10
Merged
Merged
Conversation
mmathew23
added a commit
to mmathew23/unsloth-zoo
that referenced
this pull request
May 22, 2026
Multi-reviewer pass on the autocast wrapper / norm-upcast path: - Instance-level forward (#2): an instance attribute `model.forward` (Unsloth runtime forward patching) shadows class-method overrides, so mutating __class__ silently bypassed the wrapper -> fp32 norm met a bf16 linear with no autocast and crashed. Now wrap the instance attribute when present; otherwise subclass as before. - Wrapper gating (unslothai#5, unslothai#7): install the wrapper iff fp32 norm params actually exist (from our upcast, the legacy env upcast, or an external _pre_set_compute_dtype policy) -- not on the upcast DECISION. Fixes the rollback path leaving external fp32 norms exposed, and stops wrapping models with no fp32 norm. Add _unwrap_forward_in_bf16_autocast for re-prepare (unslothai#10). - config.architectures leak (unslothai#8/unslothai#9): keep the original __name__ on the generated subclass (unique __qualname__ for registration) so save_pretrained records the base architecture. - Device detection (unslothai#11): recurse into mapping/list/tuple batches and fall back to the model's parameter device instead of defaulting to "cuda". - Legacy UNSLOTH_UPCAST_LAYERNORM (#1/#3/#4): route through the shared _cast_named_module + union matcher and honour the external-policy deferral. - Recursive external-ownership guard (unslothai#6): record descendants of tagged modules (the external policy casts recursively). - Fresh-interpreter pickle test (unslothai#12): real subprocess load. Shared helpers: _find_tensor_device_type, _call_forward_with_bf16_autocast, _canonical_module_name, _cast_named_module. Unit suite: 25 passed.
danielhanchen
added a commit
that referenced
this pull request
Jun 16, 2026
* fix fft norms to fp32 * address reviewer comments * remove old tests * fix(full-FT): make bf16 autocast wrapper picklable, broaden norm matcher Address PR review feedback: - Codex (P2): assigning a runtime type(...) subclass to model.__class__ made the model unpicklable, so torch.save(model, ...) raised PicklingError on the bf16 full-FT path. Cache the generated subclass per (base_class, compute_dtype) and register it as a module-level symbol so pickle can resolve it by module + qualname. deepcopy survival is preserved and the subclass identity is now stable across instances. - gemini-code-assist: avoid allocating list(args) + list(kwargs.values()) on every forward; iterate args then kwargs sequentially with for/else. - gemini-code-assist: drop the leading dot from the norm1/norm2 patterns so top-level norm1.weight / norm2.weight are matched, consistent with the norm. and _layernorm patterns. Add regression tests: picklability, subclass caching across instances, and top-level norm1/norm2 matching. * fix(full-FT): detect norms by module class, not just param-name substrings Multi-architecture FFT validation surfaced a real norm-coverage gap: the Gemma-4 / Gemma-3n audio tower uses RMSNorm modules named norm_out, norm_pre_attn, and norm_post_attn. None of those param names match the substring patterns (norm., _layernorm, layer_norm, norm1., norm2.), so 36 audio-tower norms were left in bf16 for bf16 full finetuning. Add owning-module class-name detection (RMSNorm/LayerNorm) as a union with the existing name patterns. This catches any norm module regardless of param naming while keeping the established patterns. Verified across Qwen3, Qwen3-VL, Qwen3.5, Gemma-3 (text+vision), Gemma-4, and gpt-oss-20B-BF16 MoE: expanded-norm misses drop to 0 with zero new false positives (gpt-oss expert biases remain fp32 from transformers load-time _keep_in_fp32, not the matcher). Add a regression test for class-based detection of unmatched norm names. * fix(full-FT): make autocast wrapper class reconstructable across processes Codex P1: the generated autocast subclass was only registered via runtime setattr in _make_bf16_autocast_subclass, so a torch.save(model) pickle could not be torch.load'd in a fresh interpreter -- unpickling tried to resolve unsloth_zoo.training_utils.<GeneratedName> before that function had ever run in the new process, raising AttributeError. Same-process load worked, masking the issue. Add __reduce__ to the generated subclass so instances serialize via their IMPORTABLE base class plus standard nn.Module state, and reconstruct through _reconstruct_bf16_autocast_model (which rebuilds the subclass on demand). No dynamically-generated symbol needs to be resolvable at unpickle time, so cross-process / cross-session checkpoint handoff works. The module-level registration is kept for direct class pickling and same-process paths. Verified with a fresh-interpreter reproduction (save in one process, load in another that never wrapped a model) and a unit test that drops the runtime registration + cache before unpickling. Full suite: 18 passed. * fix(full-FT): address reviewer findings on the bf16 autocast wrapper Multi-reviewer pass on the autocast wrapper / norm-upcast path: - Instance-level forward (#2): an instance attribute `model.forward` (Unsloth runtime forward patching) shadows class-method overrides, so mutating __class__ silently bypassed the wrapper -> fp32 norm met a bf16 linear with no autocast and crashed. Now wrap the instance attribute when present; otherwise subclass as before. - Wrapper gating (#5, #7): install the wrapper iff fp32 norm params actually exist (from our upcast, the legacy env upcast, or an external _pre_set_compute_dtype policy) -- not on the upcast DECISION. Fixes the rollback path leaving external fp32 norms exposed, and stops wrapping models with no fp32 norm. Add _unwrap_forward_in_bf16_autocast for re-prepare (#10). - config.architectures leak (#8/#9): keep the original __name__ on the generated subclass (unique __qualname__ for registration) so save_pretrained records the base architecture. - Device detection (#11): recurse into mapping/list/tuple batches and fall back to the model's parameter device instead of defaulting to "cuda". - Legacy UNSLOTH_UPCAST_LAYERNORM (#1/#3/#4): route through the shared _cast_named_module + union matcher and honour the external-policy deferral. - Recursive external-ownership guard (#6): record descendants of tagged modules (the external policy casts recursively). - Fresh-interpreter pickle test (#12): real subprocess load. Shared helpers: _find_tensor_device_type, _call_forward_with_bf16_autocast, _canonical_module_name, _cast_named_module. Unit suite: 25 passed. * fix(full-FT): end-anchor module-name suffix strip; rebind instance forward on deepcopy Two P1 regressions from the previous reviewer-fix commit: - _canonical_module_name used `.replace(".weight"/".bias", "", 1)`, which strips the FIRST occurrence anywhere in the name. Module names that merely contain those as a substring (e.g. `...self_attn.bias_proj.weight`, `...weight_scale.weight`, `...bias_layernorm.weight`) were rewritten to a bogus attribute path, so _cast_named_module could raise at runtime. Strip only a TRAILING `.weight`/`.bias` (end-anchored). - The instance-level forward wrapper closed over the original `model`, so a deep-copied model's forward ran against the original's parameters. Bind the wrapper as a types.MethodType and read the original forward from `self`, so deepcopy rebinds __self__ to the copy (EMA / model averaging / checkpoint). Tests: end-anchored canonicalizer, a model with bias_/weight_ substring module names (must not crash), and instance-forward deepcopy rebind. Suite: 28 passed. * refactor(full-FT): cast param.data directly; route instance-forward through subclass Cleaner caster + fixes the latest reviewer P1 (instance-forward unpicklable): - Replace the param-name -> module-attribute-path resolver (_canonical_module_name) and exec()-based _cast_named_module with a direct `param.data = param.data.to(dtype)`. This is what nn.Module._apply does for a dtype cast: preserves Parameter identity (tied weights stay tied), no string munging, and removes the whole class of path-resolution bugs (e.g. `.bias`/`.weight` substrings in module names). Qwen3-0.6B/1.7B losses are bit-identical before/after, confirming behavior is preserved. - Instance-level forward wrapper no longer stores a bound local function in __dict__ (which torch.load cannot resolve by import path in a fresh process). Instead it moves the original forward to `_unsloth_autocast_orig_forward`, drops the instance attribute, and routes through the generated subclass in a new `instance_mode` (forward is a module-level function read from `self`). The model stays picklable AND deepcopy-safe via the existing __reduce__ path. Tests: drop the now-removed _canonical_module_name test; add instance-forward pickle (same-process + fresh subprocess) tests; keep the substring-name integration test. Suite: 29 passed. Re-validated multi-arch FFT on transformers 4.56.2 and 5.5.4 (Qwen3, Qwen3-VL, Qwen3.5, Gemma-3/4, gpt-oss-20B-BF16): 0 norm misses, no architecture-name leak, pickle + save_pretrained OK. * Condense comments and docstrings in bf16 full-FT norm upcast code Trim the verbose comments and docstrings added for the fp32 LayerNorm upcast and the autocast forward wrapper so they are shorter and clearer without losing intent. No behavior change: the executable AST is identical and all 29 regression tests still pass. * Re-run CI for fix/fftnorms --------- Co-authored-by: Daniel Han <danielhanchen@gmail.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
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.
No description provided.