wehub-resource-sync
94 lines
3.6 KiB
Python
94 lines
3.6 KiB
Python
# Copyright (c) Microsoft Corporation.
|
|
# SPDX-License-Identifier: Apache-2.0
|
|
|
|
# DeepSpeed Team
|
|
|
|
import torch
|
|
|
|
from unit.common import DistributedTest
|
|
|
|
from transformers import GPT2Config, VisionEncoderDecoderConfig, VisionEncoderDecoderModel, ViTConfig
|
|
from transformers.integrations.deepspeed import HfDeepSpeedConfig
|
|
|
|
import deepspeed
|
|
|
|
|
|
def _create_tiny_vision_encoder_decoder_model(model_path):
|
|
encoder_config = ViTConfig(image_size=8,
|
|
patch_size=4,
|
|
num_hidden_layers=1,
|
|
hidden_size=8,
|
|
num_attention_heads=2,
|
|
intermediate_size=16)
|
|
decoder_config = GPT2Config(vocab_size=32,
|
|
n_positions=16,
|
|
n_embd=8,
|
|
n_layer=1,
|
|
n_head=2,
|
|
bos_token_id=0,
|
|
eos_token_id=1,
|
|
add_cross_attention=True,
|
|
is_decoder=True)
|
|
config = VisionEncoderDecoderConfig.from_encoder_decoder_configs(encoder_config, decoder_config)
|
|
model = VisionEncoderDecoderModel(config)
|
|
model.save_pretrained(model_path, safe_serialization=False)
|
|
|
|
|
|
class TestNestingInit(DistributedTest):
|
|
world_size = 1
|
|
|
|
def test_nesting_init(self):
|
|
ds_config = dict(train_batch_size=1, zero_optimization=dict(stage=3))
|
|
|
|
with deepspeed.zero.Init(config_dict_or_path=ds_config):
|
|
with deepspeed.zero.Init(config_dict_or_path=ds_config):
|
|
model = torch.nn.Linear(4, 4)
|
|
|
|
# ensure that zero3 processed the parameter
|
|
assert hasattr(model.weight, "ds_id")
|
|
|
|
deepspeed_engine, *_ = deepspeed.initialize(model=model, config_params=ds_config)
|
|
|
|
|
|
class TestShutdownInNestingInit(DistributedTest):
|
|
world_size = 1
|
|
|
|
def test_shutdown_in_nesting_init(self):
|
|
ds_config = dict(train_batch_size=1, zero_optimization=dict(stage=3))
|
|
|
|
with deepspeed.zero.Init(config_dict_or_path=ds_config):
|
|
with deepspeed.zero.Init(config_dict_or_path=ds_config):
|
|
model1 = torch.nn.Linear(4, 4)
|
|
|
|
assert hasattr(model1.weight, "ds_id")
|
|
deepspeed_engine1, *_ = deepspeed.initialize(model=model1, config_params=ds_config)
|
|
with deepspeed.zero.Init(config_dict_or_path=ds_config):
|
|
model2 = torch.nn.Linear(4, 4)
|
|
|
|
# ensure that zero3 processed the parameter
|
|
assert hasattr(model2.weight, "ds_id")
|
|
deepspeed_engine2, *_ = deepspeed.initialize(model=model2, config_params=ds_config)
|
|
|
|
|
|
class TestNestedParallelInit(DistributedTest):
|
|
world_size = 1
|
|
|
|
# Testing a model with composed and nested zero.Inits, with 3 zero.Init contexts, 1 parent and 2 children.
|
|
# The skeleton of the model is like so
|
|
#
|
|
# class VisionEncoderDecoderModel(...)::
|
|
# def __init__(self):
|
|
# encoder = AutoModel.from_config(config.encoder)
|
|
# decoder = AutoModelForCausalLM.from_config(config.decoder)
|
|
#
|
|
# And the user calls like below:
|
|
# VisionEncoderDecoderModel.from_pretrained(...)
|
|
# which calls this constructor inside zero.Init
|
|
|
|
def test_nested_parallel_init(self, tmp_path):
|
|
ds_config = dict(train_batch_size=1, zero_optimization=dict(stage=3))
|
|
_create_tiny_vision_encoder_decoder_model(tmp_path)
|
|
dschf = HfDeepSpeedConfig(ds_config) # keep this object alive
|
|
model = VisionEncoderDecoderModel.from_pretrained(str(tmp_path), local_files_only=True)
|
|
assert all([hasattr(p, 'ds_id') for p in model.parameters()])
|