| |
| |
|
|
| |
|
|
| import torch |
|
|
| from typing import Optional, Any, Tuple |
| from torch import Tensor |
| from packaging import version |
| import deepspeed.comm as dist |
| from deepspeed.accelerator import get_accelerator |
|
|
| try: |
| import flash_attn |
| from flash_attn.flash_attn_interface import _flash_attn_forward, _flash_attn_backward |
| flash_attn_version = version.parse(flash_attn.__version__) |
| except ImportError: |
| _flash_attn_forward = None |
| _flash_attn_backward = None |
|
|
| from einops import rearrange |
| from .layer import single_all_to_all, apply_rotary_pos_emb |
|
|
|
|
| def _rotate_half_backward(x): |
| x = rearrange(x, '... (j d) -> ... j d', j=2) |
| x1, x2 = x.unbind(dim=-2) |
| return torch.cat((x2, -x1), dim=-1) |
|
|
|
|
| def apply_rotary_pos_emb_backward(grad_output, freqs_cos, freqs_sin): |
| rot_dim = freqs_cos.shape[-1] |
| grad, grad_pass = grad_output[..., :rot_dim], grad_output[..., rot_dim:] |
| grad_t = (grad * freqs_cos) + (_rotate_half_backward(grad * freqs_sin)) |
| grad = grad_t if grad_pass.shape[-1] == 0 else torch.cat((grad_t, grad_pass), dim=-1) |
| return grad |
|
|
|
|
| def _update_out_and_lse( |
| out: torch.Tensor, |
| lse: torch.Tensor, |
| block_out: torch.Tensor, |
| block_lse: torch.Tensor, |
| ) -> Tuple[torch.Tensor, torch.Tensor]: |
|
|
| block_out = block_out.to(torch.float32) |
| block_lse = block_lse.transpose(-2, -1).unsqueeze(dim=-1) |
|
|
| new_lse = lse + torch.log1p(torch.exp(block_lse - lse)) |
|
|
| out = torch.exp(lse - new_lse) * out + torch.exp(block_lse - new_lse) * block_out |
|
|
| lse = new_lse |
| return out, lse |
|
|
|
|
| def update_out_and_lse( |
| out: Optional[torch.Tensor], |
| lse: Optional[torch.Tensor], |
| block_out: torch.Tensor, |
| block_lse: torch.Tensor, |
| slice_=None, |
| ) -> Tuple[torch.Tensor, torch.Tensor]: |
| if out is None: |
| if slice_ is not None: |
| raise RuntimeError("first update_out_and_lse should not pass slice_ args") |
| out = block_out.to(torch.float32) |
| lse = block_lse.permute(0, 2, 1).contiguous().unsqueeze(dim=-1).contiguous() |
| elif slice_ is not None: |
| slice_out, slice_lse = out[slice_], lse[slice_] |
| slice_out, slice_lse = _update_out_and_lse(slice_out, slice_lse, block_out, block_lse) |
| out[slice_], lse[slice_] = slice_out, slice_lse |
| else: |
| out, lse = _update_out_and_lse(out, lse, block_out, block_lse) |
| return out, lse |
|
|
|
|
| class FPDT_InputConstruct(torch.nn.Module): |
|
|
| def __init__(self, tokens, labels, loss_mask, attention_mask, position_ids, args, sp_size, sp_rank) -> None: |
|
|
| super(FPDT_InputConstruct, self).__init__() |
| self.tokens = tokens |
| self.labels = labels |
| self.loss_mask = loss_mask |
| self.attention_mask = attention_mask |
| self.position_ids = position_ids |
| global_seq_len = tokens.shape[1] |
| batch_size = tokens.shape[0] |
| assert global_seq_len % sp_size == 0 |
| assert global_seq_len % args.ds_sequence_parallel_fpdt_chunk_size == 0 |
| num_chunk_per_gpu = global_seq_len // args.ds_sequence_parallel_fpdt_chunk_size |
| local_seq_len = global_seq_len // sp_size |
| assert local_seq_len % num_chunk_per_gpu == 0 |
|
|
| self.num_chunk_per_gpu = num_chunk_per_gpu |
| self.chunk_size = local_seq_len // num_chunk_per_gpu |
| self.sp_size = sp_size |
| self.sp_rank = sp_rank |
| self.global_seq_len = global_seq_len |
| self.local_seq_len = local_seq_len |
| self.batch_size = batch_size |
| self.device = tokens.device |
|
|
| def generate(self): |
| device = self.device |
| totalChunks = self.global_seq_len // self.chunk_size |
| token_chunk_idx = torch.arange(self.global_seq_len, device=device, dtype=torch.int) // self.chunk_size |
| chunk_to_gpu = torch.arange(totalChunks, device=device, dtype=torch.int) |
| chunk_to_gpu = chunk_to_gpu.reshape(self.num_chunk_per_gpu, -1).t().contiguous() |
|
|
| gather_chunk = chunk_to_gpu.flatten().unsqueeze(1).contiguous() |
| mask = gather_chunk == token_chunk_idx |
|
|
| indices = mask.nonzero(as_tuple=False) |
| gather_indices = indices[:, 0] |
| token_chunk_indices = indices[:, 1] |
| indices = torch.cat([token_chunk_indices[gather_indices == i] for i in range(gather_chunk.shape[0])]) |
| load_balanced_loss_mask = self.loss_mask[:, indices] if self.loss_mask is not None else self.loss_mask |
|
|
| indices = indices.reshape(-1, self.chunk_size)[self.num_chunk_per_gpu * self.sp_rank:self.num_chunk_per_gpu * |
| (self.sp_rank + 1)].flatten().contiguous() |
| load_balanced_tokens = self.tokens[:, indices] |
| load_balanced_labels = self.labels[:, indices] if self.labels is not None else self.labels |
|
|
| load_balanced_attention_mask = self.attention_mask if self.attention_mask is not None else self.attention_mask |
| load_balanced_position_ids = self.position_ids[:, |
| indices] if self.position_ids is not None else self.position_ids |
|
|
| return load_balanced_tokens, load_balanced_labels, load_balanced_loss_mask, load_balanced_attention_mask, load_balanced_position_ids |
|
|
|
|
| class _FPDTGPUAttentionImpl_(torch.autograd.Function): |
| generate_vmap_rule = False |
|
|
| @staticmethod |
| def forward(ctx: Any, |
| layernorm_output, |
| attention_mask, |
| inference_params, |
| rotary_pos_emb, |
| spg, |
| scatter_idx, |
| gather_idx, |
| hidden_size, |
| projection_size, |
| hidden_size_per_attention_head, |
| kv_projection_size, |
| qkv_linear_weight, |
| qkv_linear_bias, |
| dropout, |
| num_chunks=8, |
| cpu_offloading=True): |
|
|
| do_save = layernorm_output.requires_grad |
|
|
| if rotary_pos_emb is not None: |
| pos_emb_cos, pos_emb_sin = rotary_pos_emb[0].permute(1, 0, 2, 3), rotary_pos_emb[1].permute(1, 0, 2, 3) |
| ctx.pos_emb_cos = pos_emb_cos |
| ctx.pos_emb_sin = pos_emb_sin |
| else: |
| ctx.pos_emb_cos = None |
| ctx.pos_emb_sin = None |
|
|
| with torch.no_grad(): |
| per_gpu_seq_len = layernorm_output.shape[0] |
| chunk_size = per_gpu_seq_len // num_chunks |
| assert chunk_size * num_chunks == per_gpu_seq_len |
| assert attention_mask is None |
| ctx.num_chunks = num_chunks |
| ctx.cpu_offloading = cpu_offloading |
| ctx.spg = spg |
| ctx.scatter_idx = scatter_idx |
| ctx.gather_idx = gather_idx |
|
|
| device = get_accelerator().current_device_name() |
| ctx.device = device |
| ctx.dtype = layernorm_output.dtype |
| ctx.projection_size = projection_size |
| ctx.kv_projection_size = kv_projection_size |
|
|
| global_q = [] |
| global_k = [] |
| global_v = [] |
|
|
| ctx.softmax_scale = hidden_size_per_attention_head**(-0.5) |
|
|
| ctx.dropout_p = dropout |
| ctx.window_size = (-1, -1) |
| ctx.alibi_slopes = None |
|
|
| batch_size = layernorm_output.shape[1] |
|
|
| global_o = [None for _ in range(num_chunks)] |
| global_lse = [None for _ in range(num_chunks)] |
|
|
| for i in range(num_chunks): |
|
|
| st = chunk_size * i |
| ed = st + chunk_size |
|
|
| qkv_chunk = torch.matmul(layernorm_output[st:ed], qkv_linear_weight.t()) + qkv_linear_bias |
|
|
| q_chunk = qkv_chunk[:, :, :projection_size].contiguous().reshape( |
| qkv_chunk.shape[0], qkv_chunk.shape[1], -1, |
| hidden_size_per_attention_head).permute(1, 0, 2, 3).contiguous() |
| q_chunk = single_all_to_all(q_chunk, scatter_idx, gather_idx, 0, spg) |
| global_q_chunk_len = q_chunk.shape[1] |
| if rotary_pos_emb is not None: |
| q_chunk = apply_rotary_pos_emb(q_chunk, |
| pos_emb_cos[:, global_q_chunk_len * i:global_q_chunk_len * (i + 1)], |
| pos_emb_sin[:, global_q_chunk_len * i:global_q_chunk_len * (i + 1)]) |
| global_q.append(q_chunk) |
|
|
| k_chunk = qkv_chunk[:, :, projection_size:projection_size + kv_projection_size].contiguous().reshape( |
| qkv_chunk.shape[0], qkv_chunk.shape[1], -1, |
| hidden_size_per_attention_head).permute(1, 0, 2, 3).contiguous() |
| k_chunk = single_all_to_all(k_chunk, scatter_idx, gather_idx, 0, spg) |
| if rotary_pos_emb is not None: |
| k_chunk = apply_rotary_pos_emb(k_chunk, |
| pos_emb_cos[:, global_q_chunk_len * i:global_q_chunk_len * (i + 1)], |
| pos_emb_sin[:, global_q_chunk_len * i:global_q_chunk_len * (i + 1)]) |
| global_k.append(k_chunk) |
|
|
| v_chunk = qkv_chunk[:, :, projection_size + kv_projection_size:].contiguous().reshape( |
| qkv_chunk.shape[0], qkv_chunk.shape[1], -1, |
| hidden_size_per_attention_head).permute(1, 0, 2, 3).contiguous() |
| v_chunk = single_all_to_all(v_chunk, scatter_idx, gather_idx, 0, spg) |
| global_v.append(v_chunk) |
|
|
| for k_i in range(len(global_k)): |
| causal_chunk = i == k_i |
| if flash_attn_version >= version.parse("2.6.0"): |
| block_out, _, _, _, _, block_lse, _, _ = _flash_attn_forward(global_q[i], |
| global_k[k_i], |
| global_v[k_i], |
| ctx.dropout_p, |
| ctx.softmax_scale, |
| causal=causal_chunk, |
| window_size=ctx.window_size, |
| softcap=0.0, |
| alibi_slopes=ctx.alibi_slopes, |
| return_softmax=False) |
| else: |
| block_out, _, _, _, _, block_lse, _, _ = _flash_attn_forward(global_q[i], |
| global_k[k_i], |
| global_v[k_i], |
| ctx.dropout_p, |
| ctx.softmax_scale, |
| causal=causal_chunk, |
| window_size=ctx.window_size, |
| alibi_slopes=ctx.alibi_slopes, |
| return_softmax=False) |
|
|
| global_o[i], global_lse[i] = update_out_and_lse(global_o[i], global_lse[i], block_out, block_lse) |
|
|
| global_o[i] = global_o[i].to(q_chunk.dtype) |
|
|
| output = [None for i in range(num_chunks)] |
|
|
| for i in range(num_chunks): |
| global_lse[i] = global_lse[i][:, :, :, 0].permute(0, 2, 1).contiguous() |
| output[i] = single_all_to_all(global_o[i].to(ctx.dtype).contiguous(), gather_idx, scatter_idx, 0, spg) |
| output = torch.cat(output, dim=1) |
|
|
| head_dim = output.shape[-1] |
|
|
| if do_save: |
| ctx.save_for_backward(layernorm_output) |
| ctx.global_q = global_q |
| ctx.global_k = global_k |
| ctx.global_v = global_v |
| ctx.attn_output = global_o |
| ctx.attn_lse = global_lse |
| ctx.head_dim = head_dim |
| ctx.batch_size = batch_size |
|
|
| ctx.qkv_linear_weight = qkv_linear_weight |
| ctx.qkv_linear_bias = qkv_linear_bias |
|
|
| return output |
|
|
| @staticmethod |
| def backward(ctx, grad_output): |
|
|
| num_chunks = ctx.num_chunks |
| device = ctx.device |
| dtype = ctx.dtype |
| spg = ctx.spg |
| scatter_idx = ctx.scatter_idx |
| gather_idx = ctx.gather_idx |
| softmax_scale = ctx.softmax_scale |
| dropout_p = ctx.dropout_p |
| window_size = ctx.window_size |
| alibi_slopes = ctx.alibi_slopes |
|
|
| projection_size = ctx.projection_size |
| kv_projection_size = ctx.kv_projection_size |
|
|
| layernorm_output = ctx.saved_tensors[0] |
|
|
| global_q = ctx.global_q |
| global_k = ctx.global_k |
| global_v = ctx.global_v |
| attn_output = ctx.attn_output |
| lse = ctx.attn_lse |
|
|
| qkv_linear_weight = ctx.qkv_linear_weight |
| qkv_linear_bias = ctx.qkv_linear_bias |
|
|
| input_chunk_size = layernorm_output.shape[0] // num_chunks |
| grad_layernorm_output = [ |
| torch.zeros((input_chunk_size, layernorm_output.shape[1], layernorm_output.shape[2]), |
| device=device, |
| dtype=dtype) for _ in range(num_chunks) |
| ] |
|
|
| grad_global_attn_output = [] |
| chunk_size = grad_output.shape[1] // num_chunks |
|
|
| for i in range(num_chunks): |
| st = chunk_size * i |
| ed = st + chunk_size |
| grad_global_attn_output.append( |
| single_all_to_all(grad_output[:, st:ed].contiguous(), scatter_idx, gather_idx, 0, spg)) |
|
|
| del grad_output |
|
|
| dq = [torch.zeros(global_q[0].shape, dtype=torch.float, device=device) for _ in range(num_chunks)] |
| dk = [torch.zeros(global_q[0].shape, dtype=torch.float, device=device) for _ in range(num_chunks)] |
| dv = [torch.zeros(global_q[0].shape, dtype=torch.float, device=device) for _ in range(num_chunks)] |
|
|
| grad_qkv_linear_weight = torch.zeros(qkv_linear_weight.shape, |
| device=qkv_linear_weight.device, |
| dtype=torch.float) |
| grad_qkv_linear_bias = torch.zeros(qkv_linear_bias.shape, device=qkv_linear_weight.device, dtype=torch.float) |
|
|
| for i in range(num_chunks): |
| k_chunk = global_k[i] |
| v_chunk = global_v[i] |
|
|
| for q_i in range(num_chunks): |
| no_computation = q_i < i |
| if no_computation: |
| continue |
|
|
| causal_chunk = q_i == i |
|
|
| q_chunk = global_q[q_i] |
| attn_output_chunk = attn_output[q_i] |
| lse_chunk = lse[q_i] |
| d_out = grad_global_attn_output[q_i] |
|
|
| dq_this = torch.zeros(global_q[0].shape, dtype=dtype, device=device) |
| dk_this = torch.zeros(global_k[0].shape, dtype=dtype, device=device) |
| dv_this = torch.zeros(global_v[0].shape, dtype=dtype, device=device) |
|
|
| if flash_attn_version >= version.parse("2.6.0"): |
| _flash_attn_backward(d_out, |
| q_chunk, |
| k_chunk, |
| v_chunk, |
| attn_output_chunk, |
| lse_chunk, |
| dq_this, |
| dk_this, |
| dv_this, |
| dropout_p, |
| softmax_scale, |
| causal_chunk, |
| window_size, |
| softcap=0.0, |
| alibi_slopes=alibi_slopes, |
| deterministic=False, |
| rng_state=None) |
| else: |
| _flash_attn_backward(d_out, |
| q_chunk, |
| k_chunk, |
| v_chunk, |
| attn_output_chunk, |
| lse_chunk, |
| dq_this, |
| dk_this, |
| dv_this, |
| dropout_p, |
| softmax_scale, |
| causal_chunk, |
| window_size, |
| alibi_slopes=alibi_slopes, |
| deterministic=False, |
| rng_state=None) |
|
|
| dq[q_i].add_(dq_this.to(torch.float)) |
| dk[i].add_(dk_this.to(torch.float)) |
| dv[i].add_(dv_this.to(torch.float)) |
|
|
| dk_seq_len = dk[i].shape[1] |
|
|
| if ctx.pos_emb_cos is not None: |
| dk[i] = apply_rotary_pos_emb_backward(dk[i].to(dtype), |
| ctx.pos_emb_cos[:, dk_seq_len * i:dk_seq_len * (i + 1)], |
| ctx.pos_emb_sin[:, dk_seq_len * i:dk_seq_len * (i + 1)]) |
| else: |
| dk[i] = dk[i].to(dtype) |
| dv[i] = dv[i].to(dtype) |
| dk[i] = single_all_to_all(dk[i].contiguous(), gather_idx, scatter_idx, 0, spg) |
| dv[i] = single_all_to_all(dv[i].contiguous(), gather_idx, scatter_idx, 0, spg) |
|
|
| input_st = i * input_chunk_size |
| input_ed = input_st + input_chunk_size |
|
|
| input_chunk = layernorm_output[input_st:input_ed].reshape(-1, layernorm_output.shape[-1]) |
|
|
| dk[i] = dk[i].flatten(2).permute(1, 0, 2) |
| dv[i] = dv[i].flatten(2).permute(1, 0, 2) |
| l, b = dk[i].shape[0], dk[i].shape[1] |
| grad_qkv_linear_weight[projection_size:projection_size + kv_projection_size].add_( |
| torch.matmul(dk[i].reshape(l * b, -1).t(), input_chunk)) |
| grad_qkv_linear_weight[projection_size + kv_projection_size:].add_( |
| torch.matmul(dv[i].reshape(l * b, -1).t(), input_chunk)) |
| grad_qkv_linear_bias[projection_size:projection_size + kv_projection_size].add_(dk[i].sum(0).sum(0)) |
| grad_qkv_linear_bias[projection_size + kv_projection_size:].add_(dv[i].sum(0).sum(0)) |
|
|
| grad_layernorm_output[i].add_( |
| torch.matmul(dk[i], qkv_linear_weight[projection_size:projection_size + kv_projection_size])) |
| grad_layernorm_output[i].add_(torch.matmul(dv[i], |
| qkv_linear_weight[projection_size + kv_projection_size:])) |
|
|
| dk[i] = None |
| dv[i] = None |
|
|
| for i in range(num_chunks): |
| dq_seq_len = dq[i].shape[1] |
| if ctx.pos_emb_cos is not None: |
| dq[i] = apply_rotary_pos_emb_backward(dq[i].to(dtype), |
| ctx.pos_emb_cos[:, dq_seq_len * i:dq_seq_len * (i + 1)], |
| ctx.pos_emb_sin[:, dq_seq_len * i:dq_seq_len * (i + 1)]) |
| else: |
| dq[i] = dq[i].to(dtype) |
| dq[i] = single_all_to_all(dq[i].to(dtype).contiguous(), gather_idx, scatter_idx, 0, spg) |
|
|
| input_chunk = layernorm_output[:input_chunk_size].reshape(-1, layernorm_output.shape[-1]) |
| layernorm_output = layernorm_output[input_chunk_size:] |
|
|
| dq[i] = dq[i].flatten(2).permute(1, 0, 2) |
| l, b = dq[i].shape[0], dq[i].shape[1] |
| grad_qkv_linear_weight[:projection_size].add_(torch.matmul(dq[i].reshape(l * b, -1).t(), input_chunk)) |
| grad_qkv_linear_bias[:projection_size].add_(dq[i].sum(0).sum(0)) |
|
|
| grad_layernorm_output[i].add_(torch.matmul(dq[i], qkv_linear_weight[:projection_size])) |
|
|
| dq[i] = None |
|
|
| return torch.cat( |
| grad_layernorm_output, |
| dim=0).to(dtype), None, None, None, None, None, None, None, None, None, None, grad_qkv_linear_weight.to( |
| dtype), grad_qkv_linear_bias.to(dtype), None, None, None |
|
|
|
|
| class SequenceChunk: |
|
|
| def __init__(self, chunk: torch.Tensor, device=None, is_in_use=False): |
|
|
| self.chunk_shape = chunk.shape |
| self.chunk_dtype = chunk.dtype |
| self.device = chunk.device if device is None else device |
|
|
| cpu_chunk = torch.empty(chunk.shape, dtype=chunk.dtype, device='cpu', pin_memory=True) |
|
|
| if get_accelerator().on_accelerator(chunk): |
| cpu_chunk.copy_(chunk, non_blocking=True) |
| else: |
| cpu_chunk = chunk |
|
|
| self.cpu_chunk = cpu_chunk |
|
|
| self.gpu_chunk = chunk if is_in_use else None |
|
|
| def load_to_gpu(self): |
| assert self.gpu_chunk is None |
| if self.gpu_chunk is not None: |
| pass |
| else: |
| gpu_chunk = torch.empty(self.chunk_shape, device=self.device, dtype=self.chunk_dtype) |
| gpu_chunk.copy_(self.cpu_chunk, non_blocking=True) |
| self.gpu_chunk = gpu_chunk |
|
|
| def get_gpu_chunk(self): |
| assert self.gpu_chunk is not None and self.gpu_chunk.device == self.device |
| return self.gpu_chunk |
|
|
| def check_gpu_chunk(self, ): |
| assert (self.gpu_chunk is not None) and ( |
| self.gpu_chunk.device == self.device |
| ), f"gpu_chunk {self.gpu_chunk is not None} shound be on {self.device}, but it is now on {self.gpu_chunk.device}" |
| return True |
|
|
| def offload(self): |
| assert self.gpu_chunk is not None and self.gpu_chunk.device == self.device |
| del self.gpu_chunk |
| self.gpu_chunk = None |
|
|
| def overwrite_to_cpu(self): |
| assert self.gpu_chunk is not None and self.gpu_chunk.device == self.device |
| self.cpu_chunk.copy_(self.gpu_chunk, non_blocking=True) |
|
|
|
|
| class _FPDTGPUOffloadingAttentionImpl_(torch.autograd.Function): |
| generate_vmap_rule = False |
|
|
| @staticmethod |
| def forward(ctx: Any, |
| layernorm_output, |
| attention_mask, |
| inference_params, |
| rotary_pos_emb, |
| spg, |
| scatter_idx, |
| gather_idx, |
| hidden_size, |
| projection_size, |
| hidden_size_per_attention_head, |
| kv_projection_size, |
| qkv_linear_weight, |
| qkv_linear_bias, |
| dropout, |
| num_chunks=8, |
| cpu_offloading=True): |
|
|
| do_save = layernorm_output.requires_grad |
|
|
| if rotary_pos_emb is not None: |
| pos_emb_cos, pos_emb_sin = rotary_pos_emb[0].permute(1, 0, 2, 3), rotary_pos_emb[1].permute(1, 0, 2, 3) |
| ctx.pos_emb_cos = pos_emb_cos |
| ctx.pos_emb_sin = pos_emb_sin |
| else: |
| ctx.pos_emb_cos = None |
| ctx.pos_emb_sin = None |
| with torch.no_grad(): |
| per_gpu_seq_len = layernorm_output.shape[0] |
| chunk_size = per_gpu_seq_len // num_chunks |
| assert chunk_size * num_chunks == per_gpu_seq_len |
| assert attention_mask is None |
| ctx.num_chunks = num_chunks |
| ctx.cpu_offloading = cpu_offloading |
| ctx.spg = spg |
| ctx.scatter_idx = scatter_idx |
| ctx.gather_idx = gather_idx |
|
|
| ctx.chunk_size = chunk_size |
| device = get_accelerator().current_device_name() |
| ctx.device = device |
| ctx.dtype = layernorm_output.dtype |
| ctx.projection_size = projection_size |
| ctx.kv_projection_size = kv_projection_size |
|
|
| global_q = [] |
| global_k = [] |
| global_v = [] |
|
|
| ctx.softmax_scale = hidden_size_per_attention_head**(-0.5) |
|
|
| ctx.dropout_p = dropout |
| ctx.window_size = (-1, -1) |
| ctx.alibi_slopes = None |
|
|
| batch_size = layernorm_output.shape[1] |
|
|
| global_o = [] |
| global_lse = [] |
|
|
| layernorm_output_cpu = [] |
| final_output = [] |
|
|
| offload_stream = get_accelerator().Stream() |
| general_offload_stream = get_accelerator().Stream() |
| compute_stream = get_accelerator().default_stream() |
|
|
| q_compute_chunk_idx = 0 |
| kv_compute_chunk_idx = 0 |
| for i in range(num_chunks): |
|
|
| qkv_chunk = torch.matmul(layernorm_output[:chunk_size], |
| qkv_linear_weight.t()) + qkv_linear_bias |
|
|
| with get_accelerator().stream(general_offload_stream): |
| layernorm_output_cpu.append(SequenceChunk(layernorm_output[:chunk_size])) |
|
|
| layernorm_output = layernorm_output[chunk_size:] |
|
|
| q_chunk = qkv_chunk[:, :, :projection_size].contiguous().reshape( |
| qkv_chunk.shape[0], qkv_chunk.shape[1], -1, |
| hidden_size_per_attention_head).permute(1, 0, 2, 3).contiguous() |
| q_chunk = single_all_to_all(q_chunk, scatter_idx, gather_idx, 0, spg) |
| global_q_chunk_len = q_chunk.shape[1] |
|
|
| k_chunk = qkv_chunk[:, :, projection_size:projection_size + kv_projection_size].contiguous().reshape( |
| qkv_chunk.shape[0], qkv_chunk.shape[1], -1, |
| hidden_size_per_attention_head).permute(1, 0, 2, 3).contiguous() |
| k_chunk = single_all_to_all(k_chunk, scatter_idx, gather_idx, 0, spg) |
|
|
| v_chunk = qkv_chunk[:, :, projection_size + kv_projection_size:].contiguous().reshape( |
| qkv_chunk.shape[0], qkv_chunk.shape[1], -1, |
| hidden_size_per_attention_head).permute(1, 0, 2, 3).contiguous() |
| v_chunk = single_all_to_all(v_chunk, scatter_idx, gather_idx, 0, spg) |
|
|
| dist.barrier() |
|
|
| if ctx.pos_emb_cos is not None: |
| pos_emb_cos_chunk = pos_emb_cos[:, global_q_chunk_len * i:global_q_chunk_len * (i + 1)] |
| pos_emb_sin_chunk = pos_emb_sin[:, global_q_chunk_len * i:global_q_chunk_len * (i + 1)] |
|
|
| q_chunk = apply_rotary_pos_emb(q_chunk, pos_emb_cos_chunk, pos_emb_sin_chunk) |
| k_chunk = apply_rotary_pos_emb(k_chunk, pos_emb_cos_chunk, pos_emb_sin_chunk) |
|
|
| compute_stream.wait_stream(offload_stream) |
| compute_stream.synchronize() |
| with get_accelerator().stream(offload_stream): |
| global_q.append(SequenceChunk(q_chunk, is_in_use=True)) |
| global_k.append(SequenceChunk(k_chunk, is_in_use=True)) |
| global_v.append(SequenceChunk(v_chunk, is_in_use=True)) |
|
|
| del qkv_chunk |
|
|
| cur_attn_output = None |
| cur_attn_lse = None |
| for k_i in range(len(global_k)): |
| causal_chunk = i == k_i |
| with get_accelerator().stream(compute_stream): |
| if flash_attn_version >= version.parse("2.6.0"): |
| block_out, _, _, _, _, block_lse, _, _ = _flash_attn_forward( |
| global_q[q_compute_chunk_idx].get_gpu_chunk(), |
| global_k[kv_compute_chunk_idx].get_gpu_chunk(), |
| global_v[kv_compute_chunk_idx].get_gpu_chunk(), |
| ctx.dropout_p, |
| ctx.softmax_scale, |
| causal=causal_chunk, |
| window_size=ctx.window_size, |
| softcap=0.0, |
| alibi_slopes=ctx.alibi_slopes, |
| return_softmax=False) |
| else: |
| block_out, _, _, _, _, block_lse, _, _ = _flash_attn_forward( |
| global_q[q_compute_chunk_idx].get_gpu_chunk(), |
| global_k[kv_compute_chunk_idx].get_gpu_chunk(), |
| global_v[kv_compute_chunk_idx].get_gpu_chunk(), |
| ctx.dropout_p, |
| ctx.softmax_scale, |
| causal=causal_chunk, |
| window_size=ctx.window_size, |
| alibi_slopes=ctx.alibi_slopes, |
| return_softmax=False) |
| cur_attn_output, cur_attn_lse = update_out_and_lse(cur_attn_output, cur_attn_lse, block_out, |
| block_lse) |
|
|
| can_offload_kv = True |
| if k_i != (len(global_k) - 1) or i != (num_chunks - 1): |
| if k_i != (len(global_k) - 1): |
| next_kv_compute_chunk_idx = k_i + 1 |
| else: |
| next_kv_compute_chunk_idx = 0 |
|
|
| if next_kv_compute_chunk_idx == kv_compute_chunk_idx: |
| can_offload_kv = False |
| else: |
| if next_kv_compute_chunk_idx != (len(global_k) - 1): |
| with get_accelerator().stream(offload_stream): |
| global_k[next_kv_compute_chunk_idx].load_to_gpu() |
| global_v[next_kv_compute_chunk_idx].load_to_gpu() |
|
|
| if i == num_chunks - 1 and k_i == num_chunks - 1: |
| with get_accelerator().stream(offload_stream): |
| global_q[0].load_to_gpu() |
| global_k[0].load_to_gpu() |
| global_v[0].load_to_gpu() |
| global_o[0].load_to_gpu() |
| global_lse[0].load_to_gpu() |
|
|
| compute_stream.wait_stream(offload_stream) |
| compute_stream.synchronize() |
|
|
| if can_offload_kv: |
| global_k[kv_compute_chunk_idx].offload() |
| global_v[kv_compute_chunk_idx].offload() |
| kv_compute_chunk_idx = next_kv_compute_chunk_idx |
|
|
| global_q[q_compute_chunk_idx].offload() |
| q_compute_chunk_idx += 1 |
|
|
| all2all_output = single_all_to_all( |
| cur_attn_output.to(ctx.dtype).contiguous(), gather_idx, scatter_idx, 0, spg) |
| final_output.append(all2all_output) |
| with get_accelerator().stream(general_offload_stream): |
| global_o.append(SequenceChunk(cur_attn_output.to(ctx.dtype))) |
| global_lse.append(SequenceChunk(cur_attn_lse[:, :, :, 0].permute(0, 2, 1).contiguous())) |
|
|
| compute_stream.wait_stream(general_offload_stream) |
| compute_stream.synchronize() |
|
|
| final_output = torch.cat(final_output, dim=1) |
|
|
| head_dim = final_output.shape[-1] |
|
|
| if do_save: |
| ctx.layernorm_output = layernorm_output_cpu |
| ctx.global_q = global_q |
| ctx.global_k = global_k |
| ctx.global_v = global_v |
| ctx.attn_output = global_o |
| ctx.attn_lse = global_lse |
| ctx.head_dim = head_dim |
| ctx.batch_size = batch_size |
|
|
| ctx.qkv_linear_weight = qkv_linear_weight |
| ctx.qkv_linear_bias = qkv_linear_bias |
|
|
| return final_output |
|
|
| @staticmethod |
| def backward(ctx, grad_output): |
| num_chunks = ctx.num_chunks |
| device = grad_output.device |
| dtype = ctx.dtype |
| spg = ctx.spg |
| scatter_idx = ctx.scatter_idx |
| gather_idx = ctx.gather_idx |
| softmax_scale = ctx.softmax_scale |
| dropout_p = ctx.dropout_p |
| window_size = ctx.window_size |
| alibi_slopes = ctx.alibi_slopes |
|
|
| projection_size = ctx.projection_size |
| kv_projection_size = ctx.kv_projection_size |
|
|
| layernorm_output = ctx.layernorm_output |
|
|
| global_q = ctx.global_q |
| global_k = ctx.global_k |
| global_v = ctx.global_v |
| attn_output = ctx.attn_output |
| lse = ctx.attn_lse |
|
|
| qkv_linear_weight = ctx.qkv_linear_weight |
| qkv_linear_bias = ctx.qkv_linear_bias |
|
|
| offload_stream = get_accelerator().Stream() |
| general_offload_stream = get_accelerator().Stream() |
| compute_stream = get_accelerator().default_stream() |
|
|
| chunk_size = grad_output.shape[1] // num_chunks |
| assert chunk_size == layernorm_output[0].cpu_chunk.shape[0] |
|
|
| grad_layernorm_output = [ |
| torch.zeros(layernorm_output[0].chunk_shape, device=device, dtype=dtype) for _ in range(num_chunks) |
| ] |
|
|
| grad_global_attn_output = [None for _ in range(num_chunks)] |
|
|
| q_compute_chunk_idx = 0 |
| kv_compute_chunk_idx = 0 |
| last_q_accum_idx = 0 |
|
|
| with get_accelerator().stream(general_offload_stream): |
| layernorm_output[0].load_to_gpu() |
| grad_qkv_linear_weight = torch.zeros(qkv_linear_weight.shape, |
| device=qkv_linear_weight.device, |
| dtype=torch.float) |
| grad_qkv_linear_bias = torch.zeros(qkv_linear_bias.shape, |
| device=qkv_linear_weight.device, |
| dtype=torch.float) |
|
|
| grad_global_attn_output_chunk = single_all_to_all(grad_output[:, :chunk_size].contiguous(), scatter_idx, |
| gather_idx, 0, spg) |
| get_accelerator().synchronize() |
| grad_output = grad_output[:, chunk_size:] |
|
|
| with get_accelerator().stream(offload_stream): |
| grad_global_attn_output[0] = SequenceChunk(grad_global_attn_output_chunk, is_in_use=True) |
| dq = [ |
| SequenceChunk(torch.zeros(global_q[0].chunk_shape, dtype=torch.float, device=device), is_in_use=True) |
| ] + [ |
| SequenceChunk(torch.zeros(global_q[0].chunk_shape, dtype=torch.float, device='cpu', pin_memory=True), |
| device) for _ in range(num_chunks - 1) |
| ] |
| dk_accum = torch.zeros(global_k[0].chunk_shape, dtype=torch.float, device=device) |
| dv_accum = torch.zeros(global_v[0].chunk_shape, dtype=torch.float, device=device) |
|
|
| for i in range(num_chunks): |
| for q_i in range(num_chunks): |
| no_computation = q_i < i |
| if no_computation: |
| continue |
|
|
| causal_chunk = q_i == i |
|
|
| dq_this = torch.zeros(global_q[0].chunk_shape, dtype=dtype, device=device) |
| dk_this = torch.zeros(global_k[0].chunk_shape, dtype=dtype, device=device) |
| dv_this = torch.zeros(global_v[0].chunk_shape, dtype=dtype, device=device) |
|
|
| with get_accelerator().stream(compute_stream): |
| if flash_attn_version >= version.parse("2.6.0"): |
| _flash_attn_backward(grad_global_attn_output[q_compute_chunk_idx].get_gpu_chunk(), |
| global_q[q_compute_chunk_idx].get_gpu_chunk(), |
| global_k[kv_compute_chunk_idx].get_gpu_chunk(), |
| global_v[kv_compute_chunk_idx].get_gpu_chunk(), |
| attn_output[q_compute_chunk_idx].get_gpu_chunk(), |
| lse[q_compute_chunk_idx].get_gpu_chunk(), |
| dq_this, |
| dk_this, |
| dv_this, |
| dropout_p, |
| softmax_scale, |
| causal_chunk, |
| window_size, |
| softcap=0.0, |
| alibi_slopes=alibi_slopes, |
| deterministic=False, |
| rng_state=None) |
| else: |
| _flash_attn_backward(grad_global_attn_output[q_compute_chunk_idx].get_gpu_chunk(), |
| global_q[q_compute_chunk_idx].get_gpu_chunk(), |
| global_k[kv_compute_chunk_idx].get_gpu_chunk(), |
| global_v[kv_compute_chunk_idx].get_gpu_chunk(), |
| attn_output[q_compute_chunk_idx].get_gpu_chunk(), |
| lse[q_compute_chunk_idx].get_gpu_chunk(), |
| dq_this, |
| dk_this, |
| dv_this, |
| dropout_p, |
| softmax_scale, |
| causal_chunk, |
| window_size, |
| alibi_slopes=alibi_slopes, |
| deterministic=False, |
| rng_state=None) |
|
|
| if i != (len(global_k) - 1): |
| if q_i != (len(global_q) - 1): |
| next_q_compute_chunk_idx = q_i + 1 |
| else: |
| next_q_compute_chunk_idx = i + 1 |
|
|
| can_offload_q = True |
|
|
| if next_q_compute_chunk_idx == q_compute_chunk_idx: |
| can_offload_q = False |
| else: |
| with get_accelerator().stream(offload_stream): |
| if i > 0 or q_i > 0: |
| if can_offload_q and last_q_accum_idx != i: |
| dq[last_q_accum_idx].offload() |
| dq[next_q_compute_chunk_idx].load_to_gpu() |
| global_q[next_q_compute_chunk_idx].load_to_gpu() |
| attn_output[next_q_compute_chunk_idx].load_to_gpu() |
| lse[next_q_compute_chunk_idx].load_to_gpu() |
| if grad_global_attn_output[next_q_compute_chunk_idx] is not None: |
| grad_global_attn_output[next_q_compute_chunk_idx].load_to_gpu() |
|
|
| if grad_global_attn_output[next_q_compute_chunk_idx] is None: |
| grad_global_attn_output_chunk = single_all_to_all(grad_output[:, :chunk_size].contiguous(), |
| scatter_idx, gather_idx, 0, spg) |
| dist.barrier() |
| grad_output = grad_output[:, chunk_size:] |
| grad_global_attn_output[next_q_compute_chunk_idx] = SequenceChunk( |
| grad_global_attn_output_chunk, is_in_use=True) |
|
|
| compute_stream.wait_stream(offload_stream) |
| compute_stream.synchronize() |
|
|
| with get_accelerator().stream(compute_stream): |
| dq[q_compute_chunk_idx].check_gpu_chunk() |
| dq[q_compute_chunk_idx].gpu_chunk.add_(dq_this) |
| dk_accum.add_(dk_this) |
| dv_accum.add_(dv_this) |
|
|
| offload_stream.wait_stream(compute_stream) |
| with get_accelerator().stream(offload_stream): |
| dq[q_compute_chunk_idx].overwrite_to_cpu() |
|
|
| if can_offload_q: |
| global_q[q_compute_chunk_idx].offload() |
| attn_output[q_compute_chunk_idx].offload() |
| lse[q_compute_chunk_idx].offload() |
| grad_global_attn_output[q_compute_chunk_idx].offload() |
|
|
| last_q_accum_idx = q_compute_chunk_idx |
| q_compute_chunk_idx = next_q_compute_chunk_idx |
|
|
| compute_stream.wait_stream(offload_stream) |
| compute_stream.synchronize() |
|
|
| dk_seq_len = dk_accum.shape[1] |
|
|
| if ctx.pos_emb_cos is not None: |
| dq_accum = apply_rotary_pos_emb_backward(dq[kv_compute_chunk_idx].get_gpu_chunk().to(dtype), |
| ctx.pos_emb_cos[:, dk_seq_len * i:dk_seq_len * (i + 1)], |
| ctx.pos_emb_sin[:, dk_seq_len * i:dk_seq_len * (i + 1)]) |
| dk_accum = apply_rotary_pos_emb_backward(dk_accum.to(dtype), |
| ctx.pos_emb_cos[:, dk_seq_len * i:dk_seq_len * (i + 1)], |
| ctx.pos_emb_sin[:, dk_seq_len * i:dk_seq_len * (i + 1)]) |
| else: |
| dq_accum = dq[kv_compute_chunk_idx].get_gpu_chunk().to(dtype) |
| dk_accum = dk_accum.to(dtype) |
| dv_accum = dv_accum.to(dtype) |
|
|
| dq_accum = single_all_to_all(dq_accum.contiguous(), gather_idx, scatter_idx, 0, spg) |
| dk_accum = single_all_to_all(dk_accum.contiguous(), gather_idx, scatter_idx, 0, spg) |
| dv_accum = single_all_to_all(dv_accum.contiguous(), gather_idx, scatter_idx, 0, spg) |
|
|
| general_offload_stream.synchronize() |
| compute_stream.wait_stream(general_offload_stream) |
| dist.barrier() |
|
|
| with get_accelerator().stream(compute_stream): |
| input_chunk = layernorm_output[i].get_gpu_chunk().reshape(-1, layernorm_output[i].chunk_shape[-1]) |
|
|
| dq_accum = dq_accum.flatten(2).permute(1, 0, 2) |
| dk_accum = dk_accum.flatten(2).permute(1, 0, 2) |
| dv_accum = dv_accum.flatten(2).permute(1, 0, 2) |
|
|
| l, b = dk_accum.shape[0], dk_accum.shape[1] |
|
|
| grad_qkv_linear_weight[:projection_size].add_( |
| torch.matmul(dq_accum.reshape(l * b, -1).t(), input_chunk)) |
| grad_qkv_linear_weight[projection_size:projection_size + kv_projection_size].add_( |
| torch.matmul(dk_accum.reshape(l * b, -1).t(), input_chunk)) |
| grad_qkv_linear_weight[projection_size + kv_projection_size:].add_( |
| torch.matmul(dv_accum.reshape(l * b, -1).t(), input_chunk)) |
|
|
| grad_qkv_linear_bias[:projection_size].add_(dq_accum.sum(0).sum(0)) |
| grad_qkv_linear_bias[projection_size:projection_size + kv_projection_size].add_(dk_accum.sum(0).sum(0)) |
| grad_qkv_linear_bias[projection_size + kv_projection_size:].add_(dv_accum.sum(0).sum(0)) |
|
|
| grad_layernorm_output[i].add_(torch.matmul(dq_accum, qkv_linear_weight[:projection_size])) |
| grad_layernorm_output[i].add_( |
| torch.matmul(dk_accum, qkv_linear_weight[projection_size:projection_size + kv_projection_size])) |
| grad_layernorm_output[i].add_( |
| torch.matmul(dv_accum, qkv_linear_weight[projection_size + kv_projection_size:])) |
|
|
| del dq_accum, dk_accum, dv_accum |
| dk_accum = torch.zeros(global_k[i].chunk_shape, dtype=torch.float, device=device) |
| dv_accum = torch.zeros(global_v[i].chunk_shape, dtype=torch.float, device=device) |
| dq[kv_compute_chunk_idx].offload() |
| dq[kv_compute_chunk_idx] = None |
|
|
| if i != (len(global_k) - 1): |
| next_kv_compute_chunk_idx = kv_compute_chunk_idx + 1 |
| with get_accelerator().stream(offload_stream): |
| global_k[next_kv_compute_chunk_idx].load_to_gpu() |
| global_v[next_kv_compute_chunk_idx].load_to_gpu() |
|
|
| with get_accelerator().stream(general_offload_stream): |
| layernorm_output[next_kv_compute_chunk_idx].load_to_gpu() |
|
|
| compute_stream.wait_stream(offload_stream) |
| compute_stream.synchronize() |
|
|
| layernorm_output[kv_compute_chunk_idx].offload() |
| global_k[kv_compute_chunk_idx].offload() |
| global_v[kv_compute_chunk_idx].offload() |
| kv_compute_chunk_idx = next_kv_compute_chunk_idx |
|
|
| return torch.cat( |
| grad_layernorm_output, |
| dim=0).to(dtype), None, None, None, None, None, None, None, None, None, None, grad_qkv_linear_weight.to( |
| dtype), grad_qkv_linear_bias.to(dtype), None, None, None |
|
|
|
|
| class FPDT_Attention(torch.nn.Module): |
|
|
| def __init__(self, |
| config, |
| first_weight, |
| first_bias, |
| second_weight, |
| second_bias, |
| sequence_process_group, |
| gather_idx: int = 0, |
| scatter_idx: int = 2, |
| return_bias=True, |
| chunk_size=65536, |
| enable_offloading=True) -> None: |
|
|
| super(FPDT_Attention, self).__init__() |
| if _flash_attn_forward is None or _flash_attn_backward is None: |
| raise ImportError( |
| "DeepSpeed FPDT requires flash-attn 2.6.3. Please install it with `pip install flash-attn --no-build-isolation`." |
| ) |
|
|
| self.spg = sequence_process_group |
| self.scatter_idx = scatter_idx |
| self.gather_idx = gather_idx |
| self.config = config |
|
|
| self.projection_size = config.kv_channels * config.num_attention_heads |
| self.hidden_size_per_attention_head = self.projection_size // config.num_attention_heads |
| self.kv_projection_size = config.kv_channels * config.num_key_value_heads |
| self.hidden_size = config.hidden_size |
|
|
| self.qkv_linear_weight = first_weight |
| self.qkv_linear_bias = first_bias |
| self.qkv_dense_weight = second_weight |
| self.qkv_dense_bias = second_bias |
|
|
| self.reture_bias = return_bias |
| self.dropout = config.attention_dropout |
|
|
| self.chunk_size = chunk_size |
| self.double_buffer = enable_offloading |
|
|
| def forward(self, |
| layernorm_output, |
| attention_mask, |
| inference_params, |
| rotary_pos_emb, |
| cpu_offloading=True) -> Tensor: |
| self.num_chunks_attn = layernorm_output.shape[0] * dist.get_world_size(self.spg) // self.chunk_size |
|
|
| if not cpu_offloading or self.num_chunks_attn == 1: |
| output = _FPDTGPUAttentionImpl_.apply(layernorm_output, attention_mask, inference_params, rotary_pos_emb, |
| self.spg, self.scatter_idx, self.gather_idx, self.hidden_size, |
| self.projection_size, self.hidden_size_per_attention_head, |
| self.kv_projection_size, self.qkv_linear_weight, |
| self.qkv_linear_bias, self.dropout, self.num_chunks_attn, |
| cpu_offloading) |
| else: |
| output = _FPDTGPUOffloadingAttentionImpl_.apply( |
| layernorm_output, attention_mask, inference_params, rotary_pos_emb, self.spg, self.scatter_idx, |
| self.gather_idx, self.hidden_size, self.projection_size, self.hidden_size_per_attention_head, |
| self.kv_projection_size, self.qkv_linear_weight, self.qkv_linear_bias, self.dropout, |
| self.num_chunks_attn, cpu_offloading) |
|
|
| output = output.flatten(2).permute(1, 0, 2).contiguous() |
|
|
| output = torch.matmul(output, self.qkv_dense_weight.t()) |
| if not self.reture_bias: |
| output += self.qkv_dense_bias |
| return output, self.qkv_dense_bias if self.reture_bias else None |
|
|
|
|
| @torch.jit.script |
| def bias_gelu(x): |
| return x * 0.5 * (1.0 + torch.tanh(0.79788456 * x * (1 + 0.044715 * x * x))) |
|
|
|
|
| @torch.jit.script |
| def bias_gelu_back(g, x): |
| tanh_out = torch.tanh(0.79788456 * x * (1 + 0.044715 * x * x)) |
| |
| ff = 0.5 * x * ((1 - tanh_out * tanh_out) * (0.79788456 + 0.1070322243 * x * x)) + 0.5 * (1 + tanh_out) |
| return ff * g |
|
|
|
|
| class FPDT_FFN(torch.autograd.Function): |
| generate_vmap_rule = False |
|
|
| @staticmethod |
| def forward(ctx: Any, x, w1, b1, w2, b2, add_bias, chunk_size): |
| do_save = x.requires_grad |
| ctx.add_bias = add_bias |
| device = x.device |
|
|
| with torch.no_grad(): |
| num_chunk = x.shape[0] // chunk_size |
| ctx.num_chunk = num_chunk |
| result = torch.empty(x.shape, device=device, dtype=x.dtype) |
| assert chunk_size * num_chunk == x.shape[0] |
| for i in range(num_chunk): |
| st = i * chunk_size |
| ed = st + chunk_size |
| x_ = torch.matmul(x[st:ed], w1.t()) + b1 |
| x_ = bias_gelu(x_) |
| if add_bias: |
| result[st:ed] = torch.matmul(x_, w2.t()) + b2 |
| else: |
| result[st:ed] = torch.matmul(x_, w2.t()) |
|
|
| del x_ |
|
|
| if do_save: |
| ctx.device = device |
| ctx.dtype = x.dtype |
| ctx.save_for_backward(x, w1, b1, w2, b2) |
| ctx.grad_x_shape = x.shape |
| return result.to(x.dtype), b2 if not add_bias else None |
|
|
| @staticmethod |
| def backward(ctx, grad_output, grad_bias): |
| x, w1, b1, w2, b2 = ctx.saved_tensors |
| device = ctx.device |
| dtype = ctx.dtype |
| add_bias = ctx.add_bias |
|
|
| num_chunk = ctx.num_chunk |
| chunk_size = x.shape[0] // num_chunk |
| assert chunk_size * num_chunk == grad_output.shape[0] |
|
|
| grad_w2 = torch.zeros(w2.shape, device=device, dtype=torch.float) |
| grad_b2 = torch.zeros(b2.shape, device=device, dtype=torch.float) |
| grad_w1 = torch.zeros(w1.shape, device=device, dtype=torch.float) |
| grad_b1 = torch.zeros(b1.shape, device=device, dtype=torch.float) |
|
|
| for i in range(num_chunk): |
| st = i * chunk_size |
| ed = st + chunk_size |
| x_chunk = x[st:ed] |
|
|
| before_act = (torch.matmul(x_chunk, w1.t()) + b1) |
| before_act_2 = before_act**2 |
| tanh_out = torch.tanh(0.79788456 * before_act * (1 + 0.044715 * before_act_2)) |
| ff = 0.5 * before_act * ((1 - tanh_out * tanh_out) * |
| (0.79788456 + 0.1070322243 * before_act_2)) + 0.5 * (1 + tanh_out) |
| grad_w2.add_( |
| torch.matmul(grad_output[st:ed].reshape(-1, grad_output.shape[2]).t(), |
| (before_act * 0.5 * (1 + tanh_out)).reshape(-1, before_act.shape[2]))) |
| del before_act, before_act_2, tanh_out |
|
|
| grad_inter = torch.matmul(grad_output[st:ed], w2) * ff |
| del ff |
|
|
| grad_w1.add_(torch.matmul( |
| grad_inter.reshape(-1, grad_inter.shape[2]).t(), x_chunk.reshape(-1, x.shape[2]))) |
| grad_b1.add_(grad_inter.sum(0).sum(0)) |
|
|
| x[st:ed].copy_(torch.matmul(grad_inter, w1)) |
|
|
| del grad_inter |
|
|
| if add_bias: |
| grad_b2.add_(grad_output[st:ed].sum(0).sum(0)) |
|
|
| return x, grad_w1.to(dtype), grad_b1.to(dtype), grad_w2.to(dtype), grad_b2.to(dtype), None, None |
|
|
|
|
| class FPDT_LogitsLoss(torch.autograd.Function): |
| generate_vmap_rule = False |
|
|
| @staticmethod |
| def forward(ctx: Any, lm_output, labels, logit_weights, rank, spg_size, spg, num_chunk): |
| labels = labels.t() |
| chunk_size = lm_output.shape[0] // num_chunk |
| assert chunk_size * num_chunk == lm_output.shape[0] |
| batch_size, local_seq_len = lm_output.shape[1], lm_output.shape[0] |
| loss = torch.empty((batch_size, local_seq_len), dtype=torch.float, device=lm_output.device) |
|
|
| ctx.num_chunk = num_chunk |
| ctx.chunk_size = chunk_size |
| ctx.device = lm_output.device |
| ctx.dtype = lm_output.dtype |
|
|
| ctx.rank = rank |
| ctx.local_seq_len = local_seq_len |
| with torch.no_grad(): |
| for i in range(num_chunk): |
| st = i * chunk_size |
| ed = st + chunk_size |
| logits_chunk = torch.matmul(lm_output[st:ed], logit_weights.t()).float() |
|
|
| vocab_size = logits_chunk.size(2) |
| |
| softmax = torch.nn.functional.softmax(logits_chunk, dim=-1) |
| loss_chunk = torch.nn.functional.nll_loss(softmax.log().reshape(-1, vocab_size).contiguous(), |
| labels[st:ed, :].reshape(-1).contiguous(), |
| reduction='none') |
| loss[:, st:ed] = loss_chunk.reshape(chunk_size, batch_size).t() |
|
|
| del logits_chunk |
| ctx.save_for_backward(lm_output.to('cpu'), labels) |
| ctx.logit_weights = logit_weights |
|
|
| seqlen = local_seq_len * spg_size |
| batch_size = loss.size(0) |
| loss = loss.t().contiguous() |
| loss_all = torch.empty(seqlen, batch_size, dtype=loss.dtype, device=loss.device).contiguous() |
|
|
| dist.allgather_fn(loss_all, loss, group=spg) |
|
|
| return loss_all |
|
|
| @staticmethod |
| def backward(ctx, grad_output): |
| lm_output, labels = ctx.saved_tensors |
| logit_weights = ctx.logit_weights |
| device = ctx.device |
| dtype = ctx.dtype |
| num_chunk = ctx.num_chunk |
| chunk_size = ctx.chunk_size |
|
|
| rank = ctx.rank |
| local_seq_len = ctx.local_seq_len |
|
|
| grad_output = grad_output[rank * local_seq_len:(rank + 1) * local_seq_len] |
| grad_lm_output = [None for _ in range(num_chunk)] |
| grad_logit_weights = torch.zeros(logit_weights.shape, device=grad_output.device, dtype=torch.float) |
| for i in range(num_chunk): |
| st = i * chunk_size |
| ed = st + chunk_size |
| lm_output_chunk = lm_output[st:ed].to(device) |
| logits_chunk = torch.matmul(lm_output_chunk, logit_weights.t()).float() |
|
|
| |
| softmax = torch.nn.functional.softmax(logits_chunk, dim=-1) |
| vocab_size = logits_chunk.size(2) |
|
|
| grad_input = softmax |
| grad_2d = grad_input.reshape(-1, vocab_size).contiguous() |
| arange_1d = torch.arange(start=0, end=grad_2d.size()[0], device=device) |
|
|
| grad_2d[arange_1d, labels[st:ed, :].reshape(-1).contiguous()] -= 1 |
| grad_input.mul_(grad_output[:chunk_size, :].unsqueeze(dim=-1)) |
| grad_input = grad_input.to(dtype) |
|
|
| grad_output = grad_output[chunk_size:].contiguous() |
|
|
| grad_lm_output_chunk = torch.matmul(grad_input, logit_weights) |
| grad_lm_output[i] = grad_lm_output_chunk |
|
|
| grad_logit_weights.add_( |
| torch.matmul( |
| grad_input.reshape(-1, grad_input.shape[2]).t(), |
| lm_output_chunk.reshape(-1, lm_output_chunk.shape[2]))) |
|
|
| return torch.cat(grad_lm_output, dim=0).to(dtype), None, grad_logit_weights.to(dtype), None, None, None, None |
|
|
