apply flagos te_groups_gemm op - #55
Conversation
| def gelu_backward(grad_output, x): | ||
| # Approximation of GELU derivative commonly used in Transformer Engine | ||
| cdf = 0.5 * (1.0 + flag_gems.erf(x / flag_gems.sqrt(2.0))) | ||
| pdf = flag_gems.exp(-0.5 * x * x) / flag_gems.sqrt(2.0 * math.pi) |
| else: | ||
| out = torch.zeros(A[i].shape[0], B[i].shape[1]) | ||
| if grad and len(bias) > i and bias[i] is not None and bias[i].numel() != 0: | ||
| bias[i].zero_() |
| n = B[i].shape[0] if transb else B[i].shape[1] | ||
| D.append(torch.empty((m, n), dtype=D[i].dtype, device=A[0].device)) | ||
|
|
||
| def gelu_backward(grad_output, x): |
There was a problem hiding this comment.
why not use flag_gems.gelu_backward
There was a problem hiding this comment.
OK，Will attempt to replace and test.
| return out1, bias_grad, gelu_input, extra_output_ret | ||
|
|
||
|
|
||
| def te_general_grouped_gemm_fl( |
There was a problem hiding this comment.
Please add a comment indicating that this function can represent both forward computation and backward computation, distinguished by the grad parameter.
| if not single_output: | ||
| # Store output | ||
| if accumulate: | ||
| D[i].add_(out.to(D[i].dtype)) |
There was a problem hiding this comment.
use flag_gems.to_copy to replace xxx.to
There was a problem hiding this comment.
OK，Will attempt to replace and test.
| for i in range(num_gemms): | ||
| m = A[i].shape[1] if transa else A[i].shape[0] | ||
| n = B[i].shape[0] if transb else B[i].shape[1] | ||
| D.append(torch.empty((m, n), dtype=D[i].dtype, device=A[0].device)) |
There was a problem hiding this comment.
why not use flag_gems.zeros
| if A[i].numel() == 0 or B[i].numel() == 0: | ||
| if not single_output: | ||
| if D[i].numel() != 0 and not accumulate: | ||
| D[i].zero_() |
There was a problem hiding this comment.
try to use flag_gems.zeros
| if D[i].numel() != 0 and not accumulate: | ||
| D[i].zero_() | ||
| else: | ||
| out = torch.zeros(A[i].shape[0], B[i].shape[1]) |
There was a problem hiding this comment.
try to use flag_gems.zeros
| and pre_gelu_out[i] is not None | ||
| and pre_gelu_out[i].numel() != 0 | ||
| ): | ||
| pre_gelu_out[i].zero_() |
There was a problem hiding this comment.
try to use flag_gems.zeros
|
|
||
| # Apply GELU epilogue if pre_gelu_out is provided | ||
| if has_pre_gelu: | ||
| pre_gelu_out[i].copy_(out) |
There was a problem hiding this comment.
try to use flag_gems.copy
|
|
||
| # Compute bias gradients if requested | ||
| if has_bias: | ||
| bias_grad = out.sum(dim=0) |
There was a problem hiding this comment.
try to use flag_gems.sum_dim
| if has_bias: | ||
| bias_grad = out.sum(dim=0) | ||
| if accumulate: | ||
| bias[i].add_(bias_grad) |
There was a problem hiding this comment.
try to use flag_gems.add
| if accumulate: | ||
| bias[i].add_(bias_grad) | ||
| else: | ||
| bias[i].copy_(bias_grad) |
There was a problem hiding this comment.
try to use flag_gems.copy
| if accumulate: | ||
| D[i].add_(out.to(D[i].dtype)) | ||
| else: | ||
| D[i].copy_(out.to(D[i].dtype)) |
| else: | ||
| D[i].copy_(out.to(D[i].dtype)) | ||
| else: | ||
| temp_D.append(out.to(D[0].dtype)) |
There was a problem hiding this comment.
try to use flag_gems.to_copy
|
|
||
| # Compute bias gradients if requested | ||
| if has_bias: | ||
| bias_grad = flag_gems.sum_dim(out, dim=0) |
There was a problem hiding this comment.
please check dim argument of flag_gems.sum_dim, is it should be a list or not
There was a problem hiding this comment.
In the flag_gems.sum_dim operator, the dim parameter supports multiple data types for input, including None, int, and List[int].
There was a problem hiding this comment.
please check
def sum_dim_comm(inp, dim=None, keepdim=False, *, dtype=None, out=None):
if dtype is None:
dtype = inp.dtype
if dtype is torch.bool:
dtype = torch.int64
if dim == []:
if not keepdim:
return sum(inp, dtype=dtype)
else:
dim_num = inp.ndim
return torch.reshape(sum(inp, dtype=dtype), [1] * dim_num)
shape = list(inp.shape)
dim = [d % inp.ndim for d in dim]
| print(f" ✓ {backend_name}") | ||
| except NotImplementedError: | ||
| self.skipped += 1 | ||
| print(f" �~J~X {backend_name} (not implemented)") |
| print(f" ✓ {backend_name}") | ||
| except NotImplementedError: | ||
| self.skipped += 1 | ||
| print(f" �~J~X {backend_name} (not implemented)") |
3 participants
te_general_grouped_gemmop for flagos backend, base on flag_gemsgrad