48 lines
1.4 KiB
Python
48 lines
1.4 KiB
Python
|
|
import os
|
|
import torch
|
|
import torch.distributed as dist
|
|
|
|
|
|
def get_global_rank() -> int:
|
|
"""
|
|
Get the global rank, the global index of the GPU.
|
|
"""
|
|
return int(os.environ.get("RANK", "0"))
|
|
|
|
|
|
def get_local_rank() -> int:
|
|
"""
|
|
Get the local rank, the local index of the GPU.
|
|
"""
|
|
return int(os.environ.get("LOCAL_RANK", "0"))
|
|
|
|
|
|
def get_world_size() -> int:
|
|
"""
|
|
Get the world size, the total amount of GPUs.
|
|
"""
|
|
return int(os.environ.get("WORLD_SIZE", "1"))
|
|
|
|
|
|
def get_device() -> torch.device:
|
|
"""
|
|
Get current rank device.
|
|
"""
|
|
return torch.device("cuda", get_local_rank())
|
|
|
|
def get_sequence_parallel_group():
|
|
"""Get the sequence parallel group the caller rank belongs to."""
|
|
return _SEQUENCE_PARALLEL_GROUP
|
|
|
|
def initialize_sequence_parallelism(sequence_parallel_size):
|
|
assert int(get_world_size()) % sequence_parallel_size == 0
|
|
sequence_parallel_num_groups = int(get_world_size()) // sequence_parallel_size
|
|
global _SEQUENCE_PARALLEL_GROUP
|
|
for i in range(sequence_parallel_num_groups):
|
|
ranks = range(i * sequence_parallel_size,
|
|
(i + 1) * sequence_parallel_size)
|
|
group = torch.distributed.new_group(ranks)
|
|
if int(get_global_rank()) in ranks:
|
|
print(f"Rank {get_global_rank()} joined group with ranks {list(ranks)}")
|
|
_SEQUENCE_PARALLEL_GROUP = group |