wehub-resource-sync
110 lines
3.8 KiB
Python
110 lines
3.8 KiB
Python
# Copyright (c) Microsoft Corporation.
|
|
# SPDX-License-Identifier: Apache-2.0
|
|
|
|
# DeepSpeed Team
|
|
|
|
import torch
|
|
import deepspeed
|
|
import pytest
|
|
from unit.common import DistributedTest
|
|
|
|
import deepspeed.comm as dist
|
|
from deepspeed.linear import LoRAConfig, init_lora
|
|
from deepspeed.linear.optimized_linear import LoRAOptimizedLinear
|
|
from unit.simple_model import random_dataloader, SimpleModel
|
|
|
|
try:
|
|
import transformers
|
|
except ImportError:
|
|
transformers = None
|
|
|
|
if transformers is None:
|
|
pytest.skip("transformers is required for this test", allow_module_level=True)
|
|
|
|
|
|
def injection_assert(model):
|
|
# pick out random linear that should have been replaced and initialized
|
|
q_proj = model.model.layers[1].self_attn.q_proj
|
|
|
|
assert isinstance(q_proj, LoRAOptimizedLinear), "injection did not happen"
|
|
assert q_proj._initialized, "lora was not initialized properly"
|
|
assert isinstance(q_proj.lora_weight_1, torch.nn.Linear)
|
|
assert isinstance(q_proj.lora_weight_2, torch.nn.Linear)
|
|
|
|
|
|
class TestEngine(DistributedTest):
|
|
world_size = 2
|
|
|
|
def test_model(self):
|
|
lora_config = LoRAConfig(lora_r=16, lora_alpha=16, base_weight_sharding=2)
|
|
quant_config = None
|
|
hidden_dim = 64
|
|
nlayers = 4
|
|
|
|
with deepspeed.linear.Init(lora_config=lora_config, quant_config=quant_config):
|
|
model = SimpleModel(hidden_dim=hidden_dim, nlayers=nlayers)
|
|
|
|
init_lora(model)
|
|
|
|
model_norms = [model.linears[i].weight.norm().item() for i in range(nlayers)]
|
|
|
|
ds_config = {
|
|
"train_batch_size": 2,
|
|
"steps_per_print": 1,
|
|
"bf16": {
|
|
"enabled": True
|
|
},
|
|
"optimizer": {
|
|
"type": "Adam",
|
|
"params": {
|
|
"lr": 0.00015
|
|
}
|
|
},
|
|
"zero_optimization": {
|
|
"stage": 1
|
|
}
|
|
}
|
|
model, *_ = deepspeed.initialize(config=ds_config, model=model, model_parameters=model.parameters())
|
|
|
|
engine_norms = [model.module.linears[i].weight.norm().item() for i in range(nlayers)]
|
|
|
|
# Ensure that sharded weights are not broadcast during engine init
|
|
assert engine_norms == model_norms, f"{dist.get_rank()=} base weight norms are not the same after engine init, {engine_norms=} != {model_norms=}"
|
|
|
|
data_loader = random_dataloader(model=model,
|
|
total_samples=50,
|
|
hidden_dim=hidden_dim,
|
|
device=model.device,
|
|
dtype=torch.bfloat16)
|
|
for n, batch in enumerate(data_loader):
|
|
loss = model(batch[0], batch[1])
|
|
model.backward(loss)
|
|
model.step()
|
|
|
|
|
|
@pytest.mark.skip(
|
|
"Skipping test for now - the context manager has an issue with ._initialized and .disabled - worked with older transformers probably because it was setting some flags with the same name"
|
|
)
|
|
class TestInitTransformers(DistributedTest):
|
|
world_size = 2
|
|
|
|
def test_pretrained_init(self):
|
|
lora_config = LoRAConfig(lora_r=16, lora_alpha=16, base_weight_sharding=2)
|
|
quant_config = None
|
|
|
|
with deepspeed.linear.Init(lora_config=lora_config, quant_config=quant_config):
|
|
model = transformers.AutoModelForCausalLM.from_pretrained("llamafactory/tiny-random-Llama-3")
|
|
|
|
injection_assert(model)
|
|
|
|
def test_config_init(self):
|
|
lora_config = LoRAConfig(lora_r=16, lora_alpha=16, base_weight_sharding=2)
|
|
quant_config = None
|
|
|
|
config = transformers.AutoConfig.from_pretrained("llamafactory/tiny-random-Llama-3")
|
|
|
|
with deepspeed.linear.Init(lora_config=lora_config, quant_config=quant_config):
|
|
model = transformers.AutoModelForCausalLM.from_config(config)
|
|
|
|
injection_assert(model)
|