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

Function get_tensor_model_parallel_group

codegeex/megatron/mpu/initialize.py:201–206  ·  view source on GitHub ↗

Get the tensor model parallel group the caller rank belongs to.

()

Source from the content-addressed store, hash-verified

199
200
201def get_tensor_model_parallel_group():
202 """Get the tensor model parallel group the caller rank belongs to."""
203 assert (
204 _TENSOR_MODEL_PARALLEL_GROUP is not None
205 ), "intra_layer_model parallel group is not initialized"
206 return _TENSOR_MODEL_PARALLEL_GROUP
207
208
209def get_pipeline_model_parallel_group():

Callers 8

forwardMethod · 0.85
gather_split_1d_tensorFunction · 0.85
broadcast_dataFunction · 0.85
_reduceFunction · 0.85
_gatherFunction · 0.85

Calls

no outgoing calls

Tested by

no test coverage detected