Uh oh!
There was an error while loading. Please reload this page.
[torchao]: fix how FqnConfig is resolved and handled - #14686
Conversation
0e34e7d to
a2b9a5fCompare| import torch | ||
| import torch.nn as nn | ||
| if is_torch_version(">=", "2.5"): |
| return None | ||
| def _fqn_to_config_weight_sizes(config: "FqnToConfig") -> tuple[set[str | None], bool]: |
There was a problem hiding this comment.
Needed to determine CUDA warmup factor from FqnConfig.
| return size_digits, leaves_modules_unquantized | ||
| def _resolve_fqn_to_config(config: "FqnToConfig", module_fqn: str, param_fqn: str): |
There was a problem hiding this comment.
Needed so that we can properly resolve config for a single linear layer because that is how we create the quantized params, i.e., iterating over the modules and calling quantize_() over them individually.
| if isinstance(target_dtype, SUPPORTED_TORCH_DTYPES_FOR_QUANTIZATION): | ||
| return target_dtype | ||
| # We need one of the supported dtypes to be selected in order for accelerate to determine | ||
| # the total size of modules/parameters for auto device placement. | ||
| possible_device_maps = ["auto", "balanced", "balanced_low_0", "sequential"] | ||
| raise ValueError( | ||
| f"You have set `device_map` as one of {possible_device_maps} on a TorchAO quantized model but a suitable target dtype " | ||
| f"could not be inferred. The supported target_dtypes are: {SUPPORTED_TORCH_DTYPES_FOR_QUANTIZATION}. If you think the " | ||
| f"dtype you are using should be supported, please open an issue at https://github.com/huggingface/diffusers/issues." | ||
| ) |
HuggingFaceDocBuilderDev
commented
Sep 2, 2026
The docs for this PR live here. All of your documentation changes will be reflected on that endpoint. The docs are available until 30 days after the last update. |
Beinsezii
left a comment
There was a problem hiding this comment.
Initial tests seem to match what i expected from my own ai generated solution. Think gtg for now if I find any other corner cases I can open another issue.
Uh oh!
There was an error while loading. Please reload this page.
Beinsezii
commented
Sep 2, 2026
interestingly it seems to completely ignore quantization configs for the MiniMax H3 "vae" but I haven't investigated why yet. |
sayakpaul
commented
Sep 3, 2026
|
Beinsezii
commented
Sep 3, 2026
not yet just something i noticed running H3, probably blacklisted in modular or has everything set to keep_in_f32. when i get time i can make a repro that doesn't need 200gb of weights.
No, think this is unrelated, shouldn't block PR |
sayakpaul
commented
Sep 3, 2026
If a given VAE has no linear layers then nothing will get applied btw because TorchAO only focuses linear layers. |
Beinsezii
commented
Sep 3, 2026
it's a 10gb transformer with the usual vae resnet downsampling bolted onto the end https://huggingface.co/MiniMaxAI/MiniMax-H3/tree/main/vae |
sayakpaul
commented
Sep 3, 2026
Once I have better coverage on #14686 (comment), I can look into it. |
| # As we perform quantization here, the repr of linear layers is set by TorchAO, so we don't have to do it ourselves | ||
| module._parameters[tensor_name] = torch.nn.Parameter(param_value).to(device=target_device) | ||
| quantize_(module, self.quantization_config.get_apply_tensor_subclass()) | ||
| module._parameters[tensor_name] = torch.nn.Parameter(param_value.to(device=target_device)) |
There was a problem hiding this comment.
The path for pre-quantized checkpoints also does the same.
Fixes#14667. It surfaced additional issues which I have fixed as well.
I have run the TorchAO test suite and it's green.