`bridge.peft.multi_lora_layers`#

Multi-adapter LoRA layer for Megatron parallel linears.

class:: MultiLoRALinear wraps a single Megatron parallel linear module with N concurrent LoRA adapters. The active adapter is selected at forward time via per-layer tokens_per_adapter set by :func:set_tokens_per_adapter_slot.

Forward stacks the raw weights of all adapters and uses torch._grouped_mm for a single fused kernel; TP/SP collectives are issued once around the two GEMMs to match the layout of the wrapped base linear.

class:: MultiLoRAGroupedExpertLinear is the MoE counterpart, wrapping a grouped expert linear (mlp.experts.linear_fc{1,2} of a TEGroupedMLP) with one low-rank pair per (adapter slot, local expert). Inside the experts the token order is the dispatcher’s expert-major permutation rather than the micro-batch’s adapter-major order, so tokens_per_adapter alone cannot segment it;
func:: install_moe_slot_routing co-permutes a per-token slot id through the dispatcher to recover the per-(slot, expert) segmentation.

Module Contents#

Classes#

`MultiLoRALinear`	Megatron parallel linear wrapped with N concurrent LoRA adapters.
`ExpertSlotRouting`	Per-forward mapping from dispatched expert tokens to adapter slots.
`MultiLoRAGroupedExpertLinear`	Grouped MoE expert linear wrapped with N concurrent LoRA adapters.

Functions#

`_narrow_token_counts_to_window`	Intersect contiguous per-slot token spans with the window `[start, start + num_rows)`.
`_iter_multi_lora_modules`
`set_tokens_per_adapter_slot`	Route a packed micro-batch to its per-slot token spans.
`_split_sizes_to_list`	Normalize a dispatcher split spec to the list form `all_to_all` expects.
`_co_permute_slot_ids`	Apply the dispatcher’s token permutation to a per-token adapter-slot vector.
`_build_expert_slot_routing`	Derive the per-(slot, expert) segmentation of one MoE layer’s dispatched tokens.
`_make_slot_routing_hook`	Build the forward pre-hook that publishes slot routing to one MoE layer’s adapters.
`install_moe_slot_routing`	Install per-MoE-layer slot routing for wrapped grouped expert linears.
`init_adapter_slot`	Claim slot `idx` across every multi-LoRA layer for an adapter.
`clear_adapter_slot`	Release slot `idx` across every multi-LoRA layer (zero alpha, re-init weights).
`load_adapter`	Load Megatron-shard format adapter weights into slot `idx`.
`expose_adapter_slot`	Context manager that temporarily exposes one adapter slot as `.adapter`.
`hide_adapters`	Context manager that temporarily hides all adapter params from the model.

Data#

_MULTI_LORA_TYPES

API#

class bridge.peft.multi_lora_layers.MultiLoRALinear( to_wrap: torch.nn.Module, n_adapters: int, dim: int, alpha: float, full_name: str, column_init_method: str = 'xavier', row_init_method: str = 'zero', dropout: float = 0.0, dropout_position: str = 'pre', a2a_experimental: bool = False, )#

Bases: megatron.bridge.peft.adapter_wrapper.AdapterWrapper

Megatron parallel linear wrapped with N concurrent LoRA adapters.

Each adapter slot is a :class:ParallelLinearAdapter stored in an nn.ModuleList. Forward uses grouped GEMM with a single set of TP/SP comms for efficiency.

For bridge export compatibility, use :func:expose_adapter_slot to temporarily expose one slot as .adapter.

Initialization

forward(

x: torch.Tensor,

*args: Any,

**kwargs: Any,

) → Tuple[torch.Tensor, Optional[torch.Tensor]]#

reset_adapter(idx: int) → None#

init_adapter_slot(idx: int, rank: int, alpha: float) → None#: Claim slot idx for an adapter: bind rank/alpha and apply the rank mask.

clear_adapter_slot(idx: int) → None#: Free slot idx: zero alpha, restore max rank, re-init weights.

_apply_rank_mask(idx: int) → None#

Zero padded rows of A and padded cols of B for slot idx.

For column-parallel base layers (linear_qkv, linear_fc1) linear_in.weight is sharded across TP — rank r owns global rows [r*L : (r+1)*L] where L = max_rank/tp. For row-parallel base it is replicated. Map the global cutoff actual_rank into the local shard before zeroing.

With both sides zero in the padded region, the autograd chain through the two GEMMs keeps the gradient zero there too — no periodic re-masking needed during training.

state_dict( destination: Optional[Dict[str, Any]] = None, prefix: str = '', keep_vars: bool = False, ) → Dict[str, Any]#

sharded_state_dict( prefix: str = '', sharded_offsets: Tuple[Tuple[int, int, int], ...] = (), metadata: Optional[Dict[str, Any]] = None, ) → Dict[str, Any]#

class bridge.peft.multi_lora_layers.ExpertSlotRouting#

Per-forward mapping from dispatched expert tokens to adapter slots.

Built once per MoE layer forward by :func:install_moe_slot_routing’s hook and shared by that layer’s linear_fc1/linear_fc2 adapters, which see the same rows in the same order.

.. attribute:: sort_idx

Permutation ordering the dispatched tokens by (slot, local_expert), so one grouped GEMM covers every (slot, expert) pair.

.. attribute:: inverse_idx

Inverse of sort_idx, restoring the base layer’s order.

.. attribute:: group_offsets

Inclusive cumsum of the n_adapters * num_local_experts group sizes, in slot-major order (group s * E + e).

.. attribute:: slot_token_counts

Tokens per slot, for per-token alpha/rank scaling.

.. attribute:: num_tokens

Row count the routing was built for; guards against a base layer that pads its input (e.g. fp8 quantization padding).

sort_idx: torch.Tensor#: None

inverse_idx: torch.Tensor#: None

group_offsets: torch.Tensor#: None

slot_token_counts: torch.Tensor#: None

num_tokens: int#: None

class bridge.peft.multi_lora_layers.MultiLoRAGroupedExpertLinear( to_wrap: torch.nn.Module, n_adapters: int, dim: int, alpha: float, full_name: str, num_local_experts: int, column_init_method: str = 'xavier', row_init_method: str = 'zero', dropout: float = 0.0, dropout_position: str = 'pre', )#

Bases: bridge.peft.multi_lora_layers.MultiLoRALinear

Grouped MoE expert linear wrapped with N concurrent LoRA adapters.

One :class:GroupedExpertLinearAdapter per slot, i.e. an independent low-rank pair per (slot, local expert). Reusing the single-LoRA adapter class keeps the packed [num_local_experts, ...] weight layout that the bridge’s grouped-expert export and distributed checkpointing already understand; this class only owns the multi-slot forward.

Subclassing :class:MultiLoRALinear is deliberate: the slot lifecycle helpers here and the isinstance-based multi-LoRA discovery in downstream consumers (per-slot optimizers, adapter-state zeroing) then pick expert layers up with no changes.

Unlike the dense layer, the wrapped base sits inside the experts, where rows are the dispatcher’s expert-major permutation of tokens from every EP rank — tokens_per_adapter does not segment it. The per-forward

Class:: ExpertSlotRouting supplies that segmentation instead.

Initialization

forward(

x: torch.Tensor,

*args: Any,

**kwargs: Any,

) → Tuple[torch.Tensor, Optional[torch.Tensor]]#

reset_adapter(idx: int) → None#

_apply_rank_mask(idx: int) → None#

Zero the padded rank rows of A and rank columns of B for slot idx.

Packed grouped-expert weights are [num_local_experts, rank, in] and [num_local_experts, out, rank], so the rank axis is 1 and -1 respectively, for every local expert at once. __init__ rejects expert TP, so the rank axis is never sharded and needs no local remapping (unlike the dense layer).

bridge.peft.multi_lora_layers._MULTI_LORA_TYPES#: ()

bridge.peft.multi_lora_layers._narrow_token_counts_to_window( counts: torch.Tensor, start: int, num_rows: int, ) → torch.Tensor#

Intersect contiguous per-slot token spans with the window [start, start + num_rows).

counts[i] tokens of slot i occupy the rows [cum[i-1], cum[i]) of the sequence-major flattened micro-batch. A base linear that consumes the sequence-parallel shard sees only num_rows of those rows starting at start, so its spans are the per-slot overlap with that window.

bridge.peft.multi_lora_layers._iter_multi_lora_modules(model)#

bridge.peft.multi_lora_layers.set_tokens_per_adapter_slot( model, tokens_per_adapter: torch.Tensor, ) → None#

Route a packed micro-batch to its per-slot token spans.

tokens_per_adapter[i] is the number of contiguous tokens in the upcoming forward that belong to adapter slot i. Must sum to the total token count of the micro-batch.

bridge.peft.multi_lora_layers._split_sizes_to_list(splits) → Optional[List[int]]#: Normalize a dispatcher split spec to the list form all_to_all expects.

bridge.peft.multi_lora_layers._co_permute_slot_ids( dispatcher, slot_ids: torch.Tensor, num_local_experts: int, ) → torch.Tensor#

Apply the dispatcher’s token permutation to a per-token adapter-slot vector.

Mirrors :class:MoEAlltoAllTokenDispatcher’s dispatch stages in order — local expert-major permute, EP all-to-all, then the local-expert sort — using the dispatcher’s own recorded metadata, so the result is aligned row-for-row with the permuted_local_hidden_states the experts receive.

The companion all-to-all carries one int32 per dispatched token, i.e. 2/hidden_size of the hidden-state exchange it shadows.

bridge.peft.multi_lora_layers._build_expert_slot_routing( dispatcher, tokens_per_expert: Union[torch.Tensor, Sequence[int]], tokens_per_adapter: torch.Tensor, n_adapters: int, num_local_experts: int, device: torch.device, ) → bridge.peft.multi_lora_layers.ExpertSlotRouting#

Derive the per-(slot, expert) segmentation of one MoE layer’s dispatched tokens.

The checks below run before the companion all-to-all in

Func:: _co_permute_slot_ids, so they can only fail together on all ranks — they depend on the micro-batch and on group sizes, which are uniform within a tensor-parallel group. A rank-local data corruption that tripped one of them on a single rank would hang its expert-parallel peers in that all-to-all rather than surfacing the error; every configuration-level requirement is therefore checked at layer construction instead, where failure is uniform.

bridge.peft.multi_lora_layers._make_slot_routing_hook( moe_layer: torch.nn.Module, expert_layers: List[bridge.peft.multi_lora_layers.MultiLoRAGroupedExpertLinear], )#

Build the forward pre-hook that publishes slot routing to one MoE layer’s adapters.

The routing is rebuilt from the layer’s current tokens_per_adapter, so an activation recompute must happen while that still describes the micro-batch being recomputed. That holds without pipelining (each micro-batch’s backward immediately follows its forward) — which is the only supported multi-LoRA configuration, since weight sync also requires pipeline_model_parallel_size == 1. A pipelined schedule would interleave a later micro-batch’s forward before the earlier one’s recompute and would need the counts carried on the graph instead.

bridge.peft.multi_lora_layers.install_moe_slot_routing(model) → int#

Install per-MoE-layer slot routing for wrapped grouped expert linears.

A forward pre-hook on each MoE layer’s experts module runs after the token dispatcher has permuted and exchanged tokens but before the expert GEMMs, which is the only point where both the dispatcher’s permutation metadata and the final row order are available.

Idempotent, and a no-op on models whose expert linears carry no adapters. Returns the number of MoE layers hooked.

bridge.peft.multi_lora_layers.init_adapter_slot(model, idx: int, rank: int, alpha: float) → None#

Claim slot idx across every multi-LoRA layer for an adapter.

A model-wide adapter is the set of slot-idx chunks across all layers; this initialises that set with the given rank/alpha. Thin iterator over the model — per-slot setup (rank/alpha bookkeeping + rank-mask invariant) lives on the layer itself in

Meth:: MultiLoRALinear.init_adapter_slot /

bridge.peft.multi_lora_layers.clear_adapter_slot(model, idx: int) → None#: Release slot idx across every multi-LoRA layer (zero alpha, re-init weights).

bridge.peft.multi_lora_layers.load_adapter( model, idx: int, state_dict: Dict[str, torch.Tensor], ) → int#

Load Megatron-shard format adapter weights into slot idx.

state_dict must use the Megatron-native names produced by saving while expose_adapter_slot(model, idx) is active — i.e. the same layout this function constructs to look them up. Each tensor is the local TP/PP shard, copied straight into the slot parameter with no gather, scatter, or rank-padding logic.

Saving from slot A and loading into slot B is fine because the slot index is stripped from the name (...adapter.linear_in.weight) while expose_adapter_slot is active.

Returns the number of tensors loaded (for logging / sanity checks). Raises KeyError when the checkpoint and the model’s adapter params do not match exactly in either direction (missing or unconsumed tensors).

bridge.peft.multi_lora_layers.expose_adapter_slot(model, idx: int)#

Context manager that temporarily exposes one adapter slot as .adapter.

Used by two consumers:

The bridge’s export_adapter_weights looks for .adapter.linear_in.weight (single-LoRA layout) on each wrapped module.
Megatron-native save/load walk model.named_parameters() and want names that don’t contain the slot index, so saving from slot A and loading into slot B produces matching keys.

Export contract: tensors are exported max-rank padded with .dim == max_rank, so the exposed .alpha is set to alpha * max_rank / rank — consumers computing alpha / dim apply the slot’s runtime scaling. Restored on exit.

MultiLoRALinear is handled via the common .adapters ModuleList — duck-typed rather than isinstance-checked so future multi-LoRA module types are picked up automatically.

bridge.peft.multi_lora_layers.hide_adapters(model)#

Context manager that temporarily hides all adapter params from the model.

Used during base checkpoint loading so the bridge doesn’t try to map adapter parameters to HF weights.

bridge.peft.multi_lora_layers#