Buckets:
| """ComfyUI-style handling for mixed-precision checkpoints (e.g. fp32 | |
| scale_shift tables next to bf16 linears in one file). | |
| Two pieces, used together on a loaded module tree: | |
| - attach_per_op_casting(root): every weighted module casts its floating-point | |
| inputs to its own weight dtype at forward (what comfy's manual-cast ops do), | |
| so activations promoted to fp32 by a stored-fp32 tensor drop back to the | |
| layer's dtype at the next op instead of erroring in a bf16 matmul. | |
| - pin_stored_fp32(root): tensors that are fp32 after load stay fp32 through | |
| any parent-level dtype cast (`module.to(dtype=...)`, `.half()`, ...) — | |
| device moves still apply. Without this, a holder's blanket `.to(dtype)` | |
| would silently downcast the deliberately-fp32 pieces. | |
| """ | |
| import types | |
| import torch | |
| from torch import nn | |
| from toolkit.util.ostris_quant import OstrisLinear | |
| def _cast_args(args, kwargs, dtype): | |
| def cast(t): | |
| if torch.is_tensor(t) and t.is_floating_point() and t.dtype != dtype: | |
| return t.to(dtype) | |
| return t | |
| return tuple(cast(a) for a in args), {k: cast(v) for k, v in kwargs.items()} | |
| def attach_per_op_casting(root: nn.Module) -> int: | |
| """Register forward-pre-hooks casting inputs to each weighted module's own | |
| weight dtype. Covers Linear/Conv/Norm-style modules that own a floating | |
| ``weight`` parameter, and OstrisLinear via its stored orig dtype (its | |
| ``weight`` property would materialize the dequantized tensor). Returns the | |
| number of modules hooked.""" | |
| hooked = 0 | |
| for module in root.modules(): | |
| if isinstance(module, OstrisLinear): | |
| def hook(mod, args, kwargs): | |
| return _cast_args(args, kwargs, mod.ostris_orig_dtype) | |
| else: | |
| weight = module._parameters.get("weight", None) | |
| if weight is None or not weight.is_floating_point(): | |
| continue | |
| def hook(mod, args, kwargs): | |
| return _cast_args(args, kwargs, mod._parameters["weight"].dtype) | |
| module.register_forward_pre_hook(hook, with_kwargs=True) | |
| hooked += 1 | |
| return hooked | |
| def _tag_fp32(root: nn.Module): | |
| for t in list(root.parameters()) + list(root.buffers()): | |
| if t.is_floating_point() and t.dtype == torch.float32: | |
| t._pin_dtype = True | |
| if isinstance(t, nn.Parameter): | |
| t.data._pin_dtype = True | |
| def pin_stored_fp32(root: nn.Module): | |
| """Mark every currently-fp32 param/buffer in ``root`` and wrap the root's | |
| ``_apply`` so parent dtype casts skip them (device moves still apply). | |
| The wrapper probes the cast function with a scalar to detect whether it | |
| changes float dtypes, so ``.to(device)`` passes through untouched.""" | |
| _tag_fp32(root) | |
| orig_apply = root._apply | |
| def _apply(self, fn, *args, **kwargs): | |
| probe = fn(torch.zeros((), dtype=torch.float32)) | |
| if probe.dtype == torch.float32: | |
| return orig_apply(fn, *args, **kwargs) | |
| # where fn sends a tensor depends on where it already is: a dtype-only | |
| # cast keeps a cuda tensor on cuda (probing from cpu would wrongly | |
| # report cpu and drag pinned tables off the gpu); an explicit device | |
| # move relocates it. Probe per source device. | |
| target_by_device = {} | |
| def _target_device(dev): | |
| if dev not in target_by_device: | |
| target_by_device[dev] = fn( | |
| torch.zeros((), dtype=torch.float32, device=dev) | |
| ).device | |
| return target_by_device[dev] | |
| def fn_pinned(t): | |
| if getattr(t, "_pin_dtype", False): | |
| target = _target_device(t.device) | |
| out = t if t.device == target else t.to(target) | |
| out._pin_dtype = True | |
| return out | |
| return fn(t) | |
| result = orig_apply(fn_pinned, *args, **kwargs) | |
| # _apply may rewrap tensors (dropping attributes); the invariant is | |
| # simple — the pinned set IS the fp32 set — so re-tag | |
| _tag_fp32(root) | |
| return result | |
| root._apply = types.MethodType(_apply, root) | |
| return root | |
Xet Storage Details
- Size:
- 4.14 kB
- Xet hash:
- 714b59a1b851ff420e3ac33c3fa3637f7852d75f3e4bf6322d5e07884e0481f4
·
Xet efficiently stores files, intelligently splitting them into unique chunks and accelerating uploads and downloads. More info.