wehub-resource-sync
289 lines
11 KiB
Python
289 lines
11 KiB
Python
# Copyright (c) Microsoft Corporation.
|
|
# SPDX-License-Identifier: Apache-2.0
|
|
|
|
# DeepSpeed Team
|
|
import pytest
|
|
import deepspeed.comm as dist
|
|
from torch.nn import Module
|
|
|
|
from unit.common import DistributedTest
|
|
from unit.simple_model import random_dataloader
|
|
|
|
import deepspeed
|
|
|
|
from deepspeed.runtime.zero.config import DeepSpeedZeroConfig
|
|
|
|
import torch.nn as nn
|
|
import torch
|
|
|
|
from transformers import AutoModelForCausalLM, AutoTokenizer
|
|
from torch.utils.data import DataLoader
|
|
|
|
import numpy as np
|
|
|
|
|
|
class NNModel(nn.Module):
|
|
|
|
def __init__(self, h_dim=1024, n_layers=2):
|
|
super(NNModel, self).__init__()
|
|
self.layers = nn.ModuleList([nn.Linear(h_dim, h_dim) for i in range(n_layers)])
|
|
self.cross_entropy_loss = nn.CrossEntropyLoss()
|
|
|
|
def forward(self, x, y):
|
|
for layer in self.layers:
|
|
x = layer(x)
|
|
return self.cross_entropy_loss(x, y)
|
|
|
|
|
|
def test_zero_hpz_partition_size_config():
|
|
config = DeepSpeedZeroConfig(**{"zero_hpz_partition_size": 4})
|
|
assert config.zero_hpz_partition_size == 4
|
|
|
|
|
|
def _assert_no_secondary_tensor_group(model: Module) -> None:
|
|
for _, param in model.named_parameters():
|
|
assert param.ds_secondary_tensor is None
|
|
assert param.ds_zero_param_process_group is None
|
|
|
|
|
|
def _check_secondary_tensor_existence(model: Module) -> None:
|
|
for _, param in model.named_parameters():
|
|
if param.ds_secondary_tensor is not None:
|
|
return True
|
|
return False
|
|
|
|
|
|
def _assert_secondary_tensor_size(model: Module) -> None:
|
|
for name, param in model.named_parameters():
|
|
assert param.ds_secondary_tensor is not None, f"param {param.ds_id}:{name} does not have secondary tensor"
|
|
assert param.ds_secondary_tensor.size()[0] % param.ds_tensor.size()[0] == 0
|
|
|
|
|
|
#Large sweep along hidden dim, num_layers, and zpg of different sizes
|
|
#Assert when zpg=1 that secondary group and tensors are invalid
|
|
@pytest.mark.sequential
|
|
@pytest.mark.parametrize("h_dim", [1024])
|
|
@pytest.mark.parametrize("n_layers", [9])
|
|
@pytest.mark.parametrize("zpg", [1, 2, 4])
|
|
class TestZeroPPConfigSweep(DistributedTest):
|
|
world_size = 4
|
|
|
|
def test(self, h_dim: int, n_layers: int, zpg: int) -> None:
|
|
config_dict = {
|
|
"train_micro_batch_size_per_gpu": 1,
|
|
"zero_optimization": {
|
|
"stage": 3,
|
|
"stage3_max_reuse_distance": 0,
|
|
"zero_hpz_partition_size": zpg,
|
|
"zero_quantized_weights": True,
|
|
"zero_quantized_gradients": True,
|
|
"contiguous_gradients": True,
|
|
"overlap_comm": True,
|
|
},
|
|
"optimizer": {
|
|
"type": "Adam",
|
|
"params": {
|
|
"lr": 1.
|
|
}
|
|
},
|
|
"fp16": {
|
|
"enabled": True,
|
|
"loss_scale": 1.,
|
|
}
|
|
}
|
|
|
|
model = NNModel(h_dim, n_layers)
|
|
model, _, _, _ = deepspeed.initialize(model=model, model_parameters=model.parameters(), config=config_dict)
|
|
data_loader = random_dataloader(model=model,
|
|
total_samples=20,
|
|
hidden_dim=h_dim,
|
|
device=model.device,
|
|
dtype=torch.float16)
|
|
dist.barrier()
|
|
if zpg == 1:
|
|
_assert_no_secondary_tensor_group(model)
|
|
|
|
for n, batch in enumerate(data_loader):
|
|
if n == 0 and zpg != 1:
|
|
_assert_secondary_tensor_size(model)
|
|
loss = model(batch[0], batch[1])
|
|
model.backward(loss)
|
|
model.step()
|
|
|
|
def test_eval(self, h_dim: int, n_layers: int, zpg: int) -> None:
|
|
# in this test case, we are testing that hpz should be enabled when eval mode is on
|
|
config_dict = {
|
|
"train_micro_batch_size_per_gpu": 1,
|
|
"zero_optimization": {
|
|
"stage": 3,
|
|
"stage3_max_reuse_distance": 0,
|
|
"zero_hpz_partition_size": zpg,
|
|
"contiguous_gradients": True,
|
|
"overlap_comm": True,
|
|
},
|
|
"optimizer": {
|
|
"type": "Adam",
|
|
"params": {
|
|
"lr": 1.
|
|
}
|
|
},
|
|
"fp16": {
|
|
"enabled": True,
|
|
"loss_scale": 1.,
|
|
}
|
|
}
|
|
|
|
model = NNModel(h_dim, n_layers)
|
|
model, _, _, _ = deepspeed.initialize(model=model, model_parameters=model.parameters(), config=config_dict)
|
|
data_loader = random_dataloader(model=model,
|
|
total_samples=20,
|
|
hidden_dim=h_dim,
|
|
device=model.device,
|
|
dtype=torch.float16)
|
|
dist.barrier()
|
|
if zpg == 1:
|
|
_assert_no_secondary_tensor_group(model)
|
|
|
|
for n, batch in enumerate(data_loader):
|
|
if zpg != 1:
|
|
# here we check that the hpz is enabled when the previous iteration does not update the model
|
|
_assert_secondary_tensor_size(model)
|
|
with torch.no_grad():
|
|
loss = model(batch[0], batch[1])
|
|
|
|
def test_gradient_accumulation(self, h_dim: int, n_layers: int, zpg: int) -> None:
|
|
# in this test case, we are testing that hpz should be enabled for the intermediate gradient accumulation steps
|
|
# In this test, we should disable loss_scale
|
|
config_dict = {
|
|
"train_micro_batch_size_per_gpu": 1,
|
|
"gradient_accumulation_steps": 3,
|
|
"zero_optimization": {
|
|
"stage": 3,
|
|
"stage3_max_reuse_distance": 0,
|
|
"zero_hpz_partition_size": zpg,
|
|
"contiguous_gradients": True,
|
|
"overlap_comm": True,
|
|
},
|
|
"optimizer": {
|
|
"type": "Adam",
|
|
"params": {
|
|
"lr": 1.
|
|
}
|
|
},
|
|
"fp16": {
|
|
"enabled": True,
|
|
"loss_scale": 0.,
|
|
}
|
|
}
|
|
|
|
model = NNModel(h_dim, n_layers)
|
|
model, _, _, _ = deepspeed.initialize(model=model, model_parameters=model.parameters(), config=config_dict)
|
|
data_loader = random_dataloader(model=model,
|
|
total_samples=20,
|
|
hidden_dim=h_dim,
|
|
device=model.device,
|
|
dtype=torch.float16)
|
|
dist.barrier()
|
|
if zpg == 1:
|
|
_assert_no_secondary_tensor_group(model)
|
|
|
|
for n, batch in enumerate(data_loader):
|
|
if n == 0 and zpg != 1:
|
|
_assert_secondary_tensor_size(model)
|
|
# here we cannot assert that secondary tensor does not exist because the gradient is likely overflowed as we use random data
|
|
if n > 0 and n % 3 != 0 and zpg != 1:
|
|
# if the previous iteration does not update the model, then the hpz should be enabled
|
|
assert _check_secondary_tensor_existence(model), f"n={n}"
|
|
loss = model(batch[0], batch[1])
|
|
model.backward(loss)
|
|
model.step()
|
|
|
|
|
|
@pytest.mark.nightly
|
|
@pytest.mark.parametrize("model_name", ["gpt2"])
|
|
class TestZeroPPConvergence(DistributedTest):
|
|
world_size = 4
|
|
|
|
def load_and_prepare_data(self, model_name):
|
|
"""Load model, tokenizer and dataset, and prepare data loader."""
|
|
from datasets import load_dataset
|
|
|
|
# Load model and tokenizer
|
|
model = AutoModelForCausalLM.from_pretrained(model_name)
|
|
tokenizer = AutoTokenizer.from_pretrained(model_name)
|
|
tokenizer.pad_token = tokenizer.eos_token
|
|
|
|
# Load and tokenize dataset
|
|
dataset = load_dataset("wikitext", 'wikitext-103-raw-v1', split='train[:1%]').filter(lambda x: x["text"])
|
|
|
|
def tokenize_function(examples):
|
|
# Tokenize and ensure 'labels' are the same as 'input_ids'
|
|
tokenized_output = tokenizer(examples["text"], padding="max_length", truncation=True, return_tensors='pt')
|
|
tokenized_output["labels"] = tokenized_output["input_ids"].clone()
|
|
return tokenized_output
|
|
|
|
tokenized_dataset = dataset.map(tokenize_function, batched=True)
|
|
tokenized_dataset.set_format('torch', columns=['input_ids', 'attention_mask', 'labels'])
|
|
|
|
# Create data loader
|
|
data_loader = DataLoader(tokenized_dataset, batch_size=1, shuffle=False)
|
|
return model, data_loader
|
|
|
|
def get_loss(self, model, data_loader, config_dict, step=500):
|
|
"""Train the model and calculate average loss."""
|
|
# Initialize DeepSpeed
|
|
model, _, _, _ = deepspeed.initialize(model=model, model_parameters=model.parameters(), config=config_dict)
|
|
dist.barrier()
|
|
model.train()
|
|
|
|
# Training loop
|
|
losses = []
|
|
for n, batch in enumerate(data_loader):
|
|
if n >= step:
|
|
break
|
|
batch = {k: v.to(model.device) for k, v in batch.items()}
|
|
outputs = model(**batch)
|
|
loss = outputs.loss
|
|
model.backward(loss)
|
|
model.step()
|
|
losses.append(loss.item())
|
|
|
|
return np.nanmean(losses[-100:])
|
|
|
|
def get_config_dict(self, use_quantized_weights=False, use_hpz=False):
|
|
"""Generate the configuration dictionary for DeepSpeed."""
|
|
config = {
|
|
"train_micro_batch_size_per_gpu": 1,
|
|
"zero_optimization": {
|
|
"stage": 3,
|
|
"stage3_max_reuse_distance": 0,
|
|
"contiguous_gradients": True,
|
|
"overlap_comm": True,
|
|
},
|
|
"optimizer": {
|
|
"type": "Adam",
|
|
"params": {
|
|
"lr": 1e-5
|
|
}
|
|
},
|
|
"fp16": {
|
|
"enabled": True
|
|
}
|
|
}
|
|
if use_quantized_weights:
|
|
config["zero_optimization"]["zero_quantized_weights"] = True
|
|
if use_hpz:
|
|
config["zero_optimization"]["zero_hpz_partition_size"] = self.world_size // 2
|
|
return config
|
|
|
|
def test(self, model_name):
|
|
torch.manual_seed(0)
|
|
model, data_loader = self.load_and_prepare_data(model_name)
|
|
zeropp_loss = self.get_loss(model, data_loader, self.get_config_dict(use_quantized_weights=True, use_hpz=True))
|
|
model, data_loader = self.load_and_prepare_data(model_name)
|
|
baseline_loss = self.get_loss(model, data_loader, self.get_config_dict())
|
|
|
|
# Output and assert
|
|
print(f"zeropp_loss={zeropp_loss}, baseline_loss={baseline_loss}")
|
|
assert zeropp_loss < baseline_loss * 1.1, f"zeropp_loss={zeropp_loss}, baseline_loss={baseline_loss}"
|