Lightning-AI / Lightning-AI/pytorch-lightning

Trainer validates `gradient_clip_algorithm` although `configure_gradient_clipping` is defined

Open
#10,919 5 comments 0 reactions 1 assignee View on GitHub

Nobody has claimed this yet.

bug
Dominant language
Python
Stars
31.4k
Forks
3.8k
Avg merge
6d 7h
Merged PRs (30d)
6

Description

🐛 Bug

I can't pass a custom gradient clipping algorithm, although I implemented configure_gradient_clipping hook, Documentation and release notes hints that you can use configure_gradient_clipping to implement your custom gradient clipping algorithm (.release notes: ... This means you can now implement state-of-the-art clipping algorithms with Lightning! ...) .

Please, allow to pass custom algorithm names in gradient_clip_algorithm when the model implements configure_gradient_clipping.

To Reproduce
from pytorch_lightning import LightningModule, Trainer

class BoringModel(LightningModule):
    def __init__(self):
        super().__init__()
        self.layer = torch.nn.Linear(32, 2)

    def forward(self, x):
        return self.layer(x)

    def training_step(self, batch, batch_idx):
        loss = self(batch).sum()
        self.log("train_loss", loss)
        return {"loss": loss}

    def validation_step(self, batch, batch_idx):
        loss = self(batch).sum()
        self.log("valid_loss", loss)

    def test_step(self, batch, batch_idx):
        loss = self(batch).sum()
        self.log("test_loss", loss)

    def configure_optimizers(self):
        return torch.optim.SGD(self.layer.parameters(), lr=0.1)

    def configure_gradient_clipping(
        self,
        optimizer: Optimizer,
        optimizer_idx: int,
        gradient_clip_val: Optional[Union[int, float]] = None,
        gradient_clip_algorithm: Optional[str] = None,
    ):     
        if gradient_clip_algorithm == "my_custom_clipping_algorithm":
            my_custom_clipping(optimizer, gradient_clip_val, gradient_clip_algorithm)
        else:  # Lightning will handle the gradient clipping
            self.clip_gradients(
                optimizer,
                gradient_clip_val=gradient_clip_val,
                gradient_clip_algorithm=gradient_clip_algorithm
            )


trainer = Trainer(
        default_root_dir=os.getcwd(),
        limit_train_batches=1,
        limit_val_batches=1,
        limit_test_batches=1,
        num_sanity_val_steps=0,
        max_epochs=1,
        enable_model_summary=False,
        gradient_clip_algorithm="my_custom_clipping_algorithm"
    )
Expected behavior
Environment
  • PyTorch Lightning Version: 1.5.4
Additional context

I want to train NFNets with Adaptive Gradient Clipping and compare with standard L2 gradient clipping.

Contributor guide

Open the contributing guide

First steps

  1. Read the whole issue, then the project's contributing guide.
  2. Comment on the issue to say you are picking it up — it saves two people doing the same work.
  3. Fork the repository and make your change on a branch.
  4. Open a pull request that references the issue number.

Assessment

This issue has not been assessed yet.

Get new issues in your inbox

A short digest of beginner-friendly GitHub issues.