wehub-resource-sync
121 lines
3.8 KiB
Python
121 lines
3.8 KiB
Python
# Copyright (c) Microsoft Corporation.
|
|
# SPDX-License-Identifier: Apache-2.0
|
|
|
|
# DeepSpeed Team
|
|
|
|
from typing import List
|
|
|
|
import pytest
|
|
import torch
|
|
|
|
from deepspeed.accelerator import get_accelerator
|
|
from deepspeed.inference.v2.model_implementations.flat_model_helpers import (
|
|
flatten_inference_model,
|
|
restore_inference_model,
|
|
)
|
|
from deepspeed.inference.v2.model_implementations.layer_container_base import LayerContainer
|
|
from .utils import SimpleParam, DummyInferenceModel
|
|
|
|
|
|
class TransformerLayerContainer(LayerContainer):
|
|
"""
|
|
Stub layer container
|
|
"""
|
|
PARAM_MAPPING = {
|
|
"param_1": "param_1.param",
|
|
"param_2": "param_2.param",
|
|
}
|
|
|
|
param_1: SimpleParam
|
|
|
|
param_2: SimpleParam
|
|
|
|
|
|
class NonTransformerContainer(LayerContainer):
|
|
"""
|
|
Stub layer container
|
|
"""
|
|
PARAM_MAPPING = {
|
|
"param_1": "param_1.param",
|
|
"param_2": "param_2.param",
|
|
"param_3": "param_3.param",
|
|
}
|
|
|
|
param_1: SimpleParam
|
|
|
|
param_2: SimpleParam
|
|
|
|
param_3: SimpleParam
|
|
|
|
|
|
@pytest.mark.inference_v2
|
|
def test_contiguify_roundtrip():
|
|
"""
|
|
Validate that contiguify round trips and reconstructions are correct.
|
|
"""
|
|
model = DummyInferenceModel()
|
|
|
|
n_layers = 2
|
|
transformer_params = []
|
|
transformer_containers = []
|
|
|
|
# Create parameters and populate them into the containers
|
|
for i in range(n_layers):
|
|
transformer_containers.append(TransformerLayerContainer(model))
|
|
layer_params = []
|
|
for j in range(2):
|
|
layer_params.append(torch.rand(16, 16))
|
|
transformer_containers[i].set_dependency(f"param_{j+1}", layer_params[j])
|
|
|
|
layer_params = [p.to(get_accelerator().current_device()) for p in layer_params]
|
|
|
|
transformer_params.append(layer_params)
|
|
assert transformer_containers[i].is_populated == True
|
|
|
|
non_transformer_params = []
|
|
non_transformer_container = NonTransformerContainer(model)
|
|
|
|
for i in range(3):
|
|
non_transformer_params.append(torch.rand(16, 16).permute(1, 0))
|
|
non_transformer_container.set_dependency(f"param_{i+1}", non_transformer_params[i])
|
|
|
|
non_transformer_params = [p.to(get_accelerator().current_device()) for p in non_transformer_params]
|
|
|
|
def validate_containers(t_containers: List[LayerContainer], n_t_containers: LayerContainer,
|
|
t_params: List[List[torch.Tensor]], n_t_params: List[torch.Tensor]):
|
|
"""
|
|
Validate params match what is on the containers.
|
|
"""
|
|
for i in range(n_layers):
|
|
l_c = t_containers[i]
|
|
|
|
assert l_c.is_initialized == True
|
|
|
|
assert torch.equal(l_c.param_1, t_params[i][0])
|
|
assert torch.equal(l_c.param_2, t_params[i][1])
|
|
|
|
assert n_t_containers.is_initialized == True
|
|
assert torch.equal(n_t_containers.param_1, n_t_params[0])
|
|
assert torch.equal(n_t_containers.param_2, n_t_params[1])
|
|
assert torch.equal(n_t_containers.param_3, n_t_params[2])
|
|
assert not n_t_containers.param_1.is_contiguous()
|
|
assert not n_t_containers.param_2.is_contiguous()
|
|
assert not n_t_containers.param_3.is_contiguous()
|
|
|
|
buffer, metadata = flatten_inference_model(transformer_containers, non_transformer_container, "NoOpPolicy")
|
|
|
|
# Validate containers before contiguify
|
|
validate_containers(transformer_containers, non_transformer_container, transformer_params, non_transformer_params)
|
|
|
|
# Validate restore pass
|
|
transformer_containers_r = []
|
|
for i in range(n_layers):
|
|
transformer_containers_r.append(TransformerLayerContainer(model))
|
|
|
|
non_transformer_container_r = NonTransformerContainer(model)
|
|
|
|
restore_inference_model(buffer, metadata, transformer_containers_r, non_transformer_container_r)
|
|
|
|
validate_containers(transformer_containers_r, non_transformer_container_r, transformer_params,
|
|
non_transformer_params)
|