wehub-resource-sync
135 lines
4.8 KiB
Python
135 lines
4.8 KiB
Python
# Copyright (c) Microsoft Corporation.
|
|
# SPDX-License-Identifier: Apache-2.0
|
|
|
|
# DeepSpeed Team
|
|
|
|
import torch
|
|
import pytest
|
|
import deepspeed
|
|
from deepspeed.runtime.zero.stage3 import DeepSpeedZeroOptimizer_Stage3
|
|
from deepspeed.utils import safe_get_local_grad, safe_set_local_grad
|
|
from deepspeed.accelerator import get_accelerator
|
|
from unit.simple_model import SimpleModel
|
|
import os
|
|
|
|
|
|
def get_config(precision, clip_value, offload_device="cpu"):
|
|
config = {
|
|
"train_batch_size": 8,
|
|
"steps_per_print": 1,
|
|
"optimizer": {
|
|
"type": "Adam",
|
|
"params": {
|
|
"lr": 1e-4
|
|
}
|
|
},
|
|
"zero_optimization": {
|
|
"stage": 3,
|
|
"offload_optimizer": {
|
|
"device": offload_device
|
|
},
|
|
"contiguous_gradients": True,
|
|
"overlap_comm": False,
|
|
},
|
|
"gradient_clipping": 1.0,
|
|
}
|
|
|
|
if precision == "fp16":
|
|
config["fp16"] = {
|
|
"enabled": True,
|
|
"loss_scale": 1024,
|
|
"initial_scale_power": 10,
|
|
}
|
|
elif precision == "bf16":
|
|
config["bf16"] = {
|
|
"enabled": True,
|
|
}
|
|
|
|
return config
|
|
|
|
|
|
@pytest.mark.parametrize("precision,clip_value,offload_device", [
|
|
("fp16", 0.5, "cpu"),
|
|
("bf16", 0.05, "cpu"),
|
|
("fp16", 0.5, "none"),
|
|
("bf16", 0.05, "none"),
|
|
])
|
|
class TestZeroGradClip():
|
|
world_size = 1
|
|
|
|
def test_grad_clip_and_norm_update(self, precision, clip_value, offload_device):
|
|
"""Test custom gradient clipping with configurations and to check if the norm_groups are updated correctly"""
|
|
config_dict = get_config(precision, clip_value, offload_device)
|
|
|
|
model = SimpleModel(hidden_dim=10)
|
|
|
|
# Set up distributed environment variables
|
|
os.environ['LOCAL_RANK'] = '0'
|
|
os.environ['RANK'] = '0'
|
|
os.environ['WORLD_SIZE'] = '1'
|
|
os.environ['MASTER_ADDR'] = '127.0.0.1'
|
|
os.environ['MASTER_PORT'] = '29500'
|
|
|
|
try:
|
|
model_engine, optimizer, _, _ = deepspeed.initialize(args=None,
|
|
model=model,
|
|
config=config_dict,
|
|
model_parameters=model.parameters(),
|
|
dist_init_required=True)
|
|
except Exception as e:
|
|
pytest.skip("Could not initialize deepspeed")
|
|
|
|
assert isinstance(optimizer, DeepSpeedZeroOptimizer_Stage3)
|
|
|
|
torch.manual_seed(1670)
|
|
inputs = torch.randn(8, 10, device=model_engine.device)
|
|
targets = torch.randn(8, 10, device=model_engine.device)
|
|
|
|
if model_engine.fp16_enabled() and get_accelerator().is_fp16_supported():
|
|
inputs = inputs.half()
|
|
targets = targets.half()
|
|
elif model_engine.bfloat16_enabled() and get_accelerator().is_bf16_supported():
|
|
inputs = inputs.bfloat16()
|
|
targets = targets.bfloat16()
|
|
else:
|
|
pytest.skip("Unsupported precision")
|
|
|
|
loss = model_engine(inputs, targets)
|
|
model_engine.backward(loss)
|
|
|
|
pre_clip_norm_groups = optimizer._get_norm_groups()
|
|
pre_clip_global_norm = torch.linalg.vector_norm(torch.stack(pre_clip_norm_groups))
|
|
|
|
modified_count = 0
|
|
|
|
for param in model_engine.parameters():
|
|
if not hasattr(param, 'ds_id'):
|
|
continue
|
|
|
|
grad = safe_get_local_grad(param)
|
|
if grad is not None:
|
|
pre_clip_norm = grad.norm().item()
|
|
clamped_grad = torch.clamp(grad, -clip_value, clip_value)
|
|
post_clip_norm = clamped_grad.norm().item()
|
|
|
|
if pre_clip_norm > clip_value:
|
|
# Checks if the post-clip norm is less than the pre-clip norm
|
|
assert post_clip_norm < pre_clip_norm, f"Post-clip norm should be < pre-clip norm for param {param.ds_id}"
|
|
|
|
safe_set_local_grad(param, clamped_grad)
|
|
modified_count += 1
|
|
|
|
# Get post-clip state
|
|
post_clip_norm_groups = optimizer._get_norm_groups()
|
|
post_clip_global_norm = torch.linalg.vector_norm(torch.stack(post_clip_norm_groups))
|
|
|
|
assert modified_count > 0, "No parameters were modified during clipping"
|
|
assert post_clip_global_norm.item() < pre_clip_global_norm.item(
|
|
), f"Post-clip norm {post_clip_global_norm.item():.6f} should be < pre-clip norm {pre_clip_global_norm.item():.6f}"
|
|
|
|
model_engine.step()
|
|
final_norm = optimizer._global_grad_norm
|
|
if pre_clip_global_norm.item() > clip_value:
|
|
assert post_clip_global_norm.item() < pre_clip_global_norm.item(
|
|
), "Global norm should be reduced after clipping when pre-clip norm > clip_value"
|