github/ColossalAI
1
0
Fork 0
mirror of https://github.com/hpcaitech/ColossalAI.git synced 2026-07-08 03:34:17 +00:00
Code Issues Packages Projects Releases Wiki Activity

[pre-commit.ci] auto fixes from pre-commit.com hooks

Browse Source 
for more information, see https://pre-commit.ci
This commit is contained in:
pre-commit-ci[bot]
2026-05-25 17:39:09 +00:00
parent ec1635477e
commit 1ee5875d61
50 changed files with 188 additions and 214 deletions
Download Patch File Download Diff File Expand all files Collapse all files

4
tests/test_checkpoint_io/test_gemini_checkpoint_io.py
View File

@@ -41,7 +41,7 @@ def exam_state_dict_with_origin(
):
from transformers import BertForSequenceClassification
(model_fn, data_gen_fn, output_transform_fn, _, _) = next(iter(model_zoo.get_sub_registry(model_name).values()))
model_fn, data_gen_fn, output_transform_fn, _, _ = next(iter(model_zoo.get_sub_registry(model_name).values()))
bert_model = model_fn()
enable_flash_attention = True if tp_size > 1 else False
@@ -101,7 +101,7 @@ def exam_state_dict(
use_async: bool,
low_cpu_mem_mode: bool,
):
(model_fn, data_gen_fn, output_transform_fn, _, _) = next(iter(model_zoo.get_sub_registry(model_name).values()))
model_fn, data_gen_fn, output_transform_fn, _, _ = next(iter(model_zoo.get_sub_registry(model_name).values()))
criterion = lambda x: x.mean()
enable_flash_attention = True if tp_size > 1 else False
enable_fused_normalization = True if tp_size > 1 else False

4
tests/test_checkpoint_io/test_gemini_torch_compability.py
View File

@@ -22,7 +22,7 @@ from tests.kit.model_zoo import model_zoo
@parameterize("shard", [False, True])
@parameterize("model_name", ["transformers_llama_for_causal_lm"])
def exam_torch_load_from_gemini(shard: bool, model_name: str):
(model_fn, data_gen_fn, output_transform_fn, _, _) = next(iter(model_zoo.get_sub_registry(model_name).values()))
model_fn, data_gen_fn, output_transform_fn, _, _ = next(iter(model_zoo.get_sub_registry(model_name).values()))
criterion = lambda x: x.mean()
plugin = GeminiPlugin(precision="fp16", initial_scale=(2**14))
booster = Booster(plugin=plugin)
@@ -88,7 +88,7 @@ def exam_torch_load_from_gemini(shard: bool, model_name: str):
@parameterize("shard", [False, True])
@parameterize("model_name", ["transformers_gpt"])
def exam_gemini_load_from_torch(shard: bool, model_name: str):
(model_fn, data_gen_fn, output_transform_fn, _, _) = next(iter(model_zoo.get_sub_registry(model_name).values()))
model_fn, data_gen_fn, output_transform_fn, _, _ = next(iter(model_zoo.get_sub_registry(model_name).values()))
criterion = lambda x: x.mean()
plugin = TorchDDPPlugin()
booster = Booster(plugin=plugin)

4
tests/test_checkpoint_io/test_hybrid_parallel_plugin_checkpoint_io.py
View File

@@ -48,9 +48,7 @@ else:
def exam_state_dict(
shard: bool, model_name: str, size_per_shard: int, test_config: dict, use_async: bool, low_cpu_mem_mode: bool
):
(model_fn, data_gen_fn, output_transform_fn, loss_fn, _) = next(
iter(model_zoo.get_sub_registry(model_name).values())
)
model_fn, data_gen_fn, output_transform_fn, loss_fn, _ = next(iter(model_zoo.get_sub_registry(model_name).values()))
criterion = loss_fn
plugin = HybridParallelPlugin(**test_config)
booster = Booster(plugin=plugin)

4
tests/test_checkpoint_io/test_plugins_huggingface_compatibility.py
View File

@@ -21,9 +21,7 @@ from tests.kit.model_zoo import model_zoo
@parameterize("model_name", ["transformers_llama_for_causal_lm"])
@parameterize("plugin_type", ["ddp", "zero", "gemini"])
def exam_from_pretrained(plugin_type: str, model_name: str, shard=True, size_per_shard=32):
(model_fn, data_gen_fn, output_transform_fn, loss_fn, _) = next(
iter(model_zoo.get_sub_registry(model_name).values())
)
model_fn, data_gen_fn, output_transform_fn, loss_fn, _ = next(iter(model_zoo.get_sub_registry(model_name).values()))
criterion = loss_fn
if plugin_type == "ddp":

8
tests/test_tensor/test_mix_gather.py
View File

@@ -13,7 +13,7 @@ from colossalai.testing import rerun_if_address_is_in_use, spawn
def check_mix_gather_S0S1(device_mesh, rank):
tensor_to_check = torch.arange(64).reshape((8, 8)).cuda()
(f, b) = (0, 1)
f, b = (0, 1)
f_target_pair = (f, [0])
b_target_pair = (b, [1])
gather_dim, logical_process_axes = mix_gather_simulator(f_target_pair, b_target_pair)
@@ -89,7 +89,7 @@ def check_two_all_gather_S0S1(device_mesh, rank):
def check_mix_gather_S1S0(device_mesh, rank):
tensor_to_check = torch.arange(64).reshape((8, 8)).cuda()
(f, b) = (0, 1)
f, b = (0, 1)
f_target_pair = (f, [1])
b_target_pair = (b, [0])
gather_dim, logical_process_axes = mix_gather_simulator(f_target_pair, b_target_pair)
@@ -165,7 +165,7 @@ def check_two_all_gather_S1S0(device_mesh, rank):
def check_mix_gather_S01R(device_mesh, rank):
tensor_to_check = torch.arange(64).reshape((8, 8)).cuda()
(f, b) = (0, 1)
f, b = (0, 1)
f_target_pair = (f, [0, 1])
b_target_pair = (b, [])
gather_dim, logical_process_axes = mix_gather_simulator(f_target_pair, b_target_pair)
@@ -231,7 +231,7 @@ def check_two_all_gather_S01R(device_mesh, rank):
def check_mix_gather_RS01(device_mesh, rank):
tensor_to_check = torch.arange(64).reshape((8, 8)).cuda()
(f, b) = (0, 1)
f, b = (0, 1)
f_target_pair = (f, [])
b_target_pair = (b, [0, 1])
gather_dim, logical_process_axes = mix_gather_simulator(f_target_pair, b_target_pair)