[Quantization] - Added uses_meta_device_weights to quant config

[Josephasafg]

Purpose

As more quant methods are starting to support online quantization we need a more robust way to check that they are loading dummy weights in the same way using process_weights_after_loading

Test Plan

Test Result

---

Essential Elements of an Effective PR Description Checklist

• The purpose of the PR, such as "Fix some issue (link existing issues this PR will resolve)".
• The test plan, such as providing test command.
• The test results, such as pasting the results comparison before and after, or e2e results
• (Optional) The necessary documentation update, such as updating supported_models.md and examples for a new model.
• (Optional) Release notes update. If your change is user facing, please update the release notes draft in the Google Doc.

[gemini-code-assist]

Code Review

This pull request introduces a uses_meta_device_weights method to the quantization configuration, providing a more robust way to handle online quantization. The change refactors a hardcoded check for fp8 quantization to a more generic mechanism. The implementation is sound, but there's a critical issue where a None value for model_config.quantization could cause a crash. I've provided a suggestion to handle this case gracefully.

[Josephasafg]

@vkuzo Thanks for the review!

Who should trigger the CI?

[mgoin]

@Josephasafg @vkuzo I would prefer to keep the information on the linear method itself rather than the top-level quant config. What do you think about this proposal

1. Add uses_meta_device: bool = False to QuantizeMethodBase in base_config.py
2. Set uses_meta_device = True on Fp8OnlineLinearMethod and Fp8OnlineMoEMethod (the methods that actually create weights on device="meta")
3. In initialize_dummy_weights, instead of checking the quant config class, iterate over model.modules() and check if any module has a quant_method with uses_meta_device = True

Something like this:

def initialize_dummy_weights(model, model_config, ...):
    meta_device_params: set[int] = set()
    for module in model.modules():
        qm = getattr(module, "quant_method", None)
        if qm is not None and getattr(qm, "uses_meta_device", False):
            for param in module.parameters(recurse=False):
                meta_device_params.add(id(param))

    for param in model.state_dict().values():
        if id(param) in meta_device_params \
                and param.device == torch.device("meta"):
            continue
        initialize_single_dummy_weight(param, low, high, seed)

[Josephasafg]

@mgoin @vkuzo I made the change but made it a little simpler. How does this look?

def initialize_dummy_weights(
    model: torch.nn.Module,
    model_config: ModelConfig,
    low: float = -1e-3,
    high: float = 1e-3,
    seed: int = 1234,
) -> None:
    def uses_meta_device(module: torch.nn.Module) -> bool:
        quant_method = getattr(module, "quant_method", None)
        return getattr(quant_method, "uses_meta_device", False)

    has_online_quant = any(uses_meta_device(m) for m in model.modules())

    for param in model.state_dict().values():
        if has_online_quant and param.device == torch.device("meta"):
            # For online quantization, weights are created on meta device and
            # dummy weight init will happen in `process_weights_after_loading`.
            continue

        initialize_single_dummy_weight(param, low, high, seed)

[mgoin]

Nice work, I'm quite happy with this!