Skip to content

Fix gradient accumulation #7

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

Open
rockerBOO wants to merge 7 commits into
base: main
Choose a base branch
from
Open

Fix gradient accumulation #7

Show file tree
Hide file tree
Changes from all commits
Commits
Show all changes
7 commits
Select commit Hold shift + click to select a range
4f08b4d
Fix gradient accumulation for IVON Optimizer
rockerBOO Jun 18, 2025
b69192a
Add test requirements
rockerBOO Jun 18, 2025
4fcfedb
Remove extra accumulation
rockerBOO Jun 18, 2025
ec231e7
Formatting
rockerBOO Jun 18, 2025
23a7e43
Update devices/dtype to be dynamic when modules are moved post initia…
rockerBOO Jun 18, 2025
81fbf8d
Fix mixed precision scenerios
rockerBOO Jun 18, 2025
3bfcc19
Ensure on same device
rockerBOO Jun 18, 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
248 changes: 181 additions & 67 deletions ivon/_ivon.py
View file
Open in desktop
Original file line number Diff line number Diff line change
Expand Up @@ -11,7 +11,17 @@


def _welford_mean(avg: Optional[Tensor], newval: Tensor, count: int) -> Tensor:
return newval if avg is None else avg + (newval - avg) / count
# Use autocast for numerical stability during arithmetic operations
with torch.autocast(device_type=newval.device.type):
if avg is None:
return newval
else:
# Ensure consistent device and dtype
if avg.device != newval.device:
avg = avg.to(device=newval.device)
if avg.dtype != newval.dtype:
newval = newval.to(dtype=avg.dtype)
return avg + (newval - avg) / count


class IVON(torch.optim.Optimizer):
Expand Down Expand Up @@ -74,7 +84,7 @@ def __init__(
self.mc_samples = mc_samples
self.hess_approx = hess_approx
self.sync = sync
self._numel, self._device, self._dtype = self._get_param_configs()
self._numel = self._get_param_configs()
self.current_step = 0
self.debias = debias
self.rescale_lr = rescale_lr
Expand All @@ -85,45 +95,55 @@ def __init__(
self._init_buffers()

def _get_param_configs(self):
all_params = []
for pg in self.param_groups:
pg["numel"] = sum(p.numel() for p in pg["params"] if p is not None)
all_params += [p for p in pg["params"] if p is not None]
if len(all_params) == 0:
return 0, torch.device("cpu"), torch.get_default_dtype()
devices = {p.device for p in all_params}
if len(devices) > 1:
raise ValueError(
"Parameters are on different devices: "
f"{[str(d) for d in devices]}"
)
device = next(iter(devices))
dtypes = {p.dtype for p in all_params}
if len(dtypes) > 1:
raise ValueError(
"Parameters are on different dtypes: "
f"{[str(d) for d in dtypes]}"
)
dtype = next(iter(dtypes))
total = sum(pg["numel"] for pg in self.param_groups)
return total, device, dtype
return total

def _reset_samples(self):
self.state['count'] = 0
self.state['avg_grad'] = None
self.state['avg_nxg'] = None
self.state['avg_gsq'] = None
@property
def _device(self):
"""Get device from first parameter"""
for pg in self.param_groups:
for p in pg["params"]:
if p is not None:
return p.device
return torch.device("cpu")

@property
def _dtype(self):
"""Get dtype from first parameter"""
for pg in self.param_groups:
for p in pg["params"]:
if p is not None:
return p.dtype
return torch.get_default_dtype()

def _reset_samples(self, reset_count=True):
# Do not reset critical state unless explicitly told to do so
if reset_count:
self.state['count'] = 0
# Intentionally keep gradient-related states as they are

def _init_buffers(self):
for group in self.param_groups:
hess_init, numel = group["hess_init"], group["numel"]

# Always use float32 for internal optimizer buffers for numerical stability
group["momentum"] = torch.zeros(
numel, device=self._device, dtype=self._dtype
numel, device=self._device, dtype=torch.float32
)
group["hess"] = torch.zeros(
numel, device=self._device, dtype=self._dtype
numel, device=self._device, dtype=torch.float32
).add(torch.as_tensor(hess_init))

def _ensure_buffers_on_device(self):
"""Ensure buffers are on the same device as parameters"""
for group in self.param_groups:
# Always use float32 for internal optimizer buffers
if "momentum" in group and group["momentum"].device != self._device:
group["momentum"] = group["momentum"].to(device=self._device, dtype=torch.float32)
if "hess" in group and group["hess"].device != self._device:
group["hess"] = group["hess"].to(device=self._device, dtype=torch.float32)

@contextmanager
def sampled_params(self, train: bool = False):
Expand All @@ -143,19 +163,53 @@ def _restore_param_average(

p_slice = slice(offset, offset + p.numel())

p.data = param_avg[p_slice].view(p.shape)
# Convert param_avg back to parameter's original dtype
p.data = param_avg[p_slice].view(p.shape).to(dtype=p.dtype)
if train:
if p.requires_grad:
# Always collect gradient, handling potential accumulated gradients
if p.grad is not None:
param_grads.append(p.grad.flatten())
else:
param_grads.append(torch.zeros_like(p).flatten())
param_grads.append(torch.zeros(
p.numel(),
device=p.device,
dtype=p.dtype
))
offset += p.numel()
assert offset == self._numel # sanity check

if train: # collect grad sample for training
if train and param_grads: # collect grad sample for training
grad_sample = torch.cat(param_grads, 0)

# Ensure state is initialized
if 'count' not in self.state:
self.state['count'] = 0

# Increment count
count = self.state["count"] + 1
self.state["count"] = count

# Compute total number of parameters
total_numel = sum(group['numel'] for group in self.param_groups)

# Initialize gradient-related states if not already done
if self.state.get('avg_grad') is None:
# Use gradient dtype for state initialization
self.state['avg_grad'] = torch.zeros(
total_numel,
device=self._device,
dtype=grad_sample.dtype
)
self.state['avg_nxg'] = torch.zeros_like(self.state['avg_grad'])
self.state['avg_gsq'] = torch.zeros_like(self.state['avg_grad'])
else:
# Ensure existing state tensors are on the correct device
if self.state['avg_grad'].device != grad_sample.device:
self.state['avg_grad'] = self.state['avg_grad'].to(device=grad_sample.device)
self.state['avg_nxg'] = self.state['avg_nxg'].to(device=grad_sample.device)
self.state['avg_gsq'] = self.state['avg_gsq'].to(device=grad_sample.device)

# Update running averages
self.state["avg_grad"] = _welford_mean(
self.state["avg_grad"], grad_sample, count
)
Expand All @@ -165,22 +219,58 @@ def _restore_param_average(
elif self.hess_approx == 'gradsq':
self.state['avg_gsq'] = _welford_mean(
self.state['avg_gsq'], grad_sample.square(), count)

# Ensure current_step is incremented
self.current_step = count

@torch.no_grad()
def step(self, closure: ClosureType = None) -> Optional[Tensor]:
# Ensure buffers are on correct device
self._ensure_buffers_on_device()

# Compute total number of parameters
total_numel = sum(group['numel'] for group in self.param_groups)

# Ensure state is properly initialized
if 'count' not in self.state or self.state['count'] is None:
self.state['count'] = 0

# Initialize gradient-related states if not already done
if self.state.get('avg_grad') is None:
# Initialize with parameter dtype
self.state['avg_grad'] = torch.zeros(
total_numel,
device=self._device,
dtype=self._dtype
)
self.state['avg_nxg'] = torch.zeros_like(self.state['avg_grad'])
self.state['avg_gsq'] = torch.zeros_like(self.state['avg_grad'])

# Set closure default if not provided
if closure is None:
loss = None
else:
# Handle multiple Monte Carlo samples if specified
losses = []
for _ in range(self.mc_samples):
with torch.enable_grad():
loss = closure()
losses.append(loss)
loss = sum(losses) / self.mc_samples
if self.sync and dist.is_initialized(): # explicit sync

# Handle distributed sync if needed
if self.sync and dist.is_initialized():
self._sync_samples()

# Track current step
self.current_step += 1

# Perform parameter update
self._update()
self._reset_samples()

# Reset running statistics, but keep count and state
self._reset_samples(reset_count=False)

return loss

def _sync_samples(self):
Expand All @@ -191,17 +281,21 @@ def _sync_samples(self):
self.state["avg_nxg"].div_(world_size)

def _sample_params(self) -> Tuple[Tensor, Tensor]:
# Ensure buffers are on correct device
self._ensure_buffers_on_device()

noise_samples = []
param_avgs = []

offset = 0
for group in self.param_groups:
gnumel = group["numel"]

# Generate noise on the same device as the hess buffer
hess_buffer = group["hess"] + group["weight_decay"]
noise_sample = (
torch.randn(gnumel, device=self._device, dtype=self._dtype)
/ (
group["ess"] * (group["hess"] + group["weight_decay"])
).sqrt()
torch.randn(gnumel, device=hess_buffer.device, dtype=self._dtype)
/ (group["ess"] * hess_buffer).sqrt()
)
noise_samples.append(noise_sample)

Expand All @@ -215,7 +309,10 @@ def _sample_params(self) -> Tuple[Tensor, Tensor]:
p_noise = noise_sample[goffset : goffset + numel]

param_avgs.append(p_avg)
p.data = (p_avg + p_noise).view(p.shape)
# Use autocast for parameter sampling arithmetic
with torch.autocast(device_type=p.device.type):
sampled_param = (p_avg + p_noise).to(dtype=p.dtype)
p.data = sampled_param.view(p.shape)
goffset += numel
offset += numel
assert goffset == group["numel"] # sanity check
Expand All @@ -224,8 +321,20 @@ def _sample_params(self) -> Tuple[Tensor, Tensor]:
return torch.cat(param_avgs, 0), torch.cat(noise_samples, 0)

def _update(self):
self.current_step += 1
# If avg_grad is None, initialize with zeros
if self.state.get('avg_grad') is None:
# Compute total number of parameters
total_numel = sum(group['numel'] for group in self.param_groups)
self.state['avg_grad'] = torch.zeros(
total_numel,
device=self._device,
dtype=self._dtype
)
self.state['avg_nxg'] = torch.zeros_like(self.state['avg_grad'])
self.state['avg_gsq'] = torch.zeros_like(self.state['avg_grad'])
self.state['count'] = 0

# Prepare for parameter update
offset = 0
for group in self.param_groups:
lr = group["lr"]
Expand All @@ -237,39 +346,44 @@ def _update(self):
[p.flatten() for p in group["params"] if p is not None], 0
)

group["momentum"] = self._new_momentum(
self.state["avg_grad"][pg_slice], group["momentum"], b1
)

group["hess"] = self._new_hess(
self.hess_approx,
group["hess"],
self.state["avg_nxg"],
self.state['avg_gsq'],
pg_slice,
group["ess"],
b2,
group["weight_decay"],
)

param_avg = self._new_param_averages(
param_avg,
group["hess"],
group["momentum"],
lr * (group["hess_init"] + group["weight_decay"]) if self.rescale_lr else lr,
group["weight_decay"],
group["clip_radius"],
1.0 - pow(b1, float(self.current_step)) if self.debias else 1.0,
group["hess_init"]
)
# Handle case where avg_grad might not have been collected
if self.state['avg_grad'] is not None:
# Update momentum with available gradient information
group["momentum"] = self._new_momentum(
self.state["avg_grad"][pg_slice], group["momentum"], b1
)

group["hess"] = self._new_hess(
self.hess_approx,
group["hess"],
self.state["avg_nxg"],
self.state['avg_gsq'],
pg_slice,
group["ess"],
b2,
group["weight_decay"],
)

param_avg = self._new_param_averages(
param_avg,
group["hess"],
group["momentum"],
lr * (group["hess_init"] + group["weight_decay"]) if self.rescale_lr else lr,
group["weight_decay"],
group["clip_radius"],
1.0 - pow(b1, float(self.current_step)) if self.debias else 1.0,
group["hess_init"]
)

# update params
pg_offset = 0
for p in group["params"]:
if p is not None:
p.data = param_avg[pg_offset : pg_offset + p.numel()].view(
# Convert back to parameter's original dtype
new_param_data = param_avg[pg_offset : pg_offset + p.numel()].view(
p.shape
)
).to(dtype=p.dtype)
p.data = new_param_data
pg_offset += p.numel()
assert pg_offset == group["numel"] # sanity check
offset += group["numel"]
Expand Down
3 changes: 3 additions & 0 deletions requirements-test.txt
View file
Open in desktop
Original file line number Diff line number Diff line change
@@ -0,0 +1,3 @@
pytest >=8.1.1
numpy < 2.0
accelerate
Loading