[torchao]: fix how FqnConfig is resolved and handled - #14686
Conversation
github-actions
Bot
added
quantization
tests
pipelines
fixes-issue
CI
size/L
0e34e7d to
a2b9a5f
Compare
github-actions
Bot
added
size/M
| 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." | ||
| ) |
|
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
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.
|
interestingly it seems to completely ignore quantization configs for the MiniMax H3 "vae" but I haven't investigated why yet. |
|
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 |
|
If a given VAE has no linear layers then nothing will get applied btw because TorchAO only focuses linear layers. |
it's a 10gb transformer with the usual vae resnet downsampling bolted onto the end https://huggingface.co/MiniMaxAI/MiniMax-H3/tree/main/vae |
|
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.
3 participants
Fixes #14667. It surfaced additional issues which I have fixed as well.
I have run the TorchAO test suite and it's green.