Fix HSDP model_rank collision when dp_shard > 1

[hancheolcho]

Summary

In DeviceMeshHandler.get_rank_info (src/l3m/helpers/dist/utils.py), when HSDP is enabled (dp_shard > 1), model_rank was multiplied by world_size (which at that point equals replicate_mesh.size()), but it should be multiplied by shard_mesh.size():

world_size = replicate_mesh.size()
model_rank = replicate_mesh.get_local_rank()
if shard_mesh.size() > 1:
    # before: stride is replicate_mesh.size() — wrong
    model_rank = model_rank * world_size + shard_mesh.get_local_rank()
    # after: stride is shard_mesh.size() — correct
    model_rank = model_rank * shard_mesh.size() + shard_mesh.get_local_rank()
    world_size *= shard_mesh.size()

Using shard_mesh.size() as the stride ensures every (replicate_rank, shard_rank) pair maps to a unique model_rank in [0, dp_replicate * dp_shard).

Example

With dp_replicate=2, dp_shard=3:

replicate_rank	shard_rank	before (× replicate_size=2)	after (× shard_size=3)
0	0	0	0
0	1	1	1
0	2	2	2
1	0	2 ❌ (collides with (0,2))	3 ✅
1	1	3	4
1	2	4	5

When replicate_size == shard_size the two expressions coincide, so the bug is only observable for non-square HSDP meshes.

Test plan

• Verified computed model_rank values are unique and contiguous over [0, dp_replicate * dp_shard) for non-square meshes.
• Run an HSDP training job with dp_shard > dp_replicate and confirm no rank-derived collisions (seeding, checkpoint sharding, logging, etc.).

🤖 Generated with Claude Code