MCPcopy Create free account
hub / github.com/zai-org/CodeGeeX / get_tensor_model_parallel_src_rank

Function get_tensor_model_parallel_src_rank

codegeex/megatron/mpu/initialize.py:345–350  ·  view source on GitHub ↗

Calculate the global rank corresponding to the first local rank in the tensor model parallel group.

()

Source from the content-addressed store, hash-verified

343
344
345def get_tensor_model_parallel_src_rank():
346 """Calculate the global rank corresponding to the first local rank
347 in the tensor model parallel group."""
348 global_rank = torch.distributed.get_rank()
349 local_world_size = get_tensor_model_parallel_world_size()
350 return (global_rank // local_world_size) * local_world_size
351
352
353def get_pipeline_model_parallel_first_rank():

Callers 2

broadcast_dataFunction · 0.85

Tested by

no test coverage detected