github/ColossalAI
1
0
Fork 0
mirror of https://github.com/hpcaitech/ColossalAI.git synced 2025-09-12 20:54:35 +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

51
tests/test_auto_parallel/test_offload/model_utils.py
View File

@@ -1,25 +1,23 @@
import torch
import torch.nn as nn
from transformers import GPT2Config, GPT2LMHeadModel
from transformers import BertConfig, BertLMHeadModel
from transformers import BertConfig, BertLMHeadModel, GPT2Config, GPT2LMHeadModel
from tests.components_to_test.registry import non_distributed_component_funcs
class GPTLMModel(nn.Module):
def __init__(self,
hidden_size=768,
num_layers=12,
num_attention_heads=12,
max_seq_len=1024,
vocab_size=50257):
class GPTLMModel(nn.Module):
def __init__(self, hidden_size=768, num_layers=12, num_attention_heads=12, max_seq_len=1024, vocab_size=50257):
super().__init__()
self.model = GPT2LMHeadModel(
GPT2Config(n_embd=hidden_size,
n_layer=num_layers,
n_head=num_attention_heads,
n_positions=max_seq_len,
n_ctx=max_seq_len,
vocab_size=vocab_size))
GPT2Config(
n_embd=hidden_size,
n_layer=num_layers,
n_head=num_attention_heads,
n_positions=max_seq_len,
n_ctx=max_seq_len,
vocab_size=vocab_size,
)
)
def forward(self, input_ids, attention_mask):
# Only return lm_logits
@@ -27,7 +25,6 @@ class GPTLMModel(nn.Module):
class LMLoss(nn.Module):
def __init__(self):
super().__init__()
self.loss_fn = nn.CrossEntropyLoss()
@@ -38,18 +35,27 @@ class LMLoss(nn.Module):
# Flatten the tokens
return self.loss_fn(shift_logits.view(-1, shift_logits.size(-1)), shift_labels.view(-1))
class BertLMModel(nn.Module):
def __init__(self, hidden_size=768, num_layers=12, num_attention_heads=32, vocab_size=30522):
super().__init__()
self.model = BertLMHeadModel(BertConfig(n_embd=hidden_size, num_hidden_layers=num_layers, hidden_size=hidden_size,
num_attention_heads=num_attention_heads, max_position_embeddings=hidden_size,
vocab_size=vocab_size))
self.model = BertLMHeadModel(
BertConfig(
n_embd=hidden_size,
num_hidden_layers=num_layers,
hidden_size=hidden_size,
num_attention_heads=num_attention_heads,
max_position_embeddings=hidden_size,
vocab_size=vocab_size,
)
)
def forward(self, input_ids, attention_mask):
# Only return lm_logits
return self.model(input_ids=input_ids, attention_mask=attention_mask, use_cache=True)[0]
@non_distributed_component_funcs.register(name='bert_')
@non_distributed_component_funcs.register(name="bert_")
def get_bert_components():
vocab_size = 1024
seq_len = 64
@@ -67,7 +73,8 @@ def get_bert_components():
return bert_model_builder, bert_data_gen
@non_distributed_component_funcs.register(name='gpt2_')
@non_distributed_component_funcs.register(name="gpt2_")
def get_gpt2_components():
vocab_size = 1024
seq_len = 8
@@ -83,4 +90,4 @@ def get_gpt2_components():
kwargs = dict(input_ids=input_ids, attention_mask=attention_mask)
return kwargs
return gpt2_model_builder, gpt2_data_gen
return gpt2_model_builder, gpt2_data_gen

58
tests/test_auto_parallel/test_offload/test_perf.py
View File

@@ -17,18 +17,22 @@ from tests.test_auto_parallel.test_offload.model_utils import *
from tests.test_tensor.common_utils import set_seed
@parameterize('model_name', ['gpt2_'])
@parameterize('memory_budget', [5000])
@parameterize('solver_name', ['asyn'])
@parameterize("model_name", ["gpt2_"])
@parameterize("memory_budget", [5000])
@parameterize("solver_name", ["asyn"])
def exam_fwd_bwd(model_name: str, memory_budget: float, solver_name: str):
# build model
get_components_func = non_distributed_component_funcs.get_callable(model_name)
model_builder, data_gen = get_components_func()
label = torch.randint(low=0, high=128, size=(
64,
8,
), device=get_current_device())
label = torch.randint(
low=0,
high=128,
size=(
64,
8,
),
device=get_current_device(),
)
criterion = LMLoss()
set_seed(42)
@@ -50,17 +54,19 @@ def exam_fwd_bwd(model_name: str, memory_budget: float, solver_name: str):
hybrid_optimizer = HybridAdam(model.model.parameters(), lr=1e-3)
optim = AMPOptimizer(hybrid_optimizer, model)
with ColoInitContext(device=torch.device('cpu')):
with ColoInitContext(device=torch.device("cpu")):
gemini_model = model_builder()
gemini_model.train()
hybrid_optimizer = HybridAdam(gemini_model.parameters(), lr=1e-3)
gemini_config = dict(strict_ddp_mode=False,
device=torch.device('cpu'),
placement_policy='cpu',
pin_memory=True,
hidden_dim=8192,
search_range_m=128)
gemini_config = dict(
strict_ddp_mode=False,
device=torch.device("cpu"),
placement_policy="cpu",
pin_memory=True,
hidden_dim=8192,
search_range_m=128,
)
gemini_model = zero_model_wrapper(gemini_model, 3, gemini_config)
optim_config = dict(reduce_bucket_size=12 * 1024 * 1024, overlap_communication=True, verbose=True)
gemini_optim = zero_optim_wrapper(gemini_model, hybrid_optimizer, optim_config=optim_config)
@@ -89,9 +95,11 @@ def exam_fwd_bwd(model_name: str, memory_budget: float, solver_name: str):
exec_time = sum(sorted(time_list)[:5]) / 5
runtime_peak_mem_alc = torch.cuda.max_memory_allocated() / 1024**2
runtime_peak_mem_res = torch.cuda.max_memory_reserved() / 1024**2
print(f'gemini | model_name: {model_name}')
print(f'| exec_time={exec_time:.3f} s | param_size={param_size:.3f} MB '
f'| runtime_peak_mem_alc={runtime_peak_mem_alc:.3f} MB| runtime_peak_mem_res={runtime_peak_mem_res:.3f} MB|')
print(f"gemini | model_name: {model_name}")
print(
f"| exec_time={exec_time:.3f} s | param_size={param_size:.3f} MB "
f"| runtime_peak_mem_alc={runtime_peak_mem_alc:.3f} MB| runtime_peak_mem_res={runtime_peak_mem_res:.3f} MB|"
)
print(time_list)
del data_args
@@ -124,24 +132,26 @@ def exam_fwd_bwd(model_name: str, memory_budget: float, solver_name: str):
exec_time = sum(sorted(time_list)[:5]) / 5
runtime_peak_mem_alc = torch.cuda.max_memory_allocated() / 1024**2
runtime_peak_mem_res = torch.cuda.max_memory_reserved() / 1024**2
print(f'solver_name: {solver_name} | model_name: {model_name}')
print(f'| exec_time={exec_time:.3f} s | param_size={param_size:.3f} MB '
f'| runtime_peak_mem_alc={runtime_peak_mem_alc:.3f} MB| runtime_peak_mem_res={runtime_peak_mem_res:.3f} MB|')
print(f"solver_name: {solver_name} | model_name: {model_name}")
print(
f"| exec_time={exec_time:.3f} s | param_size={param_size:.3f} MB "
f"| runtime_peak_mem_alc={runtime_peak_mem_alc:.3f} MB| runtime_peak_mem_res={runtime_peak_mem_res:.3f} MB|"
)
print(time_list)
def run_dist(rank, world_size, port):
config = {}
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")
exam_fwd_bwd()
@pytest.mark.skip("this test failed")
@pytest.mark.skipif(NOT_NVML, reason='pynvml is not installed')
@pytest.mark.skipif(NOT_NVML, reason="pynvml is not installed")
@rerun_if_address_is_in_use()
def test_perf():
spawn(run_dist, 1)
if __name__ == '__main__':
if __name__ == "__main__":
test_perf()

15
tests/test_auto_parallel/test_offload/test_solver.py
View File

@@ -11,13 +11,12 @@ from colossalai.testing import clear_cache_before_run, parameterize
from tests.test_auto_parallel.test_offload.model_utils import *
@pytest.mark.skipif(NOT_NVML, reason='pynvml is not installed')
@pytest.mark.skipif(NOT_NVML, reason="pynvml is not installed")
@clear_cache_before_run()
@parameterize('model_name', ['gpt2_', 'bert_'])
@parameterize('memory_budget', [4000])
@parameterize('solver_name', ['syn', 'asyn'])
@parameterize("model_name", ["gpt2_", "bert_"])
@parameterize("memory_budget", [4000])
@parameterize("solver_name", ["syn", "asyn"])
def solver_test(model_name: str, memory_budget: float, solver_name: str):
get_components_func = non_distributed_component_funcs.get_callable(model_name)
model_builder, data_gen = get_components_func()
data_args = data_gen(device="cpu")
@@ -53,15 +52,15 @@ def solver_test(model_name: str, memory_budget: float, solver_name: str):
need_offload = region.need_offload
to_prefetch = region.fwd_prefetch_region.r_id if region.fwd_prefetch_region is not None else None
print(
f'| {model_name} forward | region id: {region.r_id} | need_offload: {need_offload} | to_prefetch: {to_prefetch}'
f"| {model_name} forward | region id: {region.r_id} | need_offload: {need_offload} | to_prefetch: {to_prefetch}"
)
for region in region_list.__reversed__():
need_offload = region.need_offload
to_prefetch = region.bwd_prefetch_region.r_id if region.bwd_prefetch_region is not None else None
print(
f'| {model_name} backward | region id: {region.r_id} | need_offload: {need_offload} | to_prefetch: {to_prefetch}'
f"| {model_name} backward | region id: {region.r_id} | need_offload: {need_offload} | to_prefetch: {to_prefetch}"
)
if __name__ == '__main__':
if __name__ == "__main__":
solver_test()