Coverage for / home / jenkins / .local / lib / python3.10 / site-packages / hyper_parallel / distributed / context_parallel / mla_backward_overlap.py: 35%
180 statements
« prev ^ index » next coverage.py v7.13.1, created at 2026-11-15 07:18 +0800
1# Copyright 2026 Huawei Technologies Co., Ltd
2#
3# Licensed under the Apache License, Version 2.0 (the "License");
4# you may not use this file except in compliance with the License.
5# You may obtain a copy of the License at
6#
7# http://www.apache.org/licenses/LICENSE-2.0
8#
9# Unless required by applicable law or agreed to in writing, software
10# distributed under the License is distributed on an "AS IS" BASIS,
11# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
12# See the License for the specific language governing permissions and
13# limitations under the License.
14# ============================================================================
15"""MLA backward scheduling with paired input A2A nodes and one output projection node.
17Paired launch/wait nodes require a complete first-order training backward.
18Partial autograd.grad calls can prune the consumer and are not supported.
19"""
21from __future__ import annotations
23from dataclasses import dataclass
24from functools import wraps
25from typing import Any
27# This CP scheduler uses the native Torch autograd and communication APIs.
28import torch # pylint: disable=forbidden-backend-import
29import torch.distributed as dist # pylint: disable=forbidden-backend-import
30from torch import nn # pylint: disable=forbidden-backend-import
31from torch.autograd.function import once_differentiable # pylint: disable=forbidden-backend-import
32from torch.nn import functional as F # pylint: disable=forbidden-backend-import
34from hyper_parallel.distributed._builder.forward_rewriter import _ForwardRewriteRequest
35from hyper_parallel.distributed.context_parallel.collectives import (
36 AsyncCPCollective, async_ulysses_seq_to_head_launch,
37)
40@dataclass
41class MLAOutputProjectionState:
42 """Output layout for ordinary MLA; its projection node owns O A2A."""
44 mesh: Any
45 head_dim: int
48class _OutputProjection(torch.autograd.Function):
49 """Own the ordinary output Linear derivative and its dO A2A."""
51 @staticmethod
52 def forward(ctx: Any, attention_output: torch.Tensor, weight: torch.Tensor,
53 bias: torch.Tensor | None, state: MLAOutputProjectionState) -> torch.Tensor:
54 """Restore the FA output to local tokens, then apply Wo.
56 Args:
57 ctx: Autograd context for operands, layout and autocast settings.
58 attention_output: FA output in [batch, sequence, local_heads, value_dim].
59 weight: Managed output-projection weight.
60 bias: Optional output-projection bias.
61 state: Ulysses output layout owned by this projection node.
63 Returns:
64 Projected local hidden states.
65 """
66 output_a2a_handle = async_ulysses_seq_to_head_launch(attention_output, 2, 1, state.mesh)
67 # Function.forward owns both the numeric A2A and its backward. The
68 # collective's internal autograd wrapper is not recorded in this scope.
69 local_output = output_a2a_handle.wait().flatten(2)
70 ctx.save_for_backward(local_output, weight)
71 ctx.state, ctx.has_bias = state, bias is not None
72 ctx.device_type = local_output.device.type
73 ctx.autocast_enabled = torch.is_autocast_enabled(ctx.device_type)
74 ctx.autocast_dtype = torch.get_autocast_dtype(ctx.device_type) if ctx.autocast_enabled else None
75 return F.linear(local_output, weight, bias) # pylint: disable=not-callable
77 @staticmethod
78 @once_differentiable
79 def backward(ctx: Any, gradient: torch.Tensor) -> tuple[torch.Tensor | None, ...]:
80 """Submit dO A2A before wgrad and wait before returning the FA gradient.
82 Args:
83 ctx: Saved operands, autocast settings and output layout.
84 gradient: Gradient of the projected hidden states.
86 Returns:
87 FA output, weight, bias and no state gradients.
88 """
89 with torch.autocast(ctx.device_type, enabled=ctx.autocast_enabled, dtype=ctx.autocast_dtype):
90 return _output_projection_backward(ctx, gradient)
93def _output_projection_backward(ctx, gradient):
94 """Keep the dO communication handle local to this backward invocation."""
95 local_output, weight = ctx.saved_tensors
96 state = ctx.state
97 needs_attention_grad, needs_weight_grad, needs_bias_grad = ctx.needs_input_grad[:3]
98 grad_output_a2a = None
99 if needs_attention_grad:
100 grad_local_output = gradient.matmul(weight).to(local_output.dtype)
101 grad_local_output = grad_local_output.reshape(*local_output.shape[:2], -1, state.head_dim)
102 grad_output_a2a = async_ulysses_seq_to_head_launch(grad_local_output, 1, 2, state.mesh)
103 flat_input = local_output.reshape(-1, local_output.shape[-1])
104 flat_gradient = gradient.reshape(-1, gradient.shape[-1])
105 grad_weight = None
106 if needs_weight_grad:
107 grad_weight = flat_input.transpose(0, 1).matmul(flat_gradient).transpose(0, 1)
108 grad_bias = flat_gradient.sum(0) if ctx.has_bias and needs_bias_grad else None
109 grad_attention_output = grad_output_a2a.wait() if grad_output_a2a is not None else None
110 return grad_attention_output, grad_weight, grad_bias, None
113def mla_output_projection_request(projection: nn.Linear) -> _ForwardRewriteRequest:
114 """Keep parameter materialization hooks while allowing dgrad/wgrad scheduling.
116 Args:
117 projection: Plain, unwrapped output Linear, before FSDP installation.
119 Returns:
120 An atomic child forward rewrite; ordinary calls remain unchanged.
121 """
122 # Quantized and LoRA subclasses need their own derivative contracts.
123 if type(projection) is not nn.Linear: # pylint: disable=unidiomatic-typecheck
124 raise ValueError("MLA backward overlap requires a plain output Linear")
125 if getattr(projection.forward, "__func__", None) is not nn.Linear.forward:
126 raise ValueError("MLA backward overlap requires an unwrapped plain output Linear")
127 original = projection.forward
129 @wraps(original)
130 def output_projection(*args: torch.Tensor, mla_output_state: MLAOutputProjectionState | None = None,
131 **kwargs: Any) -> torch.Tensor:
132 """Read managed parameters after their normal module pre-hooks.
134 Args:
135 *args: Positional arguments of the original Linear.
136 mla_output_state: Optional layout selecting the scheduled ordinary output.
137 **kwargs: Keyword arguments of the original Linear.
139 Returns:
140 Projected hidden states using ordinary or scheduled autograd.
141 """
142 if mla_output_state is None:
143 return original(*args, **kwargs)
144 # Scheduled calls consume the FA output; ordinary calls retain Linear input= support.
145 if len(args) != 1 or kwargs:
146 raise TypeError("Scheduled MLA output projection requires one positional activation")
147 return _OutputProjection.apply(args[0], projection.weight, projection.bias, mla_output_state)
149 return _ForwardRewriteRequest(projection, output_projection)
152@dataclass
153class MLAInputA2AState:
154 """Own one input A2A and its inverse for a single forward invocation.
156 Attributes:
157 mesh: Ulysses device mesh.
158 forward_a2a_handle: Q/K/V transfer from token-local to head-local layout.
159 backward_a2a_handle: Gradient transfer from head-local to token-local layout.
160 needs_send_grad: Whether the original input requires its inverse transfer.
162 Each handle owns the communication work and its buffers. The paired consumer
163 waits for it exactly once before reading the result, then clears the field.
164 """
166 mesh: Any
167 forward_a2a_handle: AsyncCPCollective | None = None
168 backward_a2a_handle: AsyncCPCollective | None = None
169 needs_send_grad: bool = False
172class _InputA2ALaunch(torch.autograd.Function):
173 """Keep independent native compute between a send and its gradient wait."""
175 @staticmethod
176 def forward(ctx: Any, send_tensor: torch.Tensor, compute_input: torch.Tensor,
177 state: MLAInputA2AState) -> torch.Tensor:
178 """Launch token-to-head A2A and attach the intervening compute edge.
180 Args:
181 ctx: Autograd context holding this invocation's A2A state.
182 send_tensor: Token-local tensor ready for communication.
183 compute_input: Input of the independent projection branch.
184 state: Forward and backward handles shared with the paired wait.
186 Returns:
187 A view of compute_input retaining both gradient dependencies.
188 """
189 ctx.state = state
190 state.needs_send_grad = ctx.needs_input_grad[0]
191 state.forward_a2a_handle = async_ulysses_seq_to_head_launch(send_tensor, 1, 2, state.mesh)
192 # Pass through the value without a copy while the paired nodes bracket
193 # the independent projection's native backward graph.
194 return compute_input.view_as(compute_input)
196 @staticmethod
197 @once_differentiable
198 def backward(ctx: Any, compute_gradient: torch.Tensor) -> tuple:
199 """Consume the inverse A2A after the independent branch backward.
201 Args:
202 ctx: Autograd context containing the paired A2A state.
203 compute_gradient: Gradient returned by the independent branch.
205 Returns:
206 Exchanged source gradient, compute gradient and no state gradient.
207 """
208 state = ctx.state
209 send_gradient = None
210 if state.needs_send_grad:
211 if state.backward_a2a_handle is None:
212 raise RuntimeError("MLA input consumer did not submit its inverse A2A")
213 send_gradient = state.backward_a2a_handle.wait()
214 state.backward_a2a_handle = None
215 return send_gradient, compute_gradient, None
218class _LatentInputA2AWait(torch.autograd.Function):
219 """Launch inverse Q before releasing the local KV projection gradients."""
221 @staticmethod
222 def forward(ctx: Any, key: torch.Tensor, value: torch.Tensor, state: MLAInputA2AState) -> tuple:
223 """Attach completed Q and both KV edges to the paired consumer.
225 Args:
226 ctx: Autograd context for the paired Q A2A.
227 key: Expanded head-local keys.
228 value: Expanded head-local values.
229 state: Q A2A submitted before the KV projection.
231 Returns:
232 Attention inputs in query, key, value order.
233 """
234 ctx.state = state
235 query = state.forward_a2a_handle.wait()
236 state.forward_a2a_handle = None
237 return query, key.view_as(key), value.view_as(value)
239 @staticmethod
240 @once_differentiable
241 def backward(ctx: Any, grad_query: torch.Tensor, grad_key: torch.Tensor,
242 grad_value: torch.Tensor) -> tuple:
243 """Run KV backward while dQ is being returned to its token owner.
245 Args:
246 ctx: Autograd context containing the paired Q A2A.
247 grad_query: Attention gradient for head-local queries.
248 grad_key: Attention gradient for head-local keys.
249 grad_value: Attention gradient for head-local values.
251 Returns:
252 Key and value gradients, followed by no state gradient.
253 """
254 state = ctx.state
255 if state.needs_send_grad:
256 state.backward_a2a_handle = async_ulysses_seq_to_head_launch(grad_query, 2, 1, state.mesh)
257 return grad_key, grad_value, None
260@dataclass
261class MLAExpandedInputA2AState:
262 """Own only the Q/K/V transfers for one expanded forward invocation."""
264 query: MLAInputA2AState
265 key: MLAInputA2AState
266 value: MLAInputA2AState
268 def wait(self) -> tuple[torch.Tensor, torch.Tensor, torch.Tensor]:
269 """Consume Q/K/V in attention order after all forward work is submitted."""
270 results = []
271 for a2a_state in (self.query, self.key, self.value):
272 results.append(a2a_state.forward_a2a_handle.wait())
273 a2a_state.forward_a2a_handle = None
274 return results[0], results[1], results[2]
276 def launch_backward(self, grad_query: torch.Tensor, grad_key: torch.Tensor,
277 grad_value: torch.Tensor) -> None:
278 """Submit each required inverse before releasing the independent compute.
280 Args:
281 grad_query: Head-local query gradient from attention.
282 grad_key: Head-local key gradient from attention.
283 grad_value: Head-local value gradient from attention.
284 """
285 for a2a_state, gradient in ((self.query, grad_query), (self.key, grad_key), (self.value, grad_value)):
286 if a2a_state.needs_send_grad:
287 a2a_state.backward_a2a_handle = async_ulysses_seq_to_head_launch(gradient, 2, 1, a2a_state.mesh)
290class _ExpandedInputA2AWait(torch.autograd.Function):
291 """Finish ordinary MLA inputs and return only Q's gradient to its projection."""
293 @staticmethod
294 def forward(ctx: Any, query: torch.Tensor, state: MLAExpandedInputA2AState) -> tuple:
295 """Submit Q after the shared V/K launches, then consume all three inputs.
297 Args:
298 ctx: Autograd context holding the three input transfers.
299 query: Final token-local query, with no following gate computation.
300 state: Per-forward Q/K/V communication state.
302 Returns:
303 Completed attention inputs in query, key, value order.
304 """
305 ctx.state = state
306 state.query.needs_send_grad = ctx.needs_input_grad[0]
307 state.query.forward_a2a_handle = async_ulysses_seq_to_head_launch(query, 1, 2, state.query.mesh)
308 return state.wait()
310 @staticmethod
311 @once_differentiable
312 def backward(ctx: Any, grad_query: torch.Tensor, grad_key: torch.Tensor,
313 grad_value: torch.Tensor) -> tuple:
314 """Submit dQ/dK/dV and release Q backward after dQ has arrived.
316 Args:
317 ctx: Autograd context holding the three paired input transfers.
318 grad_query: Attention gradient for head-local queries.
319 grad_key: Attention gradient for head-local keys.
320 grad_value: Attention gradient for head-local values.
322 Returns:
323 Token-local Q gradient and no communication-state gradient.
324 """
325 state = ctx.state
326 state.launch_backward(grad_query, grad_key, grad_value)
327 query_gradient = None
328 if state.query.needs_send_grad:
329 query_gradient = state.query.backward_a2a_handle.wait()
330 state.query.backward_a2a_handle = None
331 return query_gradient, None
334def launch_input_a2a(send_tensor: torch.Tensor, compute_input: torch.Tensor,
335 a2a_state: MLAInputA2AState) -> torch.Tensor:
336 """Submit an input A2A while retaining its native collective backward.
338 Args:
339 send_tensor: Final token-local Q, K or V ready to send.
340 compute_input: Input of the independent branch to execute next.
341 a2a_state: State receiving this forward communication handle.
343 Returns:
344 Unchanged compute_input, without adding a paired autograd dependency.
345 """
346 a2a_state.forward_a2a_handle = async_ulysses_seq_to_head_launch(send_tensor, 1, 2, a2a_state.mesh)
347 return compute_input
350def finish_expanded_inputs(query: torch.Tensor, state: MLAExpandedInputA2AState,
351 overlap_backward: bool) -> tuple[torch.Tensor, torch.Tensor, torch.Tensor]:
352 """Finish ungated expanded inputs using paired or native collective gradients.
354 Args:
355 query: Final token-local query, after V and K have been submitted.
356 state: This invocation's Q/K/V communication state.
357 overlap_backward: Pair input autograd nodes for the pure-Ulysses backward.
359 Returns:
360 Completed query, key and value tensors in head-local layout.
361 """
362 if overlap_backward:
363 return _ExpandedInputA2AWait.apply(query, state)
364 state.query.forward_a2a_handle = async_ulysses_seq_to_head_launch(query, 1, 2, state.query.mesh)
365 return state.wait()
368class _LatentKVLinear(torch.autograd.Function):
369 """Schedule the gathered latent adjoint before selected-head weight gradients."""
371 @staticmethod
372 def forward(ctx: Any, payload: torch.Tensor, weight: torch.Tensor, bias: torch.Tensor | None,
373 gathered: torch.Tensor, mesh: Any) -> tuple:
374 """Reuse the completed gather inside the managed Linear module.
376 Args:
377 ctx: Autograd state for this projection.
378 payload: Local payload supplying the single gather gradient edge.
379 weight: Selected head rows of the managed projection weight.
380 bias: Optional selected projection bias.
381 gathered: Completed full-sequence latent and shared-key payload.
382 mesh: Chronological CP mesh used by the forward gather.
384 Returns:
385 Expanded KV heads and the shared key, without duplicating its gradient.
386 """
387 del payload
388 latent = gathered[..., :weight.shape[-1]].contiguous()
389 shared_key = gathered[..., weight.shape[-1]:]
390 ctx.save_for_backward(latent, weight)
391 ctx.mesh, ctx.has_bias = mesh, bias is not None
392 ctx.device_type = latent.device.type
393 ctx.autocast_enabled = torch.is_autocast_enabled(ctx.device_type)
394 ctx.autocast_dtype = torch.get_autocast_dtype(ctx.device_type) if ctx.autocast_enabled else None
395 return F.linear(latent, weight, bias), shared_key.view_as(shared_key) # pylint: disable=not-callable
397 @staticmethod
398 @once_differentiable
399 def backward(ctx: Any, gradient: torch.Tensor, shared_gradient: torch.Tensor) -> tuple:
400 """Overlap SUM ReduceScatter with the independent KV weight gradient.
402 Args:
403 ctx: Saved projection operands, mesh and autocast state.
404 gradient: Gradient of the expanded head-local KV features.
405 shared_gradient: Gradient of the gathered shared key features.
407 Returns:
408 Local payload, weight and optional bias gradients; no gather or mesh gradient.
409 """
410 with torch.autocast(ctx.device_type, enabled=ctx.autocast_enabled, dtype=ctx.autocast_dtype):
411 return _latent_projection_backward(ctx, gradient, shared_gradient)
414def _latent_projection_backward(ctx, gradient, shared_gradient):
415 """Keep both reduction buffers alive until their local gradient is consumed."""
416 activation, weight = ctx.saved_tensors
417 needs_payload_grad, needs_weight_grad, needs_bias_grad = ctx.needs_input_grad[:3]
418 work = None
419 if needs_payload_grad:
420 grad_latent = gradient.matmul(weight).to(activation.dtype)
421 grad_payload = torch.cat((grad_latent, shared_gradient), dim=-1)
422 send = grad_payload.movedim(1, 0).contiguous()
423 shape = list(send.shape)
424 shape[0] //= ctx.mesh.size()
425 receive = send.new_empty(shape)
426 work = dist.reduce_scatter_tensor(
427 receive, send, op=dist.ReduceOp.SUM, group=ctx.mesh.get_group(), async_op=True,
428 )
429 flat_gradient = gradient.reshape(-1, gradient.shape[-1])
430 grad_weight = None
431 if needs_weight_grad:
432 flat_activation = activation.reshape(-1, activation.shape[-1])
433 grad_weight = flat_gradient.transpose(0, 1).matmul(flat_activation)
434 grad_bias = flat_gradient.sum(0) if ctx.has_bias and needs_bias_grad else None
435 local = None
436 if work is not None:
437 work.wait()
438 local = receive.movedim(0, 1).contiguous()
439 return local, grad_weight, grad_bias, None, None