[Feature] Refactor ParallelContext.world_rank_matrix

For now we're storing the global ranks inside the world_rank_matrix attribute which is a numpy array of shape (expert_parallel_size, pipeline_parallel_size, data_parallel_size, tensor_parallel_size)

So in order to access a process' global rank using the world_rank_matrix right now we're using:

nanotron/tests/test_parameters_accumulate_gradient_in_fp32.py

Lines 346 to 351 in 7c01d0f

    
           parallel_context.world_rank_matrix[ 
        
               dist.get_rank(parallel_context.expert_pg), 
        
               get_pp_rank_of(target, module=mdl), 
        
               dist.get_rank(parallel_context.dp_pg), 
        
               dist.get_rank(parallel_context.tp_pg), 
        
           ],

It would be cool to make it a functional call instead such as:

parallel_context.get_global_rank(expert_parallel_rank=0, pipeline_parallel_rank=0, data_parallel_rank=0, tensor_parallel_rank=0)

Add comprehensive tensor-parallelism support

What does this PR do?

Hi! I’d like to work on this issue if it’s still available. I can implement the "get_global_rank()" helper and add/update the relevant tests. Please let me know if I should proceed.

Navigation Menu

Repository navigation

[Feature] Refactor `ParallelContext.world_rank_matrix` #77

Metadata

Assignees

Labels

Type

Projects

Milestone

Relationships

Development

	parallel_context.world_rank_matrix[
	dist.get_rank(parallel_context.expert_pg),
	get_pp_rank_of(target, module=mdl),
	dist.get_rank(parallel_context.dp_pg),
	dist.get_rank(parallel_context.tp_pg),
	],

[Feature] Refactor ParallelContext.world_rank_matrix #77

Description

Activity

info-gallary commented on Aug 25, 2026

Metadata

Metadata

Assignees

Labels

Type

Projects

Milestone

Relationships

Development

Issue actions

[Feature] Refactor `ParallelContext.world_rank_matrix` #77