Lightning-AI/pytorch-lightning · error · NotImplementedError

DeepSpeed handles gradient clipping automatically within the

Error message

DeepSpeed handles gradient clipping automatically within the optimizer. Make sure to set the `gradient_clipping` value in your Config.

What it means

DeepSpeed performs gradient clipping inside its engine/optimizer using the gradient_clipping entry of the DeepSpeed config. Lightning therefore explicitly does not implement clip_gradients_norm for this strategy and raises NotImplementedError if called.

Source

Thrown at src/lightning/fabric/strategies/deepspeed.py:563

                " or a single checkpoint file by setting `DeepSpeedStrategy(..., load_full_weights=True)`."
            )

        # `Engine.load_checkpoint` adds useless keys 'optimizer' and 'lr_scheduler' to the client state; remove
        # them to avoid name collision with user state
        keys = set(client_state) & set(state) - {"optimizer", "lr_scheduler"}
        _move_state_into(source=client_state, destination=state, keys=keys)
        return client_state

    @override
    def clip_gradients_norm(
        self,
        module: "DeepSpeedEngine",
        optimizer: Optimizer,
        max_norm: Union[float, int],
        norm_type: Union[float, int] = 2.0,
        error_if_nonfinite: bool = True,
    ) -> torch.Tensor:
        raise NotImplementedError(
            "DeepSpeed handles gradient clipping automatically within the optimizer. "
            "Make sure to set the `gradient_clipping` value in your Config."
        )

    @override
    def clip_gradients_value(
        self, module: "DeepSpeedEngine", optimizer: Optimizer, clip_val: Union[float, int]
    ) -> None:
        raise NotImplementedError(
            "DeepSpeed handles gradient clipping automatically within the optimizer. "
            "Make sure to set the `gradient_clipping` value in your Config."
        )

    @classmethod
    @override
    def register_strategies(cls, strategy_registry: _StrategyRegistry) -> None:
        strategy_registry.register("deepspeed", cls, description="Default DeepSpeed Strategy")
        strategy_registry.register("deepspeed_stage_1", cls, description="DeepSpeed with ZeRO Stage 1 enabled", stage=1)

View on GitHub (pinned to 9fed5c27d2)

Solutions

  1. Set gradient clipping in the DeepSpeed config: config['gradient_clipping'] = <float> (add gradient_clipping to the JSON/dict)
  2. Remove manual calls to clip_gradients_norm on the strategy and let the DeepSpeed engine clip during optimizer.step()
  3. If using the Trainer, ensure Trainer(gradient_clip_val=...) matches the DeepSpeed config value

Example fix

# before
strategy.clip_gradients_norm(module, optimizer, max_norm=1.0)

# after
ds_config = {"gradient_clipping": 1.0, ...}
strategy = DeepSpeedStrategy(config=ds_config)
Defensive patterns

Strategy: validation

Validate before calling

ds_config.setdefault("gradient_clipping", 1.0)
assert "gradient_clipping" in ds_config, "set gradient_clipping in DeepSpeed config"

Type guard

def supports_norm_clip(strategy) -> bool:
    return not type(strategy).__name__ == "DeepSpeedStrategy"

Prevention

When it happens

Trigger: Calling strategy.clip_gradients_norm(...) (directly, or via Fabric's Trainer/autocast gradient-clip path that is not DeepSpeed-aware) on a DeepSpeedStrategy instance.

Common situations: User code or a custom loop calls fabric.clip_gradients... or the strategy method manually; a LightningModule sets Trainer(gradient_clip_val=...) without the corresponding gradient_clipping key in the DeepSpeed config; porting code from DDPStrategy to DeepSpeed.

Related errors


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