wehub-resource-sync
406 lines
20 KiB
Python
406 lines
20 KiB
Python
# Copyright (c) Microsoft Corporation.
|
|
# SPDX-License-Identifier: Apache-2.0
|
|
|
|
# DeepSpeed Team
|
|
|
|
import pytest
|
|
import torch
|
|
import deepspeed
|
|
import deepspeed.comm as dist
|
|
from deepspeed.accelerator import get_accelerator
|
|
|
|
from unit.common import DistributedTest
|
|
from unit.simple_model import random_dataloader, SimpleModel
|
|
from deepspeed.runtime.zero.offload_config import OffloadDeviceEnum, OffloadStateTypeEnum
|
|
from deepspeed.utils import safe_get_local_fp32_param, safe_get_local_optimizer_state
|
|
from deepspeed.runtime.zero.offload_states import get_state_devices
|
|
|
|
# ==============================================================================
|
|
# ZeRO-1 and ZeRO-2 TESTS
|
|
# ==============================================================================
|
|
|
|
|
|
def validate_hp_params_device(model, device: torch.device):
|
|
"""Validates that the sharded FP32 parameters are on the specified device."""
|
|
for p in model.optimizer.single_partition_of_fp32_groups:
|
|
assert p.device.type == device.type, f"FP32 param partition is on {p.device}, expected {device}"
|
|
|
|
|
|
def validate_lp_params_device(model, device: torch.device):
|
|
"""Validates that the sharded LP parameters are on the specified device."""
|
|
for p in model.parameters():
|
|
assert p.device.type == device.type, f"LP param partition is on {p.device}, expected {device}"
|
|
|
|
|
|
def validate_adam_states_device(model, device: torch.device):
|
|
"""Validates that the sharded Adam optimizer states are on the specified device."""
|
|
for p in model.optimizer.single_partition_of_fp32_groups:
|
|
if p in model.optimizer.state:
|
|
for state_key in ['exp_avg', 'exp_avg_sq']:
|
|
if state_key in model.optimizer.state[p]:
|
|
state_tensor = model.optimizer.state[p][state_key]
|
|
assert state_tensor.device.type == device.type, f"Optimizer state '{state_key}' is on {state_tensor.device}, expected {device}"
|
|
|
|
|
|
def validate_grad_device(model, device: torch.device) -> None:
|
|
"""Validates that the sharded gradients are on the specified device."""
|
|
# This path is for before step() where gradients are in averaged_gradients
|
|
if model.optimizer.averaged_gradients:
|
|
for grad_list in model.optimizer.averaged_gradients.values():
|
|
if grad_list is not None:
|
|
for grad_tensor in grad_list:
|
|
assert grad_tensor.device.type == device.type, f"Gradient partition in averaged_gradients is on {grad_tensor.device}, expected {device}"
|
|
else:
|
|
# This path is for after step() or if grads are not in averaged_gradients
|
|
for p in model.optimizer.single_partition_of_fp32_groups:
|
|
if p.grad is not None:
|
|
assert p.grad.device.type == device.type, f"Gradient partition on hp_param.grad is on {p.grad.device}, expected {device}"
|
|
|
|
|
|
def is_offload_optimizer_enabled(config_dict):
|
|
return config_dict.get("zero_optimization", {}).get("offload_optimizer", {}).get("device", None) is not None
|
|
|
|
|
|
def is_only_offload_optimizer_states(offloaded_states, optimizer_offload_states):
|
|
if offloaded_states is None:
|
|
return False
|
|
offload_set = set(offloaded_states)
|
|
optim_states_set = set(optimizer_offload_states)
|
|
return offload_set - optim_states_set == set()
|
|
|
|
|
|
def run_model_zero12(model, param_groups, config_dict, hidden_dim, dtype, offloaded_states, pin_memory, non_blocking):
|
|
"""
|
|
This function runs a training step, offloads states, reloads them, and verifies correctness for ZeRO-1/2.
|
|
The logic is carefully structured to handle transient gradient states vs. persistent parameter/optimizer states.
|
|
"""
|
|
offload_device = OffloadDeviceEnum.cpu
|
|
offload_torch_device = torch.device(offload_device.value)
|
|
accelerator_device = torch.device(get_accelerator().current_device_name())
|
|
optimizer_device = offload_torch_device if is_offload_optimizer_enabled(config_dict) else accelerator_device
|
|
offload_only_optimizer_states = is_only_offload_optimizer_states(
|
|
offloaded_states, [OffloadStateTypeEnum.optim_states, OffloadStateTypeEnum.hp_params])
|
|
expect_memory_change = not (is_offload_optimizer_enabled(config_dict) and offload_only_optimizer_states)
|
|
|
|
model, _, _, _ = deepspeed.initialize(model=model, model_parameters=param_groups, config=config_dict)
|
|
|
|
data_loader = random_dataloader(model=model,
|
|
total_samples=10,
|
|
hidden_dim=hidden_dim,
|
|
device=model.device,
|
|
dtype=dtype)
|
|
dist.barrier()
|
|
|
|
# We only need one step to verify the logic
|
|
batch = next(iter(data_loader))
|
|
|
|
loss = model(batch[0], batch[1])
|
|
model.backward(loss)
|
|
|
|
# Determine if we are testing a transient state (gradients) or a persistent state
|
|
# REVERTED: Condition now only checks for lp_grads as it's the relevant transient state.
|
|
is_grad_test = offloaded_states is not None and OffloadStateTypeEnum.lp_grads in offloaded_states
|
|
|
|
if is_grad_test:
|
|
# --- TEST PATH FOR TRANSIENT GRADIENT STATE ---
|
|
# Gradients exist only between backward() and step(). We must test them here.
|
|
grads_expected = [[g.clone().detach() for g in grad_list]
|
|
for grad_list in model.optimizer.averaged_gradients.values() if grad_list is not None]
|
|
grad_numel = sum(sum(g.numel() for g in grad_list) for grad_list in grads_expected)
|
|
|
|
alloc_before_offload = get_accelerator().memory_allocated()
|
|
model.offload_states(include=offloaded_states,
|
|
device=offload_device,
|
|
pin_memory=pin_memory,
|
|
non_blocking=non_blocking)
|
|
alloc_after_offload = get_accelerator().memory_allocated()
|
|
|
|
if grad_numel > 0:
|
|
assert alloc_after_offload < alloc_before_offload, f"FAIL: Allocated memory for grads should decrease after offload {alloc_after_offload=} < {alloc_before_offload=}"
|
|
validate_grad_device(model, offload_torch_device)
|
|
|
|
model.reload_states()
|
|
alloc_after_reload = get_accelerator().memory_allocated()
|
|
|
|
if grad_numel > 0:
|
|
assert alloc_after_reload > alloc_after_offload, f"FAIL: Allocated memory for grads should increase after reload {alloc_after_reload=} > {alloc_after_offload=}"
|
|
validate_grad_device(model, accelerator_device)
|
|
|
|
reloaded_grads = [
|
|
grad_list for grad_list in model.optimizer.averaged_gradients.values() if grad_list is not None
|
|
]
|
|
assert len(grads_expected) == len(reloaded_grads), "FAIL: Number of gradient groups changed after reload"
|
|
for expected_list, reloaded_list in zip(grads_expected, reloaded_grads):
|
|
for expected_g, reloaded_g in zip(expected_list, reloaded_list):
|
|
assert torch.equal(expected_g, reloaded_g), "FAIL: Reloaded gradient data does not match original"
|
|
|
|
model.step()
|
|
|
|
if not is_grad_test:
|
|
# --- TEST PATH FOR PERSISTENT STATES (Params, Optimizer States) ---
|
|
# These states exist after step(), so we can test them here.
|
|
|
|
# --- Save state snapshots before offloading for data integrity check ---
|
|
lp_params_expected = [p.clone().detach() for p in model.parameters()]
|
|
hp_params_expected = [p.clone().detach() for p in model.optimizer.single_partition_of_fp32_groups]
|
|
|
|
adam_params_in_state_before = [
|
|
p for p in model.optimizer.single_partition_of_fp32_groups if p in model.optimizer.state
|
|
]
|
|
adam_exp_avg_expected = [
|
|
model.optimizer.state[p]['exp_avg'].clone().detach() for p in adam_params_in_state_before
|
|
]
|
|
adam_exp_avg_sq_expected = [
|
|
model.optimizer.state[p]['exp_avg_sq'].clone().detach() for p in adam_params_in_state_before
|
|
]
|
|
|
|
alloc_before_offload = get_accelerator().memory_allocated()
|
|
model.offload_states(include=offloaded_states,
|
|
device=offload_device,
|
|
pin_memory=pin_memory,
|
|
non_blocking=non_blocking)
|
|
alloc_after_offload = get_accelerator().memory_allocated()
|
|
|
|
if expect_memory_change:
|
|
assert alloc_after_offload < alloc_before_offload, f"FAIL: Allocated memory for persistent state {offloaded_states} should decrease after offload"
|
|
|
|
if offloaded_states is None or OffloadStateTypeEnum.lp_params in offloaded_states:
|
|
validate_lp_params_device(model, offload_torch_device)
|
|
if offloaded_states is None or OffloadStateTypeEnum.hp_params in offloaded_states:
|
|
validate_hp_params_device(model, offload_torch_device)
|
|
if offloaded_states is None or OffloadStateTypeEnum.optim_states in offloaded_states:
|
|
validate_adam_states_device(model, offload_torch_device)
|
|
|
|
model.reload_states()
|
|
alloc_after_reload = get_accelerator().memory_allocated()
|
|
|
|
if expect_memory_change:
|
|
assert alloc_after_reload > alloc_after_offload, f"FAIL: Allocated memory for persistent state {offloaded_states} should increase after reload"
|
|
|
|
# --- Verify restored data integrity ---
|
|
for expected, restored in zip(lp_params_expected, model.parameters()):
|
|
assert torch.equal(expected, restored), "FAIL: Reloaded LP param data does not match original"
|
|
|
|
for expected, restored in zip(hp_params_expected, model.optimizer.single_partition_of_fp32_groups):
|
|
assert torch.equal(expected, restored), "FAIL: Reloaded HP param data does not match original"
|
|
|
|
adam_params_in_state_after = [
|
|
p for p in model.optimizer.single_partition_of_fp32_groups if p in model.optimizer.state
|
|
]
|
|
assert len(adam_params_in_state_before) == len(
|
|
adam_params_in_state_after), "FAIL: Number of params in optimizer state changed after reload"
|
|
|
|
for expected, p in zip(adam_exp_avg_expected, adam_params_in_state_after):
|
|
assert torch.equal(
|
|
expected, model.optimizer.state[p]['exp_avg']), "FAIL: Reloaded 'exp_avg' data does not match original"
|
|
for expected, p in zip(adam_exp_avg_sq_expected, adam_params_in_state_after):
|
|
assert torch.equal(
|
|
expected,
|
|
model.optimizer.state[p]['exp_avg_sq']), "FAIL: Reloaded 'exp_avg_sq' data does not match original"
|
|
|
|
# --- FINAL VALIDATION FOR ALL TESTS ---
|
|
validate_lp_params_device(model, accelerator_device)
|
|
validate_hp_params_device(model, optimizer_device)
|
|
validate_adam_states_device(model, optimizer_device)
|
|
|
|
assert torch.any(torch.ne(list(model.parameters())[0], 0.0))
|
|
|
|
|
|
@pytest.mark.parametrize("included_state", [
|
|
OffloadStateTypeEnum.optim_states, OffloadStateTypeEnum.lp_grads, OffloadStateTypeEnum.hp_params,
|
|
OffloadStateTypeEnum.lp_params, None
|
|
])
|
|
@pytest.mark.parametrize("pin_memory", [False, True])
|
|
@pytest.mark.parametrize("non_blocking", [False, True])
|
|
@pytest.mark.parametrize("zero_stage", [1, 2])
|
|
@pytest.mark.parametrize("static_offload_optimizer", [False, True])
|
|
class TestDynamicOffloadStatesZero12(DistributedTest):
|
|
world_size = 2
|
|
|
|
def test_dynamic_offload_states_zero12(self, included_state, pin_memory, non_blocking, zero_stage,
|
|
static_offload_optimizer):
|
|
hidden_dim = 1024
|
|
config_dict = {
|
|
"train_micro_batch_size_per_gpu": 1,
|
|
"optimizer": {
|
|
"type": "Adam",
|
|
"params": {
|
|
"lr": 1e-6
|
|
}
|
|
},
|
|
"zero_optimization": {
|
|
"stage": zero_stage
|
|
},
|
|
"bf16": {
|
|
"enabled": True
|
|
}
|
|
}
|
|
if static_offload_optimizer:
|
|
config_dict["zero_optimization"]["offload_optimizer"] = {"device": "cpu"}
|
|
model = SimpleModel(hidden_dim, nlayers=4)
|
|
param_groups = [{
|
|
"params": [p for n, p in model.named_parameters() if 'bias' not in n],
|
|
"weight_decay": 0.1
|
|
}, {
|
|
"params": [p for n, p in model.named_parameters() if 'bias' in n],
|
|
"weight_decay": 0.0
|
|
}]
|
|
offloaded_states = None if included_state is None else [included_state]
|
|
run_model_zero12(model, param_groups, config_dict, hidden_dim, torch.bfloat16, offloaded_states, pin_memory,
|
|
non_blocking)
|
|
|
|
|
|
# ==============================================================================
|
|
# ZeRO-3 TESTS
|
|
# ==============================================================================
|
|
|
|
|
|
def validate_device(model, state_device: dict[OffloadStateTypeEnum, torch.device], offloaded_states) -> None:
|
|
|
|
def compare_device(state) -> bool:
|
|
devices = get_state_devices(model, state)
|
|
return len(devices) == 1 and state_device[state] in devices
|
|
|
|
for state in OffloadStateTypeEnum:
|
|
if offloaded_states is None or state in offloaded_states:
|
|
if state == OffloadStateTypeEnum.contiguous_grad_buffer and state_device[state] == torch.device("cpu"):
|
|
assert len(get_state_devices(model,
|
|
state)) == 0, f"State {state} must be removed after offload_states()"
|
|
else:
|
|
assert compare_device(state), f"State {state} is not on device {state_device[state]}"
|
|
|
|
|
|
def run_model_zero3(model, param_groups, config_dict, hidden_dim, dtype, offloaded_states, pin_memory, non_blocking):
|
|
# Currently we only support OffloadDeviceEnum.cpu
|
|
offload_device = OffloadDeviceEnum.cpu
|
|
offload_torch_device = torch.device(offload_device.value)
|
|
accelerator_device = torch.device(get_accelerator().current_device_name())
|
|
optimizer_device = offload_torch_device if is_offload_optimizer_enabled(config_dict) else accelerator_device
|
|
offload_only_optimizer_states = is_only_offload_optimizer_states(
|
|
offloaded_states,
|
|
[OffloadStateTypeEnum.optim_states, OffloadStateTypeEnum.hp_params, OffloadStateTypeEnum.lp_grads])
|
|
expect_memory_change = not (is_offload_optimizer_enabled(config_dict) and offload_only_optimizer_states)
|
|
|
|
offload_state_device: dict[OffloadStateTypeEnum, torch.device] = {
|
|
OffloadStateTypeEnum.hp_params: offload_torch_device,
|
|
OffloadStateTypeEnum.lp_params: offload_torch_device,
|
|
OffloadStateTypeEnum.optim_states: offload_torch_device,
|
|
OffloadStateTypeEnum.lp_grads: offload_torch_device,
|
|
OffloadStateTypeEnum.contiguous_grad_buffer: offload_torch_device,
|
|
}
|
|
|
|
reload_state_device: dict[OffloadStateTypeEnum, torch.device] = {
|
|
OffloadStateTypeEnum.hp_params: optimizer_device,
|
|
OffloadStateTypeEnum.lp_params: accelerator_device,
|
|
OffloadStateTypeEnum.optim_states: optimizer_device,
|
|
OffloadStateTypeEnum.lp_grads: optimizer_device,
|
|
OffloadStateTypeEnum.contiguous_grad_buffer: accelerator_device,
|
|
}
|
|
|
|
model, _, _, _ = deepspeed.initialize(model=model, model_parameters=param_groups, config=config_dict)
|
|
data_loader = random_dataloader(model=model,
|
|
total_samples=10,
|
|
hidden_dim=hidden_dim,
|
|
device=model.device,
|
|
dtype=dtype)
|
|
dist.barrier()
|
|
for batch in data_loader:
|
|
loss = model(batch[0], batch[1])
|
|
model.backward(loss)
|
|
model.step()
|
|
|
|
hp_params_expected = [safe_get_local_fp32_param(p).clone() for p in model.parameters()]
|
|
lp_params_expected = [p.ds_tensor.clone() for p in model.parameters()]
|
|
lp_grads_expected = model.optimizer.grad_partitions_flat_buffer.clone()
|
|
adam_exp_avg_expected = [safe_get_local_optimizer_state(p, "exp_avg").clone() for p in model.parameters()]
|
|
adam_exp_avg_sq = [safe_get_local_optimizer_state(p, "exp_avg_sq").clone() for p in model.parameters()]
|
|
|
|
# Start offloading
|
|
alloc_before_offload = get_accelerator().memory_allocated()
|
|
model.offload_states(include=offloaded_states,
|
|
device=offload_device,
|
|
pin_memory=pin_memory,
|
|
non_blocking=non_blocking)
|
|
alloc_after_offload = get_accelerator().memory_allocated()
|
|
|
|
if expect_memory_change:
|
|
assert alloc_after_offload < alloc_before_offload, f"FAIL: Allocated memory should decrease after offload {alloc_after_offload=} < {alloc_before_offload=}"
|
|
validate_device(model, offload_state_device, offloaded_states)
|
|
|
|
# Reload states
|
|
model.reload_states()
|
|
alloc_after_reload = get_accelerator().memory_allocated()
|
|
|
|
if expect_memory_change:
|
|
assert alloc_after_reload > alloc_after_offload, f"FAIL: Allocated memory should increase after offload back {alloc_after_reload=} > {alloc_after_offload=}"
|
|
|
|
# Verify restored states
|
|
hp_param_restored = [safe_get_local_fp32_param(p) for p in model.parameters()]
|
|
for hp_param_expected, hp_param_restored in zip(hp_params_expected, hp_param_restored):
|
|
assert torch.equal(hp_param_expected, hp_param_restored)
|
|
|
|
lp_param_restored = [p.ds_tensor for p in model.parameters()]
|
|
|
|
for lp_param_expected, lp_param_restored in zip(lp_params_expected, lp_param_restored):
|
|
assert torch.equal(lp_param_expected, lp_param_restored)
|
|
|
|
assert torch.equal(lp_grads_expected, model.optimizer.grad_partitions_flat_buffer)
|
|
|
|
adam_exp_avg_restored = [safe_get_local_optimizer_state(p, "exp_avg") for p in model.parameters()]
|
|
for adam_exp_avg_expected, adam_exp_avg_restored in zip(adam_exp_avg_expected, adam_exp_avg_restored):
|
|
assert torch.equal(adam_exp_avg_expected, adam_exp_avg_restored)
|
|
|
|
adam_exp_avg_sq_restored = [safe_get_local_optimizer_state(p, "exp_avg_sq") for p in model.parameters()]
|
|
for adam_exp_avg_sq_expected, adam_exp_avg_sq_restored in zip(adam_exp_avg_sq, adam_exp_avg_sq_restored):
|
|
assert torch.equal(adam_exp_avg_sq_expected, adam_exp_avg_sq_restored)
|
|
|
|
validate_device(model, reload_state_device, offloaded_states)
|
|
|
|
# Needed in ZeRO 3. Not doing so can give memory leak
|
|
model.destroy()
|
|
|
|
|
|
@pytest.mark.parametrize("included_state", [
|
|
OffloadStateTypeEnum.hp_params, OffloadStateTypeEnum.lp_params, OffloadStateTypeEnum.optim_states,
|
|
OffloadStateTypeEnum.lp_grads, OffloadStateTypeEnum.contiguous_grad_buffer, None
|
|
])
|
|
@pytest.mark.parametrize("pin_memory", [False, True])
|
|
@pytest.mark.parametrize("non_blocking", [False, True])
|
|
@pytest.mark.parametrize("static_offload_optimizer", [False, True])
|
|
class TestDynamicOffloadStatesZero3(DistributedTest):
|
|
# Need multiple gpus to test possible hanging
|
|
world_size = 2
|
|
|
|
def test_dynamic_offload_states_zero3(self, included_state, pin_memory, non_blocking, static_offload_optimizer):
|
|
hidden_dim = 1024
|
|
|
|
config_dict = {
|
|
"train_micro_batch_size_per_gpu": 1,
|
|
"optimizer": {
|
|
"type": "Adam",
|
|
"params": {
|
|
"lr": 1e-6
|
|
}
|
|
},
|
|
"zero_optimization": {
|
|
"stage": 3,
|
|
}
|
|
}
|
|
config_dict["bf16"] = {"enabled": True}
|
|
if static_offload_optimizer:
|
|
config_dict["zero_optimization"]["offload_optimizer"] = {"device": "cpu"}
|
|
|
|
with deepspeed.zero.Init(config_dict_or_path=config_dict):
|
|
model = SimpleModel(hidden_dim, nlayers=4)
|
|
|
|
param_groups = [{
|
|
"params": [p for n, p in model.named_parameters() if not 'bias' in n],
|
|
"weight_decay": 0.1
|
|
}, {
|
|
"params": [p for n, p in model.named_parameters() if 'bias' in n],
|
|
"weight_decay": 0.0
|
|
}]
|
|
offloaded_states = None if included_state is None else [included_state]
|
|
run_model_zero3(model, param_groups, config_dict, hidden_dim, torch.bfloat16, offloaded_states, pin_memory,
|
|
non_blocking)
|