Add LR scheduler to examples

README.md

@@ -85,6 +85,9 @@ To use with gradient release:
 opt = AdamW(model.parameters(), lr=1e-3, gradient_release=True)
 prepare_for_gradient_release(model, opt)
 
+# setup a learning rate scheduler like normal
+scheduler = CosineAnnealingLR(opt, ...)
+
 # calling backward on the model will peform the optimzier step
 loss = model(torch.randn(20, dtype=torch.bfloat16))
 loss.backward()
@@ -94,6 +97,9 @@ loss.backward()
 # opt.step()
 # opt.zero_grad()
 
+# step the learning rate scheduler like normal
+scheduler.step()
+
 # optionally remove gradient release hooks when done training
 remove_gradient_release(model)
 ```
@@ -110,6 +116,9 @@ prepare_for_gradient_release(model, opt)
 # gradients directly into the optimizer states
 accumulation_steps = 4
 
+# setup a learning rate scheduler for gradient accumulation
+scheduler = CosineAnnealingLR(opt, ...)
+
 # use existing PyTorch dataloader
 for idx, batch in enumerate(dataloader):
     # `optimizer_accumulation=True` accumulates gradients into
@@ -127,6 +136,10 @@ for idx, batch in enumerate(dataloader):
     # opt.step()
     # opt.zero_grad()
 
+    # step the learning rate scheduler after accumulating gradients
+    if not opt.optimizer_accumulation:
+        scheduler.step()
+
 # optionally remove gradient release hooks when done training
 remove_gradient_release(model)
 ```

docs/gradient_release.md

@@ -54,6 +54,9 @@ model = nn.Linear(20, 1, dtype=torch.bfloat16)
 opt = AdamW(model.parameters(), lr=1e-3, gradient_release=True)
 prepare_for_gradient_release(model, opt)
 
+# setup a learning rate scheduler like normal
+scheduler = CosineAnnealingLR(opt, ...)
+
 # calling backward on the model will peform the optimzier step
 loss = model(torch.randn(20, dtype=torch.bfloat16))
 loss.backward()
@@ -63,6 +66,9 @@ loss.backward()
 # opt.step()
 # opt.zero_grad()
 
+# step the learning rate scheduler like normal
+scheduler.step()
+
 # optionally remove gradient release hooks when done training
 remove_gradient_release(model)
 ```

docs/index.md

@@ -86,6 +86,9 @@ To use with gradient release:
 opt = AdamW(model.parameters(), lr=1e-3, gradient_release=True)
 prepare_for_gradient_release(model, opt)
 
+# setup a learning rate scheduler like normal
+scheduler = CosineAnnealingLR(opt, ...)
+
 # calling backward on the model will peform the optimzier step
 loss = model(torch.randn(20, dtype=torch.bfloat16))
 loss.backward()
@@ -95,6 +98,9 @@ loss.backward()
 # opt.step()
 # opt.zero_grad()
 
+# step the learning rate scheduler like normal
+scheduler.step()
+
 # optionally remove gradient release hooks when done training
 remove_gradient_release(model)
 ```
@@ -111,6 +117,9 @@ prepare_for_gradient_release(model, opt)
 # gradients directly into the optimizer states
 accumulation_steps = 4
 
+# setup a learning rate scheduler for gradient accumulation
+scheduler = CosineAnnealingLR(opt, ...)
+
 # use existing PyTorch dataloader
 for idx, batch in enumerate(dataloader):
     # `optimizer_accumulation=True` accumulates gradients into
@@ -128,6 +137,10 @@ for idx, batch in enumerate(dataloader):
     # opt.step()
     # opt.zero_grad()
 
+    # step the learning rate scheduler after accumulating gradients
+    if not opt.optimizer_accumulation:
+        scheduler.step()
+
 # optionally remove gradient release hooks when done training
 remove_gradient_release(model)
 ```

docs/optimizer_accumulation.md

@@ -64,6 +64,9 @@ prepare_for_gradient_release(model, opt)
 # gradients directly into the optimizer states
 accumulation_steps = 4
 
+# setup a learning rate scheduler for gradient accumulation
+scheduler = CosineAnnealingLR(opt, ...)
+
 # use existing PyTorch dataloader
 for idx, batch in enumerate(dataloader):
     # `optimizer_accumulation=True` accumulates gradients into
@@ -81,6 +84,10 @@ for idx, batch in enumerate(dataloader):
     # opt.step()
     # opt.zero_grad()
 
+    # step the learning rate scheduler after accumulating gradients
+    if not opt.optimizer_accumulation:
+        scheduler.step()
+
 # optionally remove gradient release hooks when done training
 remove_gradient_release(model)
 ```