`core.inference.shards`#

Framework-agnostic primitives for heterogeneous inference sharding: build a ProcessGroupCollection per shard, each over a contiguous rank window at its own parallelism.

Module Contents#

Classes#

InferenceShard

One shard in a multi-shard inference layout: its identity, its rank window, and this rank’s process groups for it.

Functions#

`build_inference_pg_collection`	Build a ProcessGroupCollection for one inference model.
`build_inference_pg_collections_for_shards`	Build one ProcessGroupCollection per heterogeneous inference shard.

API#

core.inference.shards.build_inference_pg_collection( world_size: int, tp_size: Optional[int] = None, pp_size: Optional[int] = None, cp_size: Optional[int] = None, ep_size: Optional[int] = None, expt_tp_size: Optional[int] = None, use_tp_pp_dp_mapping: bool = False, rank_offset: int = 0, ) → megatron.core.process_groups_config.ProcessGroupCollection#

Build a ProcessGroupCollection for one inference model.

Uses two HyperCommGrids matching mpu:

decoder_grid for dense/attention layers (tp, cp, dp, pp)
expert_grid for MoE expert layers (expt_tp, ep, expt_dp, pp)

Parameters:

world_size – Number of ranks in this inference window.
tp_size – Tensor model parallel size. Defaults to training’s TP size.
pp_size – Pipeline parallel size. Defaults to training’s PP size.
cp_size – Context parallel size. Defaults to training’s CP size.
ep_size – Expert parallel size. Defaults to training’s EP size.
expt_tp_size – Expert tensor parallel size. Defaults to training’s expert TP size.
use_tp_pp_dp_mapping – If True, use ‘tp-pp-dp’ order; otherwise ‘tp-dp-pp’.
rank_offset – Starting global rank of the window. Use 0 for collocated inference (shares ranks with training); use a non-zero offset for non-collocated setups where inference ranks are disjoint from training ranks.

Returns:

ProcessGroupCollection configured for the inference model. On ranks outside the [rank_offset, rank_offset + world_size) window every process-group field is a non-member sentinel returned by

func:: torch.distributed.new_subgroups_by_enumeration — callers should not use that instance; see
func:: build_inference_pg_collections_for_shards for the right way to filter.

class core.inference.shards.InferenceShard#

One shard in a multi-shard inference layout: its identity, its rank window, and this rank’s process groups for it.

.. attribute:: spec

This shard’s :class:~megatron.core.inference.shards_spec.InferenceShardSpec (tp/pp/ep/expt_tp/dp and optional role = prefill/decode).

.. attribute:: rank_offset

First global rank belonging to this shard.

.. attribute:: world_size

Number of ranks in this shard (tpppdp).

.. attribute:: pg_collection

The shard’s ProcessGroupCollection if the current rank is a member of this shard, else None – the is not None check is how a rank finds its own shard. Every rank still participates in the collective process-group creation for every shard (dist.new_group is world-collective); only members get a usable handle.

spec: megatron.core.inference.shards_spec.InferenceShardSpec#: None

rank_offset: int#: None

world_size: int#: None

pg_collection: Optional[megatron.core.process_groups_config.ProcessGroupCollection]#: None

core.inference.shards.build_inference_pg_collections_for_shards( total_world_size: int, shards: Union[str, Sequence[megatron.core.inference.shards_spec.InferenceShardSpec], Sequence[dict]], use_tp_pp_dp_mapping: bool = False, ) → List[core.inference.shards.InferenceShard]#

Build one ProcessGroupCollection per heterogeneous inference shard.

Partitions global ranks into contiguous non-overlapping windows, one per shard. Shard i owns ranks [offset_i, offset_i + tp_i*pp_i*dp_i).

Every rank must call this function so the collective dist.new_group calls inside :func:build_inference_pg_collection succeed for every shard. The returned pg_collection is populated only on ranks belonging to that shard; others see None.