chore: import upstream snapshot with attribution
This commit is contained in:
wehub-resource-sync
@@ -0,0 +1,511 @@
|
||||
# Copyright (c) Microsoft Corporation.
|
||||
# SPDX-License-Identifier: Apache-2.0
|
||||
|
||||
# DeepSpeed Team
|
||||
|
||||
import os
|
||||
import torch
|
||||
from typing import Any
|
||||
|
||||
from deepspeed.utils import logger
|
||||
from deepspeed.utils.tensor_fragment import map_to_flat_opt_states
|
||||
from deepspeed.runtime.utils import bwc_tensor_model_parallel_rank, see_memory_usage
|
||||
from deepspeed.runtime.torch_autocast import get_comm_dtype, is_autocast_initialized
|
||||
from deepspeed.runtime.utils import maybe_loss_for_backward
|
||||
|
||||
|
||||
class DeepSpeedOptimizer(object):
|
||||
pass
|
||||
|
||||
|
||||
def _get_universal_checkpoint_ep_info() -> tuple[int, int]:
|
||||
# Universal checkpoints use EP slicing only when an expert group exists.
|
||||
try:
|
||||
from deepspeed.utils import groups
|
||||
expert_groups = groups._get_expert_parallel_group_dict()
|
||||
if not expert_groups:
|
||||
return 0, 1
|
||||
max_ep_name = groups._get_max_expert_size_name()
|
||||
if max_ep_name not in expert_groups:
|
||||
return 0, 1
|
||||
return groups._get_expert_parallel_rank(max_ep_name), groups._get_expert_parallel_world_size(max_ep_name)
|
||||
except (RuntimeError, AttributeError, KeyError):
|
||||
return 0, 1
|
||||
|
||||
|
||||
class BackwardHookStateManager:
|
||||
"""Manages backward pass state for ZeRO optimizers.
|
||||
|
||||
This class handles the complex state management needed for gradient accumulation hooks
|
||||
to work correctly with:
|
||||
|
||||
1. **Reentrant Gradient Checkpointing** (use_reentrant=True):
|
||||
With reentrant checkpointing, gradient hooks fire in multiple phases within a
|
||||
single backward() call. For example, with model: linear1 (checkpointed) -> linear2:
|
||||
- Phase 1: Hooks for linear2 fire (non-checkpointed params)
|
||||
- Checkpoint recomputes linear1's forward
|
||||
- Phase 2: Hooks for linear1 fire (checkpointed params)
|
||||
|
||||
The challenge is that `count_used_parameters_in_backward()` only sees params
|
||||
currently in the backward graph. During Phase 1, it returns 2 (linear2's params),
|
||||
but after checkpoint recomputation, it returns 4 (all params). We must NOT run
|
||||
the epilogue prematurely after Phase 1.
|
||||
|
||||
Solution: Queue a post-backward callback on the autograd engine at the start of
|
||||
backward and run the epilogue when the graph task completes. This avoids premature
|
||||
epilogues across reentrant phases. The `_max_expected_hooks_seen` counter remains
|
||||
as a fallback when the callback API is unavailable.
|
||||
|
||||
2. **TiledFusedLogitsLoss and Similar Custom Autograd Functions**:
|
||||
Some custom autograd functions call `torch.autograd.backward()` from their
|
||||
forward pass BEFORE the user calls `engine.backward(loss)`. These internal
|
||||
backward calls trigger ZeRO's gradient hooks, but we must NOT run the epilogue
|
||||
until the user's actual backward pass.
|
||||
|
||||
Solution: Track `_backward_active_depth` which is only incremented when
|
||||
`enter_backward()` is called (from engine.backward or user code). Hooks check
|
||||
this depth before running the epilogue.
|
||||
|
||||
3. **Multiple Backward Phases with Exit/Re-entry**:
|
||||
When the epilogue runs after Phase 1 (with reentrant checkpointing), it calls
|
||||
`exit_backward()`, setting `_backward_active_depth` to 0. When Phase 2's hooks
|
||||
fire, we need to re-enter the backward context.
|
||||
|
||||
Solution: `_backward_seen_this_step` flag tracks if backward was ever active
|
||||
this step. Combined with `_backward_active_depth == 0`, this detects Phase 2
|
||||
and calls `enter_backward()` again.
|
||||
|
||||
Attributes:
|
||||
remaining_grad_acc_hooks: Count of hooks remaining before epilogue should run
|
||||
backward_active_depth: Nesting depth of backward() calls (0 = not in backward)
|
||||
backward_seen_this_step: True if enter_backward() was called this step
|
||||
epilogue_ran_this_backward: True if epilogue ran (for micro_step_id management)
|
||||
hooks_fired_this_backward: Count of gradient hooks that have fired
|
||||
max_expected_hooks_seen: Maximum expected hook count seen (grows with reentrant)
|
||||
post_backward_callback_queued: True if a post-backward callback is queued
|
||||
post_backward_callback_graph_task_id: Graph task id for the queued callback
|
||||
"""
|
||||
|
||||
def __init__(self):
|
||||
self.remaining_grad_acc_hooks = 0
|
||||
self._grad_acc_post_hooks = []
|
||||
self.backward_active_depth = 0
|
||||
self.backward_seen_this_step = False
|
||||
self.epilogue_ran_this_backward = False
|
||||
self.hooks_fired_this_backward = 0
|
||||
self.max_expected_hooks_seen = 0
|
||||
self.post_backward_callback_queued = False
|
||||
self.post_backward_callback_graph_task_id = None
|
||||
|
||||
def register_grad_acc_post_hook(self, hook):
|
||||
"""Register a callback to run when all gradient hooks have fired."""
|
||||
self._grad_acc_post_hooks.append(hook)
|
||||
|
||||
def unregister_grad_acc_post_hooks(self):
|
||||
"""Remove all registered gradient accumulation post hooks."""
|
||||
self._grad_acc_post_hooks = []
|
||||
|
||||
def run_grad_acc_post_hooks(self):
|
||||
"""Run all registered post hooks if backward is active.
|
||||
|
||||
Custom autograd Functions (e.g., TiledFusedLogitsLoss) can invoke
|
||||
`torch.autograd.backward()` from their *forward* pass before the user
|
||||
ever calls `engine.backward(loss)`. Those early backward calls still
|
||||
trigger ZeRO's grad hooks, but we must not run the engine's
|
||||
post-backward logic (which reduces/clears grads) until the outer/user
|
||||
backward is active. The depth guard filters out only those pre-user
|
||||
invocations while still allowing backward calls that happen during
|
||||
the real user backward.
|
||||
"""
|
||||
if self.backward_active_depth == 0:
|
||||
return
|
||||
for hook in self._grad_acc_post_hooks:
|
||||
hook()
|
||||
|
||||
def enter_backward(self):
|
||||
"""Enter backward context. Call at the start of backward pass."""
|
||||
# On first real backward entry of a step, reset counters that may have been
|
||||
# polluted by pre-user-backward hooks (e.g. TiledFusedLogitsLoss calling
|
||||
# torch.autograd.backward() from forward). Do NOT reset on reentrant
|
||||
# phase re-entry (backward_seen_this_step == True) so phase-to-phase
|
||||
# state remains intact.
|
||||
if self.backward_active_depth == 0 and not self.backward_seen_this_step:
|
||||
self.hooks_fired_this_backward = 0
|
||||
self.max_expected_hooks_seen = 0
|
||||
self.remaining_grad_acc_hooks = 0
|
||||
self.post_backward_callback_queued = False
|
||||
self.post_backward_callback_graph_task_id = None
|
||||
self.backward_active_depth += 1
|
||||
# Track that backward has been active at some point in this step.
|
||||
# This is used to detect subsequent gradient hook phases with reentrant checkpointing.
|
||||
self.backward_seen_this_step = True
|
||||
|
||||
def exit_backward(self):
|
||||
"""Exit backward context. Call at the end of backward pass."""
|
||||
if self.backward_active_depth > 0:
|
||||
self.backward_active_depth -= 1
|
||||
|
||||
def reset_for_new_step(self):
|
||||
"""Reset state at the start of each forward/backward step."""
|
||||
self.backward_seen_this_step = False
|
||||
self.hooks_fired_this_backward = 0
|
||||
self.max_expected_hooks_seen = 0
|
||||
self.epilogue_ran_this_backward = False
|
||||
self.post_backward_callback_queued = False
|
||||
self.post_backward_callback_graph_task_id = None
|
||||
|
||||
def should_refresh_expected_hook_count(self):
|
||||
"""Return True when count_used_parameters_in_backward() should be re-evaluated.
|
||||
|
||||
Refresh is needed in two cases:
|
||||
1. First hook of a backward (or backward phase): hooks_fired == 0.
|
||||
2. A new reentrant phase started: remaining hooks exhausted, we exited
|
||||
backward, but backward was active earlier this step.
|
||||
|
||||
The predicate must be evaluated BEFORE reenter_backward_if_needed()
|
||||
because re-entering changes backward_active_depth and hides the
|
||||
phase-boundary signal.
|
||||
"""
|
||||
return (self.hooks_fired_this_backward == 0
|
||||
or (self.remaining_grad_acc_hooks == 0 and self.backward_active_depth == 0
|
||||
and self.backward_seen_this_step))
|
||||
|
||||
def reenter_backward_if_needed(self):
|
||||
"""Re-enter backward context for subsequent phases in reentrant checkpointing.
|
||||
|
||||
With reentrant gradient checkpointing, gradient hooks can fire in multiple phases
|
||||
within a single backward call. When the epilogue runs after a phase, it calls
|
||||
exit_backward(), setting backward_active_depth to 0. When the next phase starts,
|
||||
we need to re-enter backward.
|
||||
|
||||
We detect subsequent phases by checking:
|
||||
1. remaining_grad_acc_hooks == 0 (epilogue ran or new backward)
|
||||
2. backward_active_depth == 0 (we've exited from previous phase)
|
||||
3. backward_seen_this_step == True (backward was active earlier)
|
||||
|
||||
This distinguishes from TiledFusedLogitsLoss which calls backward() during forward -
|
||||
in that case backward_seen_this_step is False because enter_backward() was never called.
|
||||
"""
|
||||
if self.remaining_grad_acc_hooks == 0:
|
||||
if self.backward_active_depth == 0 and self.backward_seen_this_step:
|
||||
self.enter_backward()
|
||||
|
||||
def queue_post_backward_callback(self):
|
||||
"""Queue post-backward hooks to run after the current graph finishes."""
|
||||
if self.post_backward_callback_queued:
|
||||
return True
|
||||
if self.backward_active_depth == 0:
|
||||
return False
|
||||
|
||||
engine = getattr(torch.autograd.Variable, "_execution_engine", None)
|
||||
if engine is None or not hasattr(engine, "queue_callback"):
|
||||
return False
|
||||
if not hasattr(torch._C, "_current_graph_task_id"):
|
||||
return False
|
||||
|
||||
graph_task_id = torch._C._current_graph_task_id()
|
||||
if graph_task_id == -1:
|
||||
return False
|
||||
|
||||
def _run_post_backward():
|
||||
self.run_grad_acc_post_hooks()
|
||||
|
||||
engine.queue_callback(_run_post_backward)
|
||||
self.post_backward_callback_queued = True
|
||||
self.post_backward_callback_graph_task_id = graph_task_id
|
||||
return True
|
||||
|
||||
def update_hook_state_and_maybe_run_epilogue(self, current_expected_count):
|
||||
"""Update hook state after a gradient hook fires and run epilogue if all hooks have fired.
|
||||
|
||||
With reentrant gradient checkpointing, count_used_parameters_in_backward() returns the
|
||||
count of params that will execute in the current backward graph. This count grows as
|
||||
checkpointed regions are recomputed. We track the MAXIMUM count seen to ensure we don't
|
||||
run the epilogue until all params that will ever participate have been processed.
|
||||
Counters are reset at forward() time via reset_for_new_step().
|
||||
|
||||
Args:
|
||||
current_expected_count: The current expected number of hooks, from
|
||||
count_used_parameters_in_backward() plus any leaf modules.
|
||||
"""
|
||||
self.hooks_fired_this_backward += 1
|
||||
self.max_expected_hooks_seen = max(self.max_expected_hooks_seen, current_expected_count)
|
||||
|
||||
# Prefer running post-backward hooks via autograd engine callback when available.
|
||||
# This avoids premature epilogues with reentrant checkpointing.
|
||||
if self.queue_post_backward_callback():
|
||||
self.remaining_grad_acc_hooks = max(self.max_expected_hooks_seen - self.hooks_fired_this_backward, 0)
|
||||
return
|
||||
|
||||
# Fallback: Run epilogue only when we've processed ALL params that will participate.
|
||||
# This is the maximum count we've seen (accounts for late-joining params
|
||||
# from reentrant checkpointing) and also excludes unused params.
|
||||
if self.hooks_fired_this_backward >= self.max_expected_hooks_seen:
|
||||
self.remaining_grad_acc_hooks = 0
|
||||
self.run_grad_acc_post_hooks()
|
||||
else:
|
||||
self.remaining_grad_acc_hooks = self.max_expected_hooks_seen - self.hooks_fired_this_backward
|
||||
|
||||
|
||||
class ZeROOptimizer(DeepSpeedOptimizer):
|
||||
"""Base class for ZeRO optimizer implementations (stages 1, 2, and 3)."""
|
||||
|
||||
def __init__(self):
|
||||
self._backward_hook_state = BackwardHookStateManager()
|
||||
|
||||
# Delegate backward hook state management to the manager.
|
||||
# These properties provide backward compatibility with code that accesses
|
||||
# these attributes directly (e.g., in stage3.py and stage_1_and_2.py).
|
||||
@property
|
||||
def _remaining_grad_acc_hooks(self):
|
||||
return self._backward_hook_state.remaining_grad_acc_hooks
|
||||
|
||||
@_remaining_grad_acc_hooks.setter
|
||||
def _remaining_grad_acc_hooks(self, value):
|
||||
self._backward_hook_state.remaining_grad_acc_hooks = value
|
||||
|
||||
@property
|
||||
def _backward_active_depth(self):
|
||||
return self._backward_hook_state.backward_active_depth
|
||||
|
||||
@_backward_active_depth.setter
|
||||
def _backward_active_depth(self, value):
|
||||
self._backward_hook_state.backward_active_depth = value
|
||||
|
||||
@property
|
||||
def _backward_seen_this_step(self):
|
||||
return self._backward_hook_state.backward_seen_this_step
|
||||
|
||||
@_backward_seen_this_step.setter
|
||||
def _backward_seen_this_step(self, value):
|
||||
self._backward_hook_state.backward_seen_this_step = value
|
||||
|
||||
@property
|
||||
def _epilogue_ran_this_backward(self):
|
||||
return self._backward_hook_state.epilogue_ran_this_backward
|
||||
|
||||
@_epilogue_ran_this_backward.setter
|
||||
def _epilogue_ran_this_backward(self, value):
|
||||
self._backward_hook_state.epilogue_ran_this_backward = value
|
||||
|
||||
@property
|
||||
def _hooks_fired_this_backward(self):
|
||||
return self._backward_hook_state.hooks_fired_this_backward
|
||||
|
||||
@_hooks_fired_this_backward.setter
|
||||
def _hooks_fired_this_backward(self, value):
|
||||
self._backward_hook_state.hooks_fired_this_backward = value
|
||||
|
||||
@property
|
||||
def _max_expected_hooks_seen(self):
|
||||
return self._backward_hook_state.max_expected_hooks_seen
|
||||
|
||||
@_max_expected_hooks_seen.setter
|
||||
def _max_expected_hooks_seen(self, value):
|
||||
self._backward_hook_state.max_expected_hooks_seen = value
|
||||
|
||||
@property
|
||||
def _grad_acc_post_hooks(self):
|
||||
return self._backward_hook_state._grad_acc_post_hooks
|
||||
|
||||
@_grad_acc_post_hooks.setter
|
||||
def _grad_acc_post_hooks(self, value):
|
||||
self._backward_hook_state._grad_acc_post_hooks = value
|
||||
|
||||
def load_hp_checkpoint_state_from_checkpoint_dir(self, lp_groups_name: str, checkpoint_dir: str) -> None:
|
||||
checkpoint_dir = os.path.join(checkpoint_dir, "zero")
|
||||
optim_state_path = os.path.join(checkpoint_dir, "optimizer_state.pt")
|
||||
assert os.path.isfile(
|
||||
optim_state_path), f'{optim_state_path} containing optimizer global state is missing! Cannot proceed.'
|
||||
optim_sd = torch.load(optim_state_path, weights_only=False)
|
||||
|
||||
self._load_global_state(optim_sd)
|
||||
|
||||
tp_rank = bwc_tensor_model_parallel_rank(mpu=self.mpu)
|
||||
if self.mpu is None:
|
||||
logger.warning("MPU is not provided, setting tp size to 1 in checkpoint loading.")
|
||||
tp_world_size = 1
|
||||
else:
|
||||
tp_world_size = self.mpu.get_slice_parallel_world_size() if hasattr(self.mpu, "get_slice_parallel_world_size") \
|
||||
else self.mpu.get_tensor_model_parallel_world_size()
|
||||
|
||||
ep_rank, ep_size = _get_universal_checkpoint_ep_info()
|
||||
|
||||
for i, (param_group,
|
||||
loaded_param_group) in enumerate(zip(self.optimizer.param_groups, optim_sd['param_groups'])):
|
||||
# We have an assumption that all params in the same param_group have the same keys
|
||||
opt_keys = set()
|
||||
steps = []
|
||||
|
||||
lp_groups = getattr(self, lp_groups_name)
|
||||
for lp in lp_groups[i]:
|
||||
if lp._hp_mapping is not None:
|
||||
#print(f"Loading {self.param_names[lp]} {tp_rank=} {tp_world_size=}")
|
||||
step = lp.load_hp_checkpoint_state(os.path.join(checkpoint_dir, self.param_names[lp]),
|
||||
tp_rank,
|
||||
tp_world_size,
|
||||
ep_rank=ep_rank,
|
||||
ep_size=ep_size)
|
||||
for key in lp._hp_mapping.get_optim_state_keys():
|
||||
opt_keys.add(key)
|
||||
steps.append(step)
|
||||
|
||||
hp_param = param_group['params'][0]
|
||||
assert all(step == steps[0] for step in steps), f"Steps {steps} are not equal"
|
||||
if steps[0] is not None:
|
||||
self.optimizer.state[hp_param]['step'] = steps[0]
|
||||
|
||||
map_to_flat_opt_states(hp_param, lp_groups[i], self.optimizer.state, opt_keys)
|
||||
|
||||
for key, value in loaded_param_group.items():
|
||||
if key == 'params':
|
||||
continue
|
||||
param_group[key] = value
|
||||
|
||||
def report_ipg_memory_usage(self, tag, param_elems, dtype=None):
|
||||
dtypes = self.ipg_buckets.keys() if dtype is None else [dtype]
|
||||
|
||||
for dt in dtypes:
|
||||
bucket = self.ipg_buckets[dt]
|
||||
elem_count = bucket.elements + param_elems
|
||||
percent_of_bucket_size = (100.0 * elem_count) // self.reduce_bucket_size
|
||||
see_memory_usage(
|
||||
f"{tag}: elems in_bucket {dt} {bucket.elements} param {param_elems} max_percent {percent_of_bucket_size}"
|
||||
)
|
||||
|
||||
def get_param_comm_dtype(self, param):
|
||||
if is_autocast_initialized():
|
||||
return get_comm_dtype(param)
|
||||
else:
|
||||
return self.communication_data_type
|
||||
|
||||
def needs_scaler(self) -> bool:
|
||||
"""
|
||||
Check if this optimizer requires loss scaling for correct backward pass.
|
||||
|
||||
Returns True if any of the following conditions are met:
|
||||
- Custom loss scaler is enabled
|
||||
- torch.autocast gradient scaler is active (fp16 only)
|
||||
- Dynamic loss scaling is enabled (fp16 with DeepSpeed's loss scaler)
|
||||
|
||||
Returns False for bf16 or fp32, which don't require gradient scaling.
|
||||
"""
|
||||
return (self.custom_loss_scaler or self.torch_autocast_gradscaler is not None
|
||||
or (hasattr(self, 'dynamic_loss_scale') and self.dynamic_loss_scale))
|
||||
|
||||
def scale_if_loss(self, value: Any) -> Any:
|
||||
"""
|
||||
Applies loss scaling to the input value if it is a loss tensor.
|
||||
"""
|
||||
if maybe_loss_for_backward(value):
|
||||
if self.custom_loss_scaler:
|
||||
return self.external_loss_scale * value
|
||||
if self.torch_autocast_gradscaler:
|
||||
return self.torch_autocast_gradscaler.scale(value)
|
||||
# Only call loss_scaler if it exists (not present in BF16_Optimizer)
|
||||
if hasattr(self, 'loss_scaler') and self.loss_scaler is not None:
|
||||
return self.loss_scaler.scale_loss(value)
|
||||
|
||||
return value
|
||||
|
||||
def backward_prologue(self):
|
||||
pass
|
||||
|
||||
def backward_epilogue(self, **kwargs):
|
||||
pass
|
||||
|
||||
def backward(self, loss, **kwargs):
|
||||
assert maybe_loss_for_backward(loss), "Optimizer's backward() only accepts a scalar tensor"
|
||||
|
||||
scaled_loss = self.backward_prologue(loss)
|
||||
retain_graph = kwargs.pop('retain_graph', False)
|
||||
self.enter_backward()
|
||||
scaled_loss.backward(retain_graph=retain_graph)
|
||||
self.backward_epilogue()
|
||||
self.exit_backward()
|
||||
|
||||
def register_grad_acc_post_hook(self, hook):
|
||||
"""Register a callback to run when all gradient hooks have fired."""
|
||||
self._backward_hook_state.register_grad_acc_post_hook(hook)
|
||||
|
||||
def unregister_grad_acc_post_hooks(self):
|
||||
"""Remove all registered gradient accumulation post hooks."""
|
||||
self._backward_hook_state.unregister_grad_acc_post_hooks()
|
||||
|
||||
def run_grad_acc_post_hooks(self):
|
||||
"""Run all registered post hooks if backward is active."""
|
||||
self._backward_hook_state.run_grad_acc_post_hooks()
|
||||
|
||||
def enter_backward(self):
|
||||
"""Enter backward context. Call at the start of backward pass."""
|
||||
self._backward_hook_state.enter_backward()
|
||||
|
||||
def exit_backward(self):
|
||||
"""Exit backward context. Call at the end of backward pass."""
|
||||
self._backward_hook_state.exit_backward()
|
||||
|
||||
def clear_backward_seen_flag(self):
|
||||
"""Clear the backward seen flag and reset hook counters at the start of each step."""
|
||||
self._backward_hook_state.reset_for_new_step()
|
||||
|
||||
def should_refresh_expected_hook_count(self):
|
||||
"""Return True when count_used_parameters_in_backward() should be re-evaluated."""
|
||||
return self._backward_hook_state.should_refresh_expected_hook_count()
|
||||
|
||||
def reenter_backward_if_needed(self):
|
||||
"""Re-enter backward context for subsequent phases in reentrant checkpointing."""
|
||||
self._backward_hook_state.reenter_backward_if_needed()
|
||||
|
||||
def update_hook_state_and_maybe_run_epilogue(self, current_expected_count):
|
||||
"""Update hook state after a gradient hook fires and run epilogue if all hooks have fired."""
|
||||
self._backward_hook_state.update_hook_state_and_maybe_run_epilogue(current_expected_count)
|
||||
|
||||
def queue_post_backward_callback(self):
|
||||
"""Queue post-backward hooks to run after autograd completes."""
|
||||
return self._backward_hook_state.queue_post_backward_callback()
|
||||
|
||||
def _configure_master_weights(self,
|
||||
fp16_master_weights_and_gradients=False,
|
||||
bf16_master_weights_and_gradients=False,
|
||||
bf16_optimizer_states=False,
|
||||
offload_enabled=False,
|
||||
fp16_offload_validator=None,
|
||||
bf16_offload_validator=None):
|
||||
"""
|
||||
Common validation and dtype selection for ZeRO optimizer master-weight settings.
|
||||
Optionally accepts callables that enforce backend-specific offload requirements.
|
||||
``offload_enabled`` tells this method whether optimizer-state offload is configured,
|
||||
so the offload requirement is also enforced for the bf16-optimizer-states + offload case.
|
||||
"""
|
||||
self.fp16_master_weights_and_gradients = fp16_master_weights_and_gradients
|
||||
self.bf16_master_weights_and_gradients = bf16_master_weights_and_gradients
|
||||
assert not (self.fp16_master_weights_and_gradients and self.bf16_master_weights_and_gradients), \
|
||||
"fp16_master_weights_and_gradients and bf16_master_weights_and_gradients are mutually exclusive."
|
||||
|
||||
self.bf16_optimizer_states = bf16_optimizer_states
|
||||
if self.bf16_optimizer_states:
|
||||
assert self.bf16_master_weights_and_gradients, \
|
||||
"bf16_optimizer_states requires bf16_master_weights_and_gradients."
|
||||
|
||||
# bf16 master weights require ZeRO-Offload + DeepSpeedCPUAdam whenever the optimizer states
|
||||
# cannot stay on the GPU: either because they remain fp32 (bf16_optimizer_states disabled),
|
||||
# or because CPU offload is explicitly requested alongside bf16 optimizer states.
|
||||
if (self.bf16_master_weights_and_gradients and bf16_offload_validator is not None
|
||||
and (not self.bf16_optimizer_states or offload_enabled)):
|
||||
bf16_offload_validator()
|
||||
# Offloaded bf16 optimizer states need the CPU optimizer to store moments in the
|
||||
# parameter (bf16) precision; otherwise they would silently expand back to fp32.
|
||||
if self.bf16_optimizer_states:
|
||||
assert not getattr(self.optimizer, 'fp32_optimizer_states', True), \
|
||||
"bf16_optimizer_states with ZeRO-Offload requires DeepSpeedCPUAdam constructed " \
|
||||
"with fp32_optimizer_states=False so optimizer moments are stored in bf16."
|
||||
|
||||
if self.fp16_master_weights_and_gradients and fp16_offload_validator is not None:
|
||||
fp16_offload_validator()
|
||||
|
||||
if self.fp16_master_weights_and_gradients:
|
||||
return torch.float16
|
||||
elif self.bf16_master_weights_and_gradients:
|
||||
return torch.bfloat16
|
||||
else:
|
||||
return torch.float32
|
||||
Reference in New Issue
Block a user