def test_max_memory_used_slowmo_memory_efficient() -> None:
world_size = 2
spawn_for_all_world_sizes(
run_max_memory_used_slowmo_memory_efficient,
world_sizes=[world_size],
args=(),
deterministic=True,
)
def test_max_memory_used_slowmo_no_sharding() -> None:
world_size = 2
spawn_for_all_world_sizes(
run_max_memory_used_slowmo_no_sharding,
world_sizes=[world_size],
args=(),
deterministic=True,
)
def test_settings(test_settings) -> None:
world_size = 2
temp_file_name = tempfile.mkstemp()[1]
print("Testing ", test_settings["test_function"], " with settings ",
test_settings["test_name"])
spawn_for_all_world_sizes(
test_settings["test_function"],
world_sizes=[world_size],
args=(test_settings["slowmo_settings"], ),
deterministic=True,
)
def test_row_parallel():
spawn_for_all_world_sizes(run_test_row_parallel_linear)
def test_embedding():
spawn_for_all_world_sizes(run_test_parallel_embedding)
def test_affine_weight():
spawn_for_all_world_sizes(run_test_initialize_affine_weight)
def test_model_parallel_cuda_manual_seed():
spawn_for_all_world_sizes(run_test_model_parallel_cuda_manual_seed)
def test_cuda_rng_tracker():
spawn_for_all_world_sizes(run_test_cuda_rng_tracker)
def test_set_cuda_rng_state():
spawn_for_all_world_sizes(run_test_set_cuda_rng_state)
def test_affine_weight():
spawn_for_all_world_sizes(run_test_initialize_affine_weight,
deterministic=True)
def test_initialize_model_parallel():
spawn_for_all_world_sizes(run_test_initialize_model_parallel)
def test_row_parallel():
spawn_for_all_world_sizes(run_test_row_parallel_linear, deterministic=True)
def test_embedding():
spawn_for_all_world_sizes(run_test_parallel_embedding, deterministic=True)
def test_eight_pipe_layer():
world_sizes = [
x for x in get_world_sizes() if x <= torch.cuda.device_count() / 2
]
spawn_for_all_world_sizes(run_test_pipe, [8])
def test_get_model_parallel_src_rank():
spawn_for_all_world_sizes(run_test_get_model_parallel_src_rank)
def spawn_and_init(fn, args=None, **spawn_kwargs):
if args is None:
args = ()
run_fn = functools.partial(init_and_run, fn, args)
spawn_for_all_world_sizes(run_fn, **spawn_kwargs)
def test_cross_entropy():
spawn_for_all_world_sizes(run_test_cross_entropy)
