| """MPS-specific kernels that preserve the model's mathematical operations.""" |
|
|
| from __future__ import annotations |
|
|
| import math |
|
|
| import torch |
| from torch import nn |
| from transformers.integrations.moe import ALL_EXPERTS_FUNCTIONS, _grouped_linear |
| from transformers.modeling_utils import ALL_ATTENTION_FUNCTIONS |
|
|
|
|
| def attach_expert_lora( |
| experts: nn.Module, |
| *, |
| rank: int, |
| alpha: int, |
| dropout: float, |
| ) -> None: |
| """Attach packed LoRA parameters to one expert module instance.""" |
| if rank <= 0: |
| raise ValueError("Expert LoRA rank must be positive.") |
| if alpha <= 0: |
| raise ValueError("Expert LoRA alpha must be positive.") |
| if not 0.0 <= dropout < 1.0: |
| raise ValueError("Expert LoRA dropout must be in [0, 1).") |
| parameter_names = ( |
| "lora_gate_up_a", |
| "lora_gate_up_b", |
| "lora_down_a", |
| "lora_down_b", |
| ) |
| existing = tuple(hasattr(experts, name) for name in parameter_names) |
| if any(existing): |
| if not all(existing): |
| raise RuntimeError("Expert LoRA parameters are only partially initialized.") |
| return |
|
|
| experts.lora_scaling = float(alpha) / rank |
| experts.lora_dropout = nn.Dropout(dropout) if dropout else nn.Identity() |
| experts.lora_gate_up_a = nn.Parameter( |
| torch.empty( |
| experts.num_experts, |
| rank, |
| experts.hidden_dim, |
| device=experts.gate_up_proj.device, |
| dtype=experts.gate_up_proj.dtype, |
| ) |
| ) |
| experts.lora_gate_up_b = nn.Parameter( |
| torch.zeros( |
| experts.num_experts, |
| 2 * experts.intermediate_dim, |
| rank, |
| device=experts.gate_up_proj.device, |
| dtype=experts.gate_up_proj.dtype, |
| ) |
| ) |
| experts.lora_down_a = nn.Parameter( |
| torch.empty( |
| experts.num_experts, |
| rank, |
| experts.intermediate_dim, |
| device=experts.down_proj.device, |
| dtype=experts.down_proj.dtype, |
| ) |
| ) |
| experts.lora_down_b = nn.Parameter( |
| torch.zeros( |
| experts.num_experts, |
| experts.hidden_dim, |
| rank, |
| device=experts.down_proj.device, |
| dtype=experts.down_proj.dtype, |
| ) |
| ) |
| nn.init.kaiming_uniform_(experts.lora_gate_up_a, a=math.sqrt(5)) |
| nn.init.kaiming_uniform_(experts.lora_down_a, a=math.sqrt(5)) |
| experts.gate_up_proj.requires_grad_(False) |
| experts.down_proj.requires_grad_(False) |
|
|
|
|
| def lora_aware_eager_experts_forward( |
| self: nn.Module, |
| hidden_states: torch.Tensor, |
| top_k_index: torch.Tensor, |
| top_k_weights: torch.Tensor, |
| ) -> torch.Tensor: |
| final_hidden_states = torch.zeros_like(hidden_states) |
| with torch.no_grad(): |
| expert_mask = torch.nn.functional.one_hot(top_k_index, num_classes=self.num_experts) |
| expert_mask = expert_mask.permute(2, 1, 0) |
| expert_hit = torch.greater(expert_mask.sum(dim=(-1, -2)), 0).nonzero() |
|
|
| lora_dropout = getattr(self, "lora_dropout", nn.Identity()) |
| lora_scaling = getattr(self, "lora_scaling", 1.0) |
|
|
| for expert_idx in expert_hit: |
| expert_idx = expert_idx[0] |
| if expert_idx == self.num_experts: |
| continue |
| top_k_pos, token_idx = torch.where(expert_mask[expert_idx]) |
| current_state = hidden_states[token_idx] |
| gate_up = nn.functional.linear(current_state, self.gate_up_proj[expert_idx]) |
| if hasattr(self, "lora_gate_up_a"): |
| update = nn.functional.linear( |
| nn.functional.linear(lora_dropout(current_state), self.lora_gate_up_a[expert_idx]), |
| self.lora_gate_up_b[expert_idx], |
| ) |
| gate_up = gate_up + lora_scaling * update |
| gate, up = gate_up.chunk(2, dim=-1) |
| current_hidden_states = self.act_fn(gate) * up |
| expert_output = nn.functional.linear(current_hidden_states, self.down_proj[expert_idx]) |
| if hasattr(self, "lora_down_a"): |
| update = nn.functional.linear( |
| nn.functional.linear(lora_dropout(current_hidden_states), self.lora_down_a[expert_idx]), |
| self.lora_down_b[expert_idx], |
| ) |
| expert_output = expert_output + lora_scaling * update |
| current_hidden_states = expert_output * top_k_weights[token_idx, top_k_pos, None] |
| final_hidden_states.index_add_(0, token_idx, current_hidden_states.to(final_hidden_states.dtype)) |
|
|
| return final_hidden_states |
|
|
|
|
| def grouped_expert_offsets( |
| sorted_expert_ids: torch.Tensor, |
| num_experts: int, |
| ) -> torch.Tensor: |
| """Build grouped-mm offsets without MPS ``histc`` host synchronization.""" |
| counts = torch.bincount(sorted_expert_ids, minlength=num_experts)[:num_experts] |
| return torch.cumsum(counts, dim=0, dtype=torch.int32) |
|
|
|
|
| def mps_grouped_mm_experts_forward( |
| self: torch.nn.Module, |
| hidden_states: torch.Tensor, |
| top_k_index: torch.Tensor, |
| top_k_weights: torch.Tensor, |
| ) -> torch.Tensor: |
| """Transformers grouped-mm MoE with histogram-free MPS routing offsets.""" |
| num_top_k = top_k_index.size(-1) |
| num_tokens = hidden_states.size(0) |
| hidden_dim = hidden_states.size(-1) |
| sample_weights = top_k_weights.reshape(-1) |
| expert_ids = top_k_index.reshape(-1) |
|
|
| expert_ids_g, perm = torch.sort(expert_ids) |
| selected_hidden_states_g = hidden_states[perm // num_top_k] |
| sample_weights_g = sample_weights[perm] |
| offsets = grouped_expert_offsets(expert_ids_g, self.num_experts) |
|
|
| sentinel_mask = (expert_ids_g >= self.num_experts).unsqueeze(-1) |
| expert_ids_g.clamp_(max=self.num_experts - 1) |
| selected_hidden_states_g.masked_fill_(sentinel_mask, 0.0) |
|
|
| gate_up_weights = self.gate_up_proj if self.has_gate else self.up_proj |
| gate_up_biases = ( |
| self.gate_up_proj_bias[expert_ids_g] |
| if self.has_gate and self.has_bias |
| else self.up_proj_bias[expert_ids_g] |
| if self.has_bias |
| else None |
| ) |
| projected = _grouped_linear( |
| selected_hidden_states_g, |
| gate_up_weights, |
| offsets, |
| bias=gate_up_biases, |
| is_transposed=self.is_transposed, |
| ) |
| projected = self._apply_gate(projected) if self.has_gate else self.act_fn(projected) |
|
|
| down_biases = self.down_proj_bias[expert_ids_g] if self.has_bias else None |
| projected = _grouped_linear( |
| projected, |
| self.down_proj, |
| offsets, |
| bias=down_biases, |
| is_transposed=self.is_transposed, |
| ) |
| weighted = projected * sample_weights_g.unsqueeze(-1) |
| weighted.masked_fill_(sentinel_mask, 0.0) |
|
|
| |
| |
| |
| reordered = torch.empty_like(weighted) |
| reordered.index_copy_(0, perm, weighted) |
| weighted = reordered |
| return weighted.view(num_tokens, num_top_k, hidden_dim).sum(dim=1).to(hidden_states.dtype) |
|
|
|
|
| class _IndependentExpertLoRA(torch.autograd.Function): |
| """Capacity-padded reference path for independent expert LoRA.""" |
|
|
| @staticmethod |
| def forward( |
| ctx, |
| inputs: torch.Tensor, |
| weight_a: torch.Tensor, |
| weight_b: torch.Tensor, |
| expert_ids: torch.LongTensor, |
| counts: torch.LongTensor, |
| capacity: int, |
| ) -> torch.Tensor: |
| offsets = counts.cumsum(0) |
| starts = offsets - counts |
| slots = torch.arange(inputs.shape[0], device=inputs.device) - starts[expert_ids] |
| padded = inputs.new_zeros((weight_a.shape[0], capacity, inputs.shape[-1])) |
| padded[expert_ids, slots] = inputs |
| low_rank = torch.bmm(padded, weight_a.transpose(1, 2)) |
| updates = torch.bmm(low_rank, weight_b.transpose(1, 2)) |
| ctx.save_for_backward( |
| padded, |
| low_rank, |
| weight_a, |
| weight_b, |
| expert_ids, |
| slots, |
| ) |
| return updates[expert_ids, slots] |
|
|
| @staticmethod |
| def backward(ctx, grad_output: torch.Tensor): |
| padded, low_rank, weight_a, weight_b, expert_ids, slots = ctx.saved_tensors |
| grad_padded = grad_output.new_zeros( |
| (weight_a.shape[0], padded.shape[1], grad_output.shape[-1]) |
| ) |
| grad_padded[expert_ids, slots] = grad_output |
| grad_weight_b = torch.bmm(grad_padded.transpose(1, 2), low_rank) |
| grad_low_rank = torch.bmm(grad_padded, weight_b) |
| grad_weight_a = torch.bmm(grad_low_rank.transpose(1, 2), padded) |
| grad_inputs_padded = torch.bmm(grad_low_rank, weight_a) |
| grad_inputs = grad_inputs_padded[expert_ids, slots] |
| return grad_inputs, grad_weight_a, grad_weight_b, None, None, None |
|
|
|
|
| _INDEPENDENT_EXPERT_LORA_METAL_SOURCE = r""" |
| #include <metal_stdlib> |
| using namespace metal; |
| |
| constant uint kRank = 8; |
| constant uint kThreads = 256; |
| |
| kernel void expert_lora_forward_a( |
| device const bfloat* inputs [[buffer(0)]], |
| device const bfloat* weight_a [[buffer(1)]], |
| device const long* expert_ids [[buffer(2)]], |
| device bfloat* low_rank [[buffer(3)]], |
| constant uint& row_count [[buffer(4)]], |
| constant uint& input_dim [[buffer(5)]], |
| uint row [[threadgroup_position_in_grid]], |
| uint lane [[thread_index_in_threadgroup]]) { |
| if (row >= row_count) return; |
| const uint expert = uint(expert_ids[row]); |
| threadgroup float partial[kThreads * kRank]; |
| for (uint rank = 0; rank < kRank; ++rank) { |
| float value = 0.0f; |
| const uint a_base = (expert * kRank + rank) * input_dim; |
| const uint x_base = row * input_dim; |
| for (uint column = lane; column < input_dim; column += kThreads) { |
| value += float(inputs[x_base + column]) * float(weight_a[a_base + column]); |
| } |
| partial[lane * kRank + rank] = value; |
| } |
| threadgroup_barrier(mem_flags::mem_threadgroup); |
| for (uint stride = kThreads / 2; stride > 0; stride >>= 1) { |
| if (lane < stride) { |
| for (uint rank = 0; rank < kRank; ++rank) { |
| partial[lane * kRank + rank] += |
| partial[(lane + stride) * kRank + rank]; |
| } |
| } |
| threadgroup_barrier(mem_flags::mem_threadgroup); |
| } |
| if (lane == 0) { |
| for (uint rank = 0; rank < kRank; ++rank) { |
| low_rank[row * kRank + rank] = bfloat(partial[rank]); |
| } |
| } |
| } |
| |
| kernel void expert_lora_forward_b( |
| device const bfloat* low_rank [[buffer(0)]], |
| device const bfloat* weight_b [[buffer(1)]], |
| device const long* expert_ids [[buffer(2)]], |
| device bfloat* output [[buffer(3)]], |
| constant uint& row_count [[buffer(4)]], |
| constant uint& output_dim [[buffer(5)]], |
| uint index [[thread_position_in_grid]]) { |
| const uint total = row_count * output_dim; |
| if (index >= total) return; |
| const uint row = index / output_dim; |
| const uint column = index - row * output_dim; |
| const uint expert = uint(expert_ids[row]); |
| const uint b_base = (expert * output_dim + column) * kRank; |
| float value = 0.0f; |
| for (uint rank = 0; rank < kRank; ++rank) { |
| value += float(low_rank[row * kRank + rank]) |
| * float(weight_b[b_base + rank]); |
| } |
| output[index] = bfloat(value); |
| } |
| |
| kernel void expert_lora_backward_low_rank( |
| device const bfloat* grad_output [[buffer(0)]], |
| device const bfloat* weight_b [[buffer(1)]], |
| device const long* expert_ids [[buffer(2)]], |
| device bfloat* grad_low_rank [[buffer(3)]], |
| constant uint& row_count [[buffer(4)]], |
| constant uint& output_dim [[buffer(5)]], |
| uint row [[threadgroup_position_in_grid]], |
| uint lane [[thread_index_in_threadgroup]]) { |
| if (row >= row_count) return; |
| const uint expert = uint(expert_ids[row]); |
| threadgroup float partial[kThreads * kRank]; |
| for (uint rank = 0; rank < kRank; ++rank) { |
| float value = 0.0f; |
| const uint grad_base = row * output_dim; |
| for (uint column = lane; column < output_dim; column += kThreads) { |
| const uint b_index = |
| (expert * output_dim + column) * kRank + rank; |
| value += float(grad_output[grad_base + column]) |
| * float(weight_b[b_index]); |
| } |
| partial[lane * kRank + rank] = value; |
| } |
| threadgroup_barrier(mem_flags::mem_threadgroup); |
| for (uint stride = kThreads / 2; stride > 0; stride >>= 1) { |
| if (lane < stride) { |
| for (uint rank = 0; rank < kRank; ++rank) { |
| partial[lane * kRank + rank] += |
| partial[(lane + stride) * kRank + rank]; |
| } |
| } |
| threadgroup_barrier(mem_flags::mem_threadgroup); |
| } |
| if (lane == 0) { |
| for (uint rank = 0; rank < kRank; ++rank) { |
| grad_low_rank[row * kRank + rank] = bfloat(partial[rank]); |
| } |
| } |
| } |
| |
| kernel void expert_lora_backward_inputs( |
| device const bfloat* grad_low_rank [[buffer(0)]], |
| device const bfloat* weight_a [[buffer(1)]], |
| device const long* expert_ids [[buffer(2)]], |
| device bfloat* grad_inputs [[buffer(3)]], |
| constant uint& row_count [[buffer(4)]], |
| constant uint& input_dim [[buffer(5)]], |
| uint index [[thread_position_in_grid]]) { |
| const uint total = row_count * input_dim; |
| if (index >= total) return; |
| const uint row = index / input_dim; |
| const uint column = index - row * input_dim; |
| const uint expert = uint(expert_ids[row]); |
| float value = 0.0f; |
| for (uint rank = 0; rank < kRank; ++rank) { |
| const uint a_index = |
| (expert * kRank + rank) * input_dim + column; |
| value += float(grad_low_rank[row * kRank + rank]) |
| * float(weight_a[a_index]); |
| } |
| grad_inputs[index] = bfloat(value); |
| } |
| |
| kernel void expert_lora_backward_a( |
| device const bfloat* grad_low_rank [[buffer(0)]], |
| device const bfloat* inputs [[buffer(1)]], |
| device const long* offsets [[buffer(2)]], |
| device bfloat* grad_weight_a [[buffer(3)]], |
| constant uint& expert_count [[buffer(4)]], |
| constant uint& input_dim [[buffer(5)]], |
| uint index [[thread_position_in_grid]]) { |
| const uint total = expert_count * kRank * input_dim; |
| if (index >= total) return; |
| const uint column = index % input_dim; |
| const uint rank_expert = index / input_dim; |
| const uint rank = rank_expert % kRank; |
| const uint expert = rank_expert / kRank; |
| const uint start = uint(offsets[expert]); |
| const uint stop = uint(offsets[expert + 1]); |
| float value = 0.0f; |
| for (uint row = start; row < stop; ++row) { |
| value += float(grad_low_rank[row * kRank + rank]) |
| * float(inputs[row * input_dim + column]); |
| } |
| grad_weight_a[index] = bfloat(value); |
| } |
| |
| kernel void expert_lora_backward_b( |
| device const bfloat* grad_output [[buffer(0)]], |
| device const bfloat* low_rank [[buffer(1)]], |
| device const long* offsets [[buffer(2)]], |
| device bfloat* grad_weight_b [[buffer(3)]], |
| constant uint& expert_count [[buffer(4)]], |
| constant uint& output_dim [[buffer(5)]], |
| uint index [[thread_position_in_grid]]) { |
| const uint total = expert_count * output_dim * kRank; |
| if (index >= total) return; |
| const uint rank = index % kRank; |
| const uint output_expert = index / kRank; |
| const uint column = output_expert % output_dim; |
| const uint expert = output_expert / output_dim; |
| const uint start = uint(offsets[expert]); |
| const uint stop = uint(offsets[expert + 1]); |
| float value = 0.0f; |
| for (uint row = start; row < stop; ++row) { |
| value += float(grad_output[row * output_dim + column]) |
| * float(low_rank[row * kRank + rank]); |
| } |
| grad_weight_b[index] = bfloat(value); |
| } |
| """ |
|
|
|
|
| _independent_expert_lora_metal_library = None |
|
|
|
|
| def _get_independent_expert_lora_metal_library(): |
| global _independent_expert_lora_metal_library |
| if _independent_expert_lora_metal_library is None: |
| _independent_expert_lora_metal_library = torch.mps.compile_shader( |
| _INDEPENDENT_EXPERT_LORA_METAL_SOURCE |
| ) |
| return _independent_expert_lora_metal_library |
|
|
|
|
| class _MetalIndependentExpertLoRA(torch.autograd.Function): |
| """Independent rank-8 expert A/B with direct BF16 Metal forward/backward.""" |
|
|
| @staticmethod |
| def forward( |
| ctx, |
| inputs: torch.Tensor, |
| weight_a: torch.Tensor, |
| weight_b: torch.Tensor, |
| expert_ids: torch.LongTensor, |
| counts: torch.LongTensor, |
| ) -> torch.Tensor: |
| inputs = inputs.contiguous() |
| weight_a = weight_a.contiguous() |
| weight_b = weight_b.contiguous() |
| expert_ids = expert_ids.contiguous() |
| expert_count, rank, input_dim = weight_a.shape |
| row_count = inputs.shape[0] |
| output_dim = weight_b.shape[1] |
| if rank != 8: |
| raise ValueError("The direct Metal expert LoRA kernel requires rank 8.") |
| offsets = torch.cat( |
| ( |
| counts.new_zeros(1, dtype=torch.int64), |
| counts.cumsum(0, dtype=torch.int64), |
| ) |
| ).contiguous() |
| low_rank = inputs.new_empty((row_count, rank)) |
| output = inputs.new_empty((row_count, output_dim)) |
| library = _get_independent_expert_lora_metal_library() |
| library.expert_lora_forward_a( |
| inputs, |
| weight_a, |
| expert_ids, |
| low_rank, |
| row_count, |
| input_dim, |
| threads=(row_count * 256, 1, 1), |
| group_size=(256, 1, 1), |
| ) |
| library.expert_lora_forward_b( |
| low_rank, |
| weight_b, |
| expert_ids, |
| output, |
| row_count, |
| output_dim, |
| threads=(row_count * output_dim, 1, 1), |
| group_size=(256, 1, 1), |
| ) |
| ctx.save_for_backward( |
| inputs, |
| low_rank, |
| weight_a, |
| weight_b, |
| expert_ids, |
| offsets, |
| ) |
| ctx.dimensions = (expert_count, row_count, input_dim, output_dim) |
| return output |
|
|
| @staticmethod |
| def backward(ctx, grad_output: torch.Tensor): |
| inputs, low_rank, weight_a, weight_b, expert_ids, offsets = ctx.saved_tensors |
| expert_count, row_count, input_dim, output_dim = ctx.dimensions |
| grad_output = grad_output.contiguous() |
| grad_low_rank = low_rank.new_empty(low_rank.shape) |
| grad_inputs = inputs.new_empty(inputs.shape) |
| grad_weight_a = weight_a.new_empty(weight_a.shape) |
| grad_weight_b = weight_b.new_empty(weight_b.shape) |
| library = _get_independent_expert_lora_metal_library() |
| library.expert_lora_backward_low_rank( |
| grad_output, |
| weight_b, |
| expert_ids, |
| grad_low_rank, |
| row_count, |
| output_dim, |
| threads=(row_count * 256, 1, 1), |
| group_size=(256, 1, 1), |
| ) |
| library.expert_lora_backward_inputs( |
| grad_low_rank, |
| weight_a, |
| expert_ids, |
| grad_inputs, |
| row_count, |
| input_dim, |
| threads=(row_count * input_dim, 1, 1), |
| group_size=(256, 1, 1), |
| ) |
| library.expert_lora_backward_a( |
| grad_low_rank, |
| inputs, |
| offsets, |
| grad_weight_a, |
| expert_count, |
| input_dim, |
| threads=(expert_count * 8 * input_dim, 1, 1), |
| group_size=(256, 1, 1), |
| ) |
| library.expert_lora_backward_b( |
| grad_output, |
| low_rank, |
| offsets, |
| grad_weight_b, |
| expert_count, |
| output_dim, |
| threads=(expert_count * output_dim * 8, 1, 1), |
| group_size=(256, 1, 1), |
| ) |
| return grad_inputs, grad_weight_a, grad_weight_b, None, None |
|
|
|
|
| def _independent_expert_lora( |
| inputs: torch.Tensor, |
| weight_a: torch.Tensor, |
| weight_b: torch.Tensor, |
| expert_ids: torch.LongTensor, |
| counts: torch.LongTensor, |
| ) -> torch.Tensor: |
| """Dispatch the compact independent-expert operator without sharing A/B.""" |
|
|
| if inputs.shape[0] == 0: |
| return inputs.new_empty((0, weight_b.shape[1])) |
| if ( |
| inputs.device.type == "mps" |
| and inputs.dtype == torch.bfloat16 |
| and weight_a.dtype == torch.bfloat16 |
| and weight_b.dtype == torch.bfloat16 |
| and weight_a.shape[1] == 8 |
| and hasattr(torch.mps, "compile_shader") |
| ): |
| return _MetalIndependentExpertLoRA.apply( |
| inputs, |
| weight_a, |
| weight_b, |
| expert_ids, |
| counts, |
| ) |
| return _IndependentExpertLoRA.apply( |
| inputs, |
| weight_a, |
| weight_b, |
| expert_ids, |
| counts, |
| int(counts.max().detach().cpu()), |
| ) |
|
|
|
|
| def mps_segmented_experts_forward( |
| self: torch.nn.Module, |
| hidden_states: torch.Tensor, |
| top_k_index: torch.Tensor, |
| top_k_weights: torch.Tensor, |
| ) -> torch.Tensor: |
| """Run GPU-grouped base experts plus independent Metal LoRA A/B.""" |
| if hidden_states.device.type != "mps": |
| return lora_aware_eager_experts_forward( |
| self, |
| hidden_states, |
| top_k_index, |
| top_k_weights, |
| ) |
|
|
| num_top_k = top_k_index.size(-1) |
| num_tokens, hidden_dim = hidden_states.shape |
| sample_weights = top_k_weights.reshape(-1) |
| expert_ids = top_k_index.reshape(-1) |
| expert_ids_g, permutation = torch.sort(expert_ids) |
| selected_hidden_g = hidden_states[permutation // num_top_k] |
| sample_weights_g = sample_weights[permutation] |
|
|
| counts = torch.bincount( |
| expert_ids_g, |
| minlength=self.num_experts, |
| )[: self.num_experts] |
| offsets = torch.cumsum(counts, dim=0, dtype=torch.int32) |
| gate_up_weights = self.gate_up_proj if self.has_gate else self.up_proj |
| gate_up_bias = None |
| if self.has_bias: |
| packed_bias = self.gate_up_proj_bias if self.has_gate else self.up_proj_bias |
| gate_up_bias = packed_bias[expert_ids_g] |
| gate_up_g = _grouped_linear( |
| selected_hidden_g, |
| gate_up_weights, |
| offsets, |
| bias=gate_up_bias, |
| is_transposed=self.is_transposed, |
| ) |
| if hasattr(self, "lora_gate_up_a") and self.training: |
| gate_up_update = _independent_expert_lora( |
| self.lora_dropout(selected_hidden_g), |
| self.lora_gate_up_a, |
| self.lora_gate_up_b, |
| expert_ids_g, |
| counts, |
| ) |
| gate_up_g = gate_up_g + gate_up_update * self.lora_scaling |
| elif hasattr(self, "lora_gate_up_a"): |
| gate_up_update = _independent_expert_lora( |
| selected_hidden_g, |
| self.lora_gate_up_a, |
| self.lora_gate_up_b, |
| expert_ids_g, |
| counts, |
| ) |
| gate_up_g = gate_up_g + gate_up_update * self.lora_scaling |
| activated_g = ( |
| self._apply_gate(gate_up_g) |
| if self.has_gate |
| else self.act_fn(gate_up_g) |
| ) |
| down_bias = self.down_proj_bias[expert_ids_g] if self.has_bias else None |
| projected_g = _grouped_linear( |
| activated_g, |
| self.down_proj, |
| offsets, |
| bias=down_bias, |
| is_transposed=self.is_transposed, |
| ) |
| if hasattr(self, "lora_down_a") and self.training: |
| down_update = _independent_expert_lora( |
| self.lora_dropout(activated_g), |
| self.lora_down_a, |
| self.lora_down_b, |
| expert_ids_g, |
| counts, |
| ) |
| projected_g = projected_g + down_update * self.lora_scaling |
| elif hasattr(self, "lora_down_a"): |
| down_update = _independent_expert_lora( |
| activated_g, |
| self.lora_down_a, |
| self.lora_down_b, |
| expert_ids_g, |
| counts, |
| ) |
| projected_g = projected_g + down_update * self.lora_scaling |
| weighted_g = projected_g * sample_weights_g[:, None] |
| weighted = torch.empty_like(weighted_g) |
| weighted.index_copy_(0, permutation, weighted_g) |
| return weighted.view(num_tokens, num_top_k, hidden_dim).sum(dim=1).to(hidden_states.dtype) |
|
|
|
|
| _SDPA_MASK_FLOOR = -1.0e9 |
| _LONG_KV_CHUNK = 256 |
| _LONG_KV_THRESHOLD = 2048 |
| _original_sdpa_attention_forward = ALL_ATTENTION_FUNCTIONS.get_interface("sdpa", None) |
|
|
|
|
| class _MpsGqaChunkedAttention(torch.autograd.Function): |
| """Exact GQA attention that never repeats long KV onto every query head. |
| |
| Transformers' SDPA path `repeat_kv`s 2 global heads to 16 for full |
| attention. On an 8k prefix that materializes ~2GB of K and V per layer; |
| mlp-moe retains all 5 full layers, and the second denoise still holds the |
| first latent graph. MPS then returns NaN LoRA grads with a finite loss. |
| Saving compact GQA tensors and reducing K in chunks keeps the math |
| identical while staying inside unified memory. |
| """ |
|
|
| @staticmethod |
| def forward( |
| ctx, |
| query: torch.Tensor, |
| key: torch.Tensor, |
| value: torch.Tensor, |
| attn_mask: torch.Tensor | None, |
| scale: float, |
| chunk_size: int, |
| ) -> torch.Tensor: |
| query_fp32 = query.float() |
| key_fp32 = key.float() |
| value_fp32 = value.float() |
| batch, query_heads, queries, head_dim = query_fp32.shape |
| key_heads = int(key_fp32.shape[1]) |
| key_len = int(key_fp32.shape[-2]) |
| repeats = query_heads // key_heads |
| packed = batch * query_heads |
| query_flat = query_fp32.reshape(packed, queries, head_dim) |
| negative_inf = query_fp32.new_full((), _SDPA_MASK_FLOOR) |
| running_max = query_fp32.new_full((packed, queries, 1), float("-inf")) |
| running_sum = torch.zeros_like(running_max) |
| acc = torch.zeros( |
| packed, queries, head_dim, |
| device=query_fp32.device, dtype=query_fp32.dtype, |
| ) |
| chunk = max(int(chunk_size), 1) |
|
|
| def _pack_kv(tensor: torch.Tensor, start: int, stop: int) -> torch.Tensor: |
| width = stop - start |
| packed_kv = tensor[:, :, start:stop, :].unsqueeze(2).expand( |
| batch, key_heads, repeats, width, head_dim |
| ) |
| return packed_kv.reshape(packed, width, head_dim) |
|
|
| for start in range(0, key_len, chunk): |
| stop = min(start + chunk, key_len) |
| key_flat = _pack_kv(key_fp32, start, stop) |
| value_flat = _pack_kv(value_fp32, start, stop) |
| scores = torch.bmm(query_flat, key_flat.transpose(1, 2)) * scale |
| if attn_mask is not None: |
| mask_chunk = attn_mask[..., start:stop] |
| if mask_chunk.dtype == torch.bool: |
| scores = scores.masked_fill( |
| ~mask_chunk.expand(batch, query_heads, queries, stop - start).reshape( |
| packed, queries, stop - start |
| ), |
| negative_inf, |
| ) |
| else: |
| scores = scores + mask_chunk.float().expand( |
| batch, query_heads, queries, stop - start |
| ).reshape(packed, queries, stop - start) |
| block_max = scores.amax(dim=-1, keepdim=True) |
| new_max = torch.maximum(running_max, block_max) |
| weights = torch.exp(scores - new_max) |
| rescale = torch.exp(running_max - new_max) |
| running_sum = running_sum * rescale + weights.sum(dim=-1, keepdim=True) |
| acc = acc * rescale + torch.bmm(weights, value_flat) |
| running_max = new_max |
| output = acc / running_sum.clamp_min(torch.finfo(running_sum.dtype).tiny) |
| output = output.view(batch, query_heads, queries, head_dim) |
| ctx.save_for_backward( |
| query, key, value, |
| attn_mask if attn_mask is not None else query.new_empty(0), |
| ) |
| ctx.scale = float(scale) |
| ctx.chunk_size = chunk |
| ctx.has_mask = attn_mask is not None |
| return output.to(dtype=query.dtype) |
|
|
| @staticmethod |
| def backward(ctx, grad_output: torch.Tensor): |
| query, key, value, mask_or_empty = ctx.saved_tensors |
| attn_mask = mask_or_empty if ctx.has_mask else None |
| scale = ctx.scale |
| chunk = ctx.chunk_size |
| query_fp32 = query.float() |
| key_fp32 = key.float() |
| value_fp32 = value.float() |
| grad_out = grad_output.float() |
| batch, query_heads, queries, head_dim = query_fp32.shape |
| key_heads = int(key_fp32.shape[1]) |
| key_len = int(key_fp32.shape[-2]) |
| repeats = query_heads // key_heads |
| packed = batch * query_heads |
| query_flat = query_fp32.reshape(packed, queries, head_dim) |
| grad_out_flat = grad_out.reshape(packed, queries, head_dim) |
| negative_inf = query_fp32.new_full((), _SDPA_MASK_FLOOR) |
| running_max = query_fp32.new_full((packed, queries, 1), float("-inf")) |
| running_sum = torch.zeros_like(running_max) |
|
|
| def _pack_kv(tensor: torch.Tensor, start: int, stop: int) -> torch.Tensor: |
| width = stop - start |
| packed_kv = tensor[:, :, start:stop, :].unsqueeze(2).expand( |
| batch, key_heads, repeats, width, head_dim |
| ) |
| return packed_kv.reshape(packed, width, head_dim) |
|
|
| def _apply_mask(scores: torch.Tensor, start: int, stop: int) -> torch.Tensor: |
| if attn_mask is None: |
| return scores |
| mask_chunk = attn_mask[..., start:stop] |
| width = stop - start |
| if mask_chunk.dtype == torch.bool: |
| return scores.masked_fill( |
| ~mask_chunk.expand(batch, query_heads, queries, width).reshape( |
| packed, queries, width |
| ), |
| negative_inf, |
| ) |
| return scores + mask_chunk.float().expand( |
| batch, query_heads, queries, width |
| ).reshape(packed, queries, width) |
|
|
| for start in range(0, key_len, chunk): |
| stop = min(start + chunk, key_len) |
| scores = _apply_mask( |
| torch.bmm(query_flat, _pack_kv(key_fp32, start, stop).transpose(1, 2)) * scale, |
| start, |
| stop, |
| ) |
| block_max = scores.amax(dim=-1, keepdim=True) |
| new_max = torch.maximum(running_max, block_max) |
| running_sum = running_sum * torch.exp(running_max - new_max) + torch.exp( |
| scores - new_max |
| ).sum(dim=-1, keepdim=True) |
| running_max = new_max |
| grad_query_flat = torch.zeros_like(query_flat) |
| grad_key = torch.zeros_like(key_fp32) |
| grad_value = torch.zeros_like(value_fp32) |
| for start in range(0, key_len, chunk): |
| stop = min(start + chunk, key_len) |
| width = stop - start |
| key_flat = _pack_kv(key_fp32, start, stop) |
| value_flat = _pack_kv(value_fp32, start, stop) |
| scores = _apply_mask( |
| torch.bmm(query_flat, key_flat.transpose(1, 2)) * scale, |
| start, |
| stop, |
| ) |
| probs = torch.exp(scores - running_max) / running_sum.clamp_min( |
| torch.finfo(running_sum.dtype).tiny |
| ) |
| grad_value_packed = torch.bmm(probs.transpose(1, 2), grad_out_flat) |
| grad_value[:, :, start:stop, :] = grad_value_packed.view( |
| batch, key_heads, repeats, width, head_dim |
| ).sum(dim=2) |
| grad_probs = torch.bmm(grad_out_flat, value_flat.transpose(1, 2)) |
| delta = grad_probs - (grad_probs * probs).sum(dim=-1, keepdim=True) |
| grad_scores = probs * delta |
| grad_query_flat = grad_query_flat + torch.bmm(grad_scores, key_flat) * scale |
| grad_key_packed = torch.bmm(grad_scores.transpose(1, 2), query_flat) * scale |
| grad_key[:, :, start:stop, :] = grad_key_packed.view( |
| batch, key_heads, repeats, width, head_dim |
| ).sum(dim=2) |
| return ( |
| grad_query_flat.view(batch, query_heads, queries, head_dim).to(dtype=query.dtype), |
| grad_key.to(dtype=key.dtype), |
| grad_value.to(dtype=value.dtype), |
| None, |
| None, |
| None, |
| ) |
|
|
|
|
| def mps_safe_sdpa_attention_forward( |
| module: nn.Module, |
| query: torch.Tensor, |
| key: torch.Tensor, |
| value: torch.Tensor, |
| attention_mask: torch.Tensor | None, |
| dropout: float = 0.0, |
| scaling: float | None = None, |
| is_causal: bool | None = None, |
| position_bias: torch.Tensor | None = None, |
| **kwargs, |
| ) -> tuple[torch.Tensor, None]: |
| """Use compact GQA attention for long MPS decoder backwards; else official SDPA.""" |
|
|
| kv_length = int(key.shape[-2]) |
| use_chunked = ( |
| query.device.type == "mps" |
| and torch.is_grad_enabled() |
| and query.requires_grad |
| and kv_length >= _LONG_KV_THRESHOLD |
| and dropout == 0.0 |
| and position_bias is None |
| and not kwargs.get("output_attentions", False) |
| ) |
| if not use_chunked: |
| return _original_sdpa_attention_forward( |
| module, query, key, value, attention_mask, |
| dropout=dropout, scaling=scaling, is_causal=is_causal, |
| position_bias=position_bias, **kwargs, |
| ) |
| scale = 1.0 if scaling is None else float(scaling) |
| output = _MpsGqaChunkedAttention.apply( |
| query, key, value, attention_mask, scale, _LONG_KV_CHUNK, |
| ) |
| return output.transpose(1, 2).contiguous(), None |
|
|
|
|
| def register_mps_backends() -> None: |
| """Register the schema25 expert and attention implementations explicitly.""" |
|
|
| ALL_EXPERTS_FUNCTIONS.register("lora_eager", lora_aware_eager_experts_forward) |
| ALL_EXPERTS_FUNCTIONS.register("modilify_mk2_lora", inference_experts_forward) |
| ALL_EXPERTS_FUNCTIONS.register("mps_grouped_mm", mps_grouped_mm_experts_forward) |
| ALL_EXPERTS_FUNCTIONS.register("mps_segmented", mps_segmented_experts_forward) |
| ALL_ATTENTION_FUNCTIONS.register("sdpa", mps_safe_sdpa_attention_forward) |
|
|
|
|
| __all__ = [ |
| "attach_expert_lora", |
| "grouped_expert_offsets", |
| "lora_aware_eager_experts_forward", |
| "mps_grouped_mm_experts_forward", |
| "mps_safe_sdpa_attention_forward", |
| "mps_segmented_experts_forward", |
| "register_mps_backends", |
| ] |
|
|
|
|
| def inference_experts_forward(self, hidden_states, top_k_index, top_k_weights): |
| """Use the reference MPS path locally and the Torch fallback elsewhere.""" |
| implementation = (mps_segmented_experts_forward if hidden_states.device.type == "mps" |
| else lora_aware_eager_experts_forward) |
| return implementation(self, hidden_states, top_k_index, top_k_weights) |
|
|
