Lightning-AI/pytorch-lightning · error · RuntimeError

Instantiating your model under the `init_module` context man

Error message

Instantiating your model under the `init_module` context manager is not supported when used with `BitsandbytesPrecision(..., ignore_modules={self.ignore_modules})` as this may initialize the layers on-device, defeating the purpose of quantization. You can remove `ignore_modules` or remove the `init_module` context manager.

What it means

BitsandbytesPrecision quantizes by monkey-patching nn.Linear during module creation under init_module. If ignore_modules is set, the patch cannot be applied selectively, so un-ignored Linears could be initialized on-device unquantized — defeating the purpose. Lightning therefore forbids the combination.

Source

Thrown at src/lightning/fabric/plugins/precision/bitsandbytes.py:133

            # this will not quantize the model but only replace the layer classes
            _convert_layers(module, self._linear_cls, self.ignore_modules)

        # set the compute dtype if necessary
        for m in module.modules():
            if isinstance(m, bnb.nn.Linear4bit):
                m.compute_dtype = self.dtype
                m.compute_type_is_set = False
        return module

    @override
    def tensor_init_context(self) -> AbstractContextManager:
        return _DtypeContextManager(self.dtype)

    @override
    def module_init_context(self) -> AbstractContextManager:
        if self.ignore_modules:
            # cannot patch the Linear class if the user wants to skip some submodules
            raise RuntimeError(
                "Instantiating your model under the `init_module` context manager is not supported when used with"
                f" `BitsandbytesPrecision(..., ignore_modules={self.ignore_modules})` as this"
                " may initialize the layers on-device, defeating the purpose of quantization. You can remove"
                " `ignore_modules` or remove the `init_module` context manager."
            )
        dtype_ctx = self.tensor_init_context()
        # TODO: this could also support replacing `Embedding` and `Conv1D`
        context_manager = _ClassReplacementContextManager({"torch.nn.Linear": self._linear_cls})
        stack = ExitStack()
        stack.enter_context(dtype_ctx)
        stack.enter_context(context_manager)
        return stack

    @override
    def forward_context(self) -> AbstractContextManager:
        return _DtypeContextManager(self.dtype)

    @override

View on GitHub (pinned to 9fed5c27d2)

Solutions

1. Remove ignore_modules and rely on layer-type replacement only (keep init_module)
2. Drop the init_module context manager (instantiate the model normally) if you must keep ignore_modules
3. Instantiate the model first, then manually replace/quantize modules, keeping ignore_modules semantics

Example fix

# before
plugin = BitsandbytesPrecision(mode="nf4", ignore_modules={"lm_head"})
with fabric.init_module():
    model = GPT2(config)  # RuntimeError

# after
plugin = BitsandbytesPrecision(mode="nf4")
with fabric.init_module():
    model = GPT2(config)

plugin = BitsandbytesPrecision(mode="nf4")
use_init_module = not plugin.ignore_modules
if use_init_module:
    with fabric.init_module():
        model = build_model()
else:
    model = build_model()

def can_use_init_module(plugin) -> bool:
    return not bool(plugin.ignore_modules)

When it happens

Trigger: BitsandbytesPrecision(..., ignore_modules={...}) used together with fabric.init_module() context manager when instantiating the model.

Common situations: Users excluding some submodules (e.g. lm_head) from quantization while also wanting meta-device/empty init via init_module to save memory before loading weights.

Related errors

• {mode!r} only works with `dtype=torch.float16`, but you chos
• You are using the bitsandbytes precision plugin, but your mo
• str(_BITSANDBYTES_AVAILABLE)
• Bitsandbytes is only supported on CUDA GPUs.
• Device should be CPU, got {device} instead.

AI-assisted analysis of Lightning-AI/pytorch-lightning@9fed5c27d2 (2026-08-28). Data as JSON: /api/errors/2c4aabebd0e9b953. Report an issue: GitHub.