Skip to content

Fix #7733: Replace torch.sqrt with math.sqrt in scale_lr for sqrt method #7735

New issue

Have a question about this project? Sign up for a free GitHub account to open an issue and contact its maintainers and the community.

Sign up for GitHub

By clicking “Sign up for GitHub”, you agree to our terms of service and privacy statement. We’ll occasionally send you account related emails.

Already on GitHub? Sign in to your account

Merged
sfc-gh-truwase merged 5 commits into from
Dec 19, 2025
+2 −2
Merged

Fix #7733: Replace torch.sqrt with math.sqrt in scale_lr for sqrt method #7735

Changes from all commits
Commits
Show all changes
5 commits
Select commit Hold shift + click to select a range
895ac3f
Fix #7733: Replace torch.sqrt with math.sqrt in scale_lr for sqrt method
Rakshit-gen Dec 18, 2025
7286dd9
removed unused imports
Rakshit-gen Dec 19, 2025
eea3ef8
Fix rare hang in DeepSpeed Async I/O wait by releasing the Python GIL…
xylian86 Dec 18, 2025
994da47
revert unneccesary change
Rakshit-gen Dec 19, 2025
c18e5a4
Merge branch 'master' into fix-scale-lr-sqrt-bug
sfc-gh-truwase Dec 19, 2025
File filter

Filter by extension

Filter by extension
Conversations
Failed to load comments.
Loading
Jump to
Jump to file
Failed to load files.
Loading
Diff view
Diff view
4 changes: 2 additions & 2 deletions deepspeed/runtime/data_pipeline/data_sampling/variable_batch_size_and_lr.py
View file
Open in desktop
Original file line number Diff line number Diff line change
Expand Up @@ -6,8 +6,8 @@
# support/questions/maintenance: github user @brunomaga or @deepspeedai/deepspeed

import random
import torch
import os
import math
import numpy as np
from torch.optim.lr_scheduler import LRScheduler
from torch.optim.optimizer import Optimizer
Expand Down Expand Up @@ -156,7 +156,7 @@ def scale_lr(base_batch_size, batch_size, base_lr=1, method="linear"):
# Square Root scaling: "when multiplying the batch size by k, multiply the learning rate
# by √k, to keep the variance in the gradient expectation constant"
# (A. Krizhevsky. One weird trick for parallelizing convolutional neural networks)
return base_lr * torch.sqrt(batch_size / base_batch_size)
return base_lr * math.sqrt(batch_size / base_batch_size)
elif method == None or method.upper() == "NONE":
return base_lr
raise ValueError("Unknown scaling method: {}".format(method))
Expand Down