mirror of https://github.com/hpcaitech/ColossalAI
aibig-modeldata-parallelismdeep-learningdistributed-computingfoundation-modelsheterogeneous-traininghpcinferencelarge-scalemodel-parallelismpipeline-parallelism
You can not select more than 25 topics
Topics must start with a letter or number, can include dashes ('-') and can be up to 35 characters long.
78 lines
3.2 KiB
78 lines
3.2 KiB
import pytest |
|
import torch |
|
|
|
import colossalai |
|
from colossalai.booster import Booster |
|
from colossalai.booster.plugin import LowLevelZeroPlugin |
|
from colossalai.moe.manager import MOE_MANAGER |
|
from colossalai.testing import rerun_if_address_is_in_use, spawn |
|
from colossalai.testing.random import seed_all |
|
from tests.test_moe.moe_utils import MoeModel, delete_moe_info, run_fwd_bwd, sync_local_from_ep |
|
|
|
|
|
def run_zero_test(local_rank, stage=1): |
|
criterion = torch.nn.CrossEntropyLoss() |
|
|
|
MOE_MANAGER.__init__() |
|
MOE_MANAGER.setup(parallel="EP") |
|
moe_model = MoeModel().bfloat16() |
|
moe_optimizer = torch.optim.Adam(moe_model.parameters()) |
|
moe_plugin = LowLevelZeroPlugin(stage=stage, precision="bf16") |
|
moe_booster = Booster(plugin=moe_plugin) |
|
moe_model, moe_optimizer, _, _, _ = moe_booster.boost(moe_model, moe_optimizer) |
|
|
|
MOE_MANAGER.__init__() |
|
MOE_MANAGER.setup(parallel=None) |
|
zero_model = MoeModel().bfloat16() |
|
delete_moe_info(zero_model) |
|
zero_optimizer = torch.optim.Adam(zero_model.parameters()) |
|
zero_plugin = LowLevelZeroPlugin(stage=stage, precision="bf16") |
|
zero_booster = Booster(plugin=zero_plugin) |
|
zero_model, zero_optimizer, _, _, _ = zero_booster.boost(zero_model, zero_optimizer) |
|
sync_local_from_ep(zero_model, moe_model) |
|
|
|
data = torch.randn(16, 4).bfloat16().cuda() |
|
label = torch.randint(0, 4, (16,)).cuda() |
|
|
|
zero_out = run_fwd_bwd(zero_model, data, label, criterion, zero_optimizer) |
|
moe_out = run_fwd_bwd(moe_model, data, label, criterion, moe_optimizer) |
|
assert torch.allclose(zero_out, moe_out) |
|
|
|
for (moe_name, moe_param), (zero_name, zero_param) in zip( |
|
moe_model.module.named_parameters(), zero_model.module.named_parameters() |
|
): |
|
assert moe_name == zero_name |
|
moe_grad_list = moe_optimizer._grad_store.get_partitioned_gradients_by_param_id(0, id(moe_param)) |
|
zero_grad_list = zero_optimizer._grad_store.get_partitioned_gradients_by_param_id(0, id(zero_param)) |
|
if hasattr(moe_param, "moe_info"): |
|
assert len(moe_grad_list) == 0 |
|
if stage == 1: |
|
zero_grad = zero_grad_list[local_rank].view(moe_param.grad.shape) |
|
else: |
|
zero_grad = zero_grad_list[0].view(moe_param.grad.shape) |
|
assert torch.allclose( |
|
moe_param.grad, zero_grad, atol=1e-5 |
|
), f"zero grad:\n{moe_param.grad}\ntorch grad:\n{zero_grad}\nmax diff: {(moe_param.grad - zero_grad).abs().max()}, mean diff: {(moe_param.grad - zero_grad).abs().mean()}" |
|
else: |
|
assert len(moe_grad_list) > 0 |
|
assert len(moe_grad_list) == len(zero_grad_list) |
|
for moe_grad, zero_grad in zip(moe_grad_list, zero_grad_list): |
|
assert torch.allclose(moe_grad, zero_grad) |
|
|
|
|
|
def run_dist(rank, world_size, port, stage): |
|
colossalai.launch(rank=rank, world_size=world_size, host="localhost", port=port, backend="nccl") |
|
seed_all(42 + rank) |
|
run_zero_test(rank, stage=stage) |
|
|
|
|
|
@pytest.mark.dist |
|
@pytest.mark.parametrize("world_size", [2]) |
|
@pytest.mark.parametrize("stage", [1, 2]) |
|
@rerun_if_address_is_in_use() |
|
def test_moe_zero_model(world_size, stage): |
|
spawn(run_dist, world_size, stage=stage) |
|
|
|
|
|
if __name__ == "__main__": |
|
test_moe_zero_model(world_size=2, stage=1)
|
|
|