Skip to content

MPL2 License#10

Merged
danielhanchen merged 2 commits into
mainfrom
nightly
Nov 1, 2024
Merged

MPL2 License#10
danielhanchen merged 2 commits into
mainfrom
nightly

Conversation

@danielhanchen

Copy link
Copy Markdown
Member

No description provided.

@danielhanchen danielhanchen merged commit 0119b45 into main Nov 1, 2024
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>
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