github/ColossalAI
1
0
Fork 0
mirror of https://github.com/hpcaitech/ColossalAI.git synced 2025-09-02 01:28:31 +00:00
Code Issues Packages Projects Releases Wiki Activity

[misc] update pre-commit and run all files (#4752)

Browse Source 
* [misc] update pre-commit

* [misc] run pre-commit

* [misc] remove useless configuration files

* [misc] ignore cuda for clang-format
This commit is contained in:
Hongxin Liu
2023-09-19 14:20:26 +08:00
committed by GitHub
parent 3c6b831c26
commit 079bf3cb26
1268 changed files with 50037 additions and 38444 deletions
Download Patch File Download Diff File Expand all files Collapse all files

6
tests/test_moe/test_grad_handler.py
View File

@@ -17,11 +17,11 @@ CONFIG = dict()
def run_test(rank, world_size, port):
colossalai.launch(config=CONFIG, rank=rank, world_size=world_size, host='localhost', port=port, backend='nccl')
colossalai.launch(config=CONFIG, rank=rank, world_size=world_size, host="localhost", port=port, backend="nccl")
expert_module = nn.Linear
expert_factor = dict(in_features=DIM, out_features=DIM, device=get_current_device())
MOE_CONTEXT.setup(42) # MOE initialization
MOE_CONTEXT.setup(42) # MOE initialization
noisy_func = UniformNoiseGenerator()
router = Top1Router(noisy_func=noisy_func)
num_experts_list = [1, 2, 4]
@@ -67,5 +67,5 @@ def test_grad_handler():
spawn(run_test, 4)
if __name__ == '__main__':
if __name__ == "__main__":
test_grad_handler()

14
tests/test_moe/test_kernel.py
View File

@@ -23,12 +23,12 @@ def run_routing(rank, world_size, port, rs=2, hidden_size=128, data_type=torch.f
# Here we do not need TF32, since it brings absolute error on results
torch.backends.cuda.matmul.allow_tf32 = False
colossalai.launch(config=CONFIG, rank=rank, world_size=world_size, host='localhost', port=port, backend='nccl')
colossalai.launch(config=CONFIG, rank=rank, world_size=world_size, host="localhost", port=port, backend="nccl")
local_rank = gpc.get_local_rank(ParallelMode.GLOBAL)
MOE_CONTEXT.setup(42) # MOE environment initialization
MOE_CONTEXT.setup(42) # MOE environment initialization
MOE_CONTEXT.reset_loss()
torch.manual_seed(rs + local_rank) # set each process has different random seed
torch.manual_seed(rs + local_rank) # set each process has different random seed
# get randomized data
tokens = torch.randn(BATCH_SIZE, hidden_size, dtype=data_type, device=get_current_device(), requires_grad=True)
@@ -46,7 +46,7 @@ def run_routing(rank, world_size, port, rs=2, hidden_size=128, data_type=torch.f
old_out, _ = layer(tokens)
ech = old_out.shape
grad = torch.randn(ech, device=get_current_device())
old_out.backward(grad) # get gradient
old_out.backward(grad) # get gradient
# save all results
o_tk_grad = tokens.grad.data.clone()
@@ -57,7 +57,7 @@ def run_routing(rank, world_size, port, rs=2, hidden_size=128, data_type=torch.f
layer.gate_weight.grad.zero_()
layer.use_kernel = True
new_out, _ = layer(tokens) # get outputs through colossal kernel
new_out, _ = layer(tokens) # get outputs through colossal kernel
if data_type == torch.float32:
check_equal(old_out, new_out)
@@ -65,7 +65,7 @@ def run_routing(rank, world_size, port, rs=2, hidden_size=128, data_type=torch.f
check_equal(old_out, new_out, 1e-2)
# forward function passed
new_out.backward(grad) # get new type gradient
new_out.backward(grad) # get new type gradient
n_tk_grad = tokens.grad.data.clone()
n_gt_grad = layer.gate_weight.grad.data.clone()
@@ -92,5 +92,5 @@ def test_moe_kernel(rs, hidden_size, data_type, router):
spawn(run_routing, 4, rs=rs, hidden_size=hidden_size, data_type=data_type, router=router)
if __name__ == '__main__':
if __name__ == "__main__":
test_moe_kernel(2, 256, torch.float16, Top2Router)

10
tests/test_moe/test_moe_checkpoint.py
View File

@@ -17,11 +17,11 @@ from tests.test_zero.test_legacy.common import CONFIG
def exam_moe_checkpoint():
with ColoInitContext(device=get_current_device()):
model = MoeModel(checkpoint=True)
save_moe_model(model, 'temp_path.pth')
save_moe_model(model, "temp_path.pth")
with ColoInitContext(device=get_current_device()):
other_model = MoeModel(checkpoint=True)
load_moe_model(other_model, 'temp_path.pth')
load_moe_model(other_model, "temp_path.pth")
state_0 = model.state_dict()
state_1 = other_model.state_dict()
@@ -30,11 +30,11 @@ def exam_moe_checkpoint():
assert torch.equal(u.data, v.data)
if dist.get_rank() == 0:
os.remove('temp_path.pth')
os.remove("temp_path.pth")
def _run_dist(rank, world_size, port):
colossalai.launch(config=CONFIG, rank=rank, world_size=world_size, host='localhost', port=port, backend='nccl')
colossalai.launch(config=CONFIG, rank=rank, world_size=world_size, host="localhost", port=port, backend="nccl")
MOE_CONTEXT.setup(seed=42)
exam_moe_checkpoint()
@@ -46,5 +46,5 @@ def test_moe_checkpoint(world_size):
spawn(_run_dist)
if __name__ == '__main__':
if __name__ == "__main__":
test_moe_checkpoint(world_size=4)

11
tests/test_moe/test_moe_colo_init.py
View File

@@ -9,17 +9,16 @@ from colossalai.testing import parameterize, rerun_if_address_is_in_use, spawn
from colossalai.utils import get_current_device
from colossalai.zero import ColoInitContext
from tests.test_moe.test_moe_zero_init import MoeModel
from tests.test_tensor.common_utils import debug_print
from tests.test_zero.test_legacy.common import CONFIG
@parameterize("init_device_type", ['cpu', 'cuda'])
@parameterize("init_device_type", ["cpu", "cuda"])
def exam_moe_colo_init(init_device_type):
world_size = dist.get_world_size()
if init_device_type == 'cuda':
if init_device_type == "cuda":
init_device = get_current_device()
elif init_device_type == 'cpu':
elif init_device_type == "cpu":
init_device = torch.device("cpu")
else:
raise NotImplementedError("Unknown device found.")
@@ -40,7 +39,7 @@ def exam_moe_colo_init(init_device_type):
def _run_dist(rank, world_size, port):
colossalai.launch(config=CONFIG, rank=rank, world_size=world_size, host='localhost', port=port, backend='nccl')
colossalai.launch(config=CONFIG, rank=rank, world_size=world_size, host="localhost", port=port, backend="nccl")
MOE_CONTEXT.setup(seed=42)
exam_moe_colo_init()
@@ -52,5 +51,5 @@ def test_moe_colo_init(world_size):
spawn(_run_dist, world_size)
if __name__ == '__main__':
if __name__ == "__main__":
test_moe_colo_init(world_size=4)

6
tests/test_moe/test_moe_group.py
View File

@@ -16,11 +16,11 @@ CONFIG = dict()
def run_test(rank, world_size, port):
world_size = 4
colossalai.launch(config=CONFIG, rank=rank, world_size=world_size, host='localhost', port=port, backend='nccl')
colossalai.launch(config=CONFIG, rank=rank, world_size=world_size, host="localhost", port=port, backend="nccl")
expert_module = nn.Linear
expert_factor = dict(in_features=D_MODEL, out_features=D_FF, device=get_current_device())
MOE_CONTEXT.setup(42) # MOE environment initialization
MOE_CONTEXT.setup(42) # MOE environment initialization
exp0 = Experts(expert_module, 1, **expert_factor)
exp1 = Experts(expert_module, 2, **expert_factor)
exp2 = Experts(expert_module, 4, **expert_factor)
@@ -64,5 +64,5 @@ def test_moe_initialization():
spawn(run_test, 4)
if __name__ == '__main__':
if __name__ == "__main__":
test_moe_initialization()

46
tests/test_moe/test_moe_zero_init.py
View File

@@ -15,20 +15,15 @@ from tests.test_zero.test_legacy.common import CONFIG
class MoeModel(nn.Module):
def __init__(self, checkpoint: bool = False):
class TestSubModule(CheckpointModule):
def __init__(self):
super().__init__(checkpoint)
expert_cls = nn.Linear
expert_args_dict = dict(in_features=16, out_features=16)
self.moe = MoeModule(dim_model=16,
num_experts=8,
use_residual=True,
expert_cls=expert_cls,
**expert_args_dict)
self.moe = MoeModule(
dim_model=16, num_experts=8, use_residual=True, expert_cls=expert_cls, **expert_args_dict
)
self.proj = nn.Linear(16, 4)
def _forward(self, x):
@@ -50,49 +45,52 @@ class MoeModel(nn.Module):
return x
@parameterize("init_device_type", ['cpu', 'cuda'])
@parameterize("init_device_type", ["cpu", "cuda"])
@parameterize("shard_strategy_class", [TensorShardStrategy, BucketTensorShardStrategy])
def run_moe_zero_init(init_device_type, shard_strategy_class):
logger = get_dist_logger("test_moe_zero_init")
get_dist_logger("test_moe_zero_init")
if init_device_type == 'cuda':
if init_device_type == "cuda":
init_device = get_current_device()
elif init_device_type == 'cpu':
elif init_device_type == "cpu":
init_device = torch.device("cpu")
else:
raise NotImplementedError("Unknown device found.")
model_numel_tensor = torch.zeros(1, dtype=torch.int)
with ZeroInitContext(target_device=init_device,
shard_strategy=shard_strategy_class(),
shard_param=True,
model_numel_tensor=model_numel_tensor):
with ZeroInitContext(
target_device=init_device,
shard_strategy=shard_strategy_class(),
shard_param=True,
model_numel_tensor=model_numel_tensor,
):
model = MoeModel(checkpoint=True)
for name, param in model.named_parameters():
assert hasattr(param, 'colo_attr')
assert hasattr(param, "colo_attr")
# the parameters in moe experts and its gate should not be sharded
if ('experts' in name) or ('gate' in name) or ('residual_combine' in name):
if ("experts" in name) or ("gate" in name) or ("residual_combine" in name):
assert not param.colo_attr.sharded_data_tensor.is_sharded, "`{}` parameter has problem".format(name)
else:
assert param.colo_attr.sharded_data_tensor.is_sharded
# the parameters in moe experts is not replicated
if 'experts' in name:
if "experts" in name:
assert not param.colo_attr.is_replicated
else:
assert param.colo_attr.is_replicated
if param.colo_attr.param_is_sharded:
assert param.colo_attr.data_payload.device.type == init_device.type, \
f'{param.colo_attr.data_payload.device.type} vs. {init_device.type}'
assert (
param.colo_attr.data_payload.device.type == init_device.type
), f"{param.colo_attr.data_payload.device.type} vs. {init_device.type}"
else:
assert param.colo_attr.data_payload.device.type == 'cuda'
assert param.colo_attr.data_payload.device.type == "cuda"
def _run_dist(rank, world_size, port):
colossalai.launch(config=CONFIG, rank=rank, world_size=world_size, host='localhost', port=port, backend='nccl')
colossalai.launch(config=CONFIG, rank=rank, world_size=world_size, host="localhost", port=port, backend="nccl")
MOE_CONTEXT.setup(seed=42)
run_moe_zero_init()
@@ -104,5 +102,5 @@ def test_moe_zero_init(world_size):
spawn(_run_dist, world_size)
if __name__ == '__main__':
if __name__ == "__main__":
test_moe_zero_init(world_size=2)

12
tests/test_moe/test_moe_zero_model.py
View File

@@ -21,13 +21,13 @@ from tests.test_zero.test_legacy.common import CONFIG, check_grads_padding, run_
def run_model_test(enable_autocast, shard_strategy_class):
shard_strategy = shard_strategy_class()
get_components_func = non_distributed_component_funcs.get_callable('hanging_param_model')
get_components_func = non_distributed_component_funcs.get_callable("hanging_param_model")
_, train_dataloader, _, optimizer_class, _ = get_components_func()
criterion = MoeLoss(aux_weight=0.01, loss_fn=torch.nn.CrossEntropyLoss)
with ZeroInitContext(target_device=torch.device('cuda', torch.cuda.current_device()),
shard_strategy=shard_strategy,
shard_param=True):
with ZeroInitContext(
target_device=torch.device("cuda", torch.cuda.current_device()), shard_strategy=shard_strategy, shard_param=True
):
zero_model = MoeModel(checkpoint=True)
zero_model = ShardedModelV2(zero_model, shard_strategy)
@@ -54,7 +54,7 @@ def run_model_test(enable_autocast, shard_strategy_class):
def run_dist(rank, world_size, port):
colossalai.launch(config=CONFIG, rank=rank, world_size=world_size, host='localhost', port=port, backend='nccl')
colossalai.launch(config=CONFIG, rank=rank, world_size=world_size, host="localhost", port=port, backend="nccl")
MOE_CONTEXT.setup(seed=42)
run_model_test()
@@ -66,5 +66,5 @@ def test_moe_zero_model(world_size):
spawn(run_dist, world_size)
if __name__ == '__main__':
if __name__ == "__main__":
test_moe_zero_model(world_size=2)

43
tests/test_moe/test_moe_zero_optim.py
View File

@@ -43,31 +43,33 @@ def _run_step(model, optimizer, data, label, criterion, grad_handler):
@parameterize("cpu_offload", [True])
@parameterize("use_cpuadam", [True]) # We do not use Hybrid Adam right now, since it has a little bug
@parameterize("use_cpuadam", [True]) # We do not use Hybrid Adam right now, since it has a little bug
@parameterize("reuse_fp16_shard", [True, False])
@parameterize("shard_strategy_class", [TensorShardStrategy, BucketTensorShardStrategy])
def _run_test_sharded_optim_v2(cpu_offload,
shard_strategy_class,
use_cpuadam,
reuse_fp16_shard,
gpu_margin_mem_ratio=0.0):
def _run_test_sharded_optim_v2(
cpu_offload, shard_strategy_class, use_cpuadam, reuse_fp16_shard, gpu_margin_mem_ratio=0.0
):
shard_strategy = shard_strategy_class()
if use_cpuadam and cpu_offload is False:
return
MOE_CONTEXT.reset_loss()
get_components_func = non_distributed_component_funcs.get_callable('hanging_param_model')
get_components_func = non_distributed_component_funcs.get_callable("hanging_param_model")
_, train_dataloader, _, optimizer_class, _ = get_components_func()
criterion = MoeLoss(aux_weight=0.01, loss_fn=torch.nn.CrossEntropyLoss)
with ZeroInitContext(target_device=torch.device('cpu') if cpu_offload else get_current_device(),
shard_strategy=shard_strategy,
shard_param=True):
with ZeroInitContext(
target_device=torch.device("cpu") if cpu_offload else get_current_device(),
shard_strategy=shard_strategy,
shard_param=True,
):
zero_model = MoeModel(checkpoint=True)
zero_model = ShardedModelV2(zero_model,
shard_strategy,
tensor_placement_policy='cpu' if cpu_offload else 'cuda',
reuse_fp16_shard=reuse_fp16_shard)
zero_model = ShardedModelV2(
zero_model,
shard_strategy,
tensor_placement_policy="cpu" if cpu_offload else "cuda",
reuse_fp16_shard=reuse_fp16_shard,
)
# check whether parameters are identical in ddp
for name, p in zero_model.named_parameters():
@@ -82,12 +84,11 @@ def _run_test_sharded_optim_v2(cpu_offload,
optimizer_class = CPUAdam
optim = optimizer_class(model.parameters(), lr=1e-3)
sharded_optim = optimizer_class(zero_model.parameters(), lr=1e-3)
sharded_optim = ShardedOptimizerV2(zero_model,
sharded_optim,
initial_scale=2**5,
gpu_margin_mem_ratio=gpu_margin_mem_ratio)
sharded_optim = ShardedOptimizerV2(
zero_model, sharded_optim, initial_scale=2**5, gpu_margin_mem_ratio=gpu_margin_mem_ratio
)
amp_config = dict(opt_level='O2', keep_batchnorm_fp32=False)
amp_config = dict(opt_level="O2", keep_batchnorm_fp32=False)
apex_model, apex_optimizer = convert_to_apex_amp(model, optim, amp_config)
apex_grad_handler = MoeGradientHandler(model)
@@ -103,7 +104,7 @@ def _run_test_sharded_optim_v2(cpu_offload,
def _run_dist(rank, world_size, port):
colossalai.launch(config=CONFIG, rank=rank, world_size=world_size, host='localhost', port=port, backend='nccl')
colossalai.launch(config=CONFIG, rank=rank, world_size=world_size, host="localhost", port=port, backend="nccl")
MOE_CONTEXT.setup(seed=42)
_run_test_sharded_optim_v2()
@@ -116,5 +117,5 @@ def test_moe_zero_optim(world_size):
spawn(_run_dist, world_size)
if __name__ == '__main__':
if __name__ == "__main__":
test_moe_zero_optim(world_size=4)