Merged
feat: add npugraph_ex support for engram_grad_w_reduce #55
wangjiaxingAdd on 13 days ago
feat: add npugraph_ex support for engram_grad_w_reduce #55
Merged
Total: 2 files changed+165-8
| @@ -28,20 +28,13 @@ def _check_tensor( | |||
| 28 | raise ValueError(f'{name} must be contiguous.') | 28 | raise ValueError(f'{name} must be contiguous.') |
| 29 | 29 | ||
| 30 | 30 | ||
| 31 | -def engram_grad_w_reduce( | 31 | +def _validate_inputs( |
| 32 | grad_w_partial: torch.Tensor, | 32 | grad_w_partial: torch.Tensor, |
| 33 | weight_hidden: torch.Tensor, | 33 | weight_hidden: torch.Tensor, |
| 34 | weight_embed: torch.Tensor, | 34 | weight_embed: torch.Tensor, |
| 35 | grad_weight_hidden: torch.Tensor, | 35 | grad_weight_hidden: torch.Tensor, |
| 36 | grad_weight_embed: torch.Tensor, | 36 | grad_weight_embed: torch.Tensor, |
| 37 | ) -> None: | 37 | ) -> None: |
| 38 | - """Accumulate partial fused-weight gradients into both unfused weights. | ||
| 39 | - | ||
| 40 | - ``grad_weight_hidden`` and ``grad_weight_embed`` are FP32 accumulation | ||
| 41 | - buffers and are updated in place. The function returns ``None``. | ||
| 42 | - """ | ||
| 43 | - | ||
| 44 | - require_ascend() | ||
| 45 | _check_tensor('grad_w_partial', grad_w_partial, dtype=torch.float32) | 38 | _check_tensor('grad_w_partial', grad_w_partial, dtype=torch.float32) |
| 46 | _check_tensor('weight_hidden', weight_hidden, dtype=torch.bfloat16, device=grad_w_partial.device) | 39 | _check_tensor('weight_hidden', weight_hidden, dtype=torch.bfloat16, device=grad_w_partial.device) |
| 47 | _check_tensor('weight_embed', weight_embed, dtype=torch.bfloat16, device=grad_w_partial.device) | 40 | _check_tensor('weight_embed', weight_embed, dtype=torch.bfloat16, device=grad_w_partial.device) |
| @@ -77,7 +70,59 @@ def engram_grad_w_reduce( | |||
| 77 | if tensor.shape != expected_shape: | 70 | if tensor.shape != expected_shape: |
| 78 | raise ValueError(f'{name} must have shape {expected_shape}, got {tuple(tensor.shape)}.') | 71 | raise ValueError(f'{name} must have shape {expected_shape}, got {tuple(tensor.shape)}.') |
| 79 | 72 | ||
| 73 | + | ||
| 74 | +def _engram_grad_w_reduce_impl( | ||
| 75 | + grad_w_partial: torch.Tensor, | ||
| 76 | + weight_hidden: torch.Tensor, | ||
| 77 | + weight_embed: torch.Tensor, | ||
| 78 | + grad_weight_hidden: torch.Tensor, | ||
| 79 | + grad_weight_embed: torch.Tensor, | ||
| 80 | +) -> None: | ||
| 81 | + require_ascend() | ||
| 82 | + _validate_inputs(grad_w_partial, weight_hidden, weight_embed, grad_weight_hidden, grad_weight_embed) | ||
| 83 | + num_partials, hc_mult, hidden_size = grad_w_partial.shape | ||
| 80 | kernel = get_engram_grad_w_reduce_kernel(hidden_size, num_partials, hc_mult) | 84 | kernel = get_engram_grad_w_reduce_kernel(hidden_size, num_partials, hc_mult) |
| 81 | if should_print_kernel_source(): | 85 | if should_print_kernel_source(): |
| 82 | print(kernel.get_kernel_source()) | 86 | print(kernel.get_kernel_source()) |
| 83 | kernel(grad_w_partial, weight_hidden, weight_embed, grad_weight_hidden, grad_weight_embed) | 87 | kernel(grad_w_partial, weight_hidden, weight_embed, grad_weight_hidden, grad_weight_embed) |
| 88 | + | ||
| 89 | + | ||
| 90 | +@torch.library.custom_op( | ||
| 91 | + 'cann_ops_tilelang::engram_grad_w_reduce', | ||
| 92 | + mutates_args=('grad_weight_hidden', 'grad_weight_embed'), | ||
| 93 | +) | ||
| 94 | +def _compiled_engram_grad_w_reduce( | ||
| 95 | + grad_w_partial: torch.Tensor, | ||
| 96 | + weight_hidden: torch.Tensor, | ||
| 97 | + weight_embed: torch.Tensor, | ||
| 98 | + grad_weight_hidden: torch.Tensor, | ||
| 99 | + grad_weight_embed: torch.Tensor, | ||
| 100 | +) -> None: | ||
| 101 | + _engram_grad_w_reduce_impl(grad_w_partial, weight_hidden, weight_embed, grad_weight_hidden, grad_weight_embed) | ||
| 102 | + | ||
| 103 | + | ||
| 104 | +@_compiled_engram_grad_w_reduce.register_fake | ||
| 105 | +def _engram_grad_w_reduce_fake( | ||
| 106 | + grad_w_partial: torch.Tensor, | ||
| 107 | + weight_hidden: torch.Tensor, | ||
| 108 | + weight_embed: torch.Tensor, | ||
| 109 | + grad_weight_hidden: torch.Tensor, | ||
| 110 | + grad_weight_embed: torch.Tensor, | ||
| 111 | +) -> None: | ||
| 112 | + _validate_inputs(grad_w_partial, weight_hidden, weight_embed, grad_weight_hidden, grad_weight_embed) | ||
| 113 | + | ||
| 114 | + | ||
| 115 | +def engram_grad_w_reduce( | ||
| 116 | + grad_w_partial: torch.Tensor, | ||
| 117 | + weight_hidden: torch.Tensor, | ||
| 118 | + weight_embed: torch.Tensor, | ||
| 119 | + grad_weight_hidden: torch.Tensor, | ||
| 120 | + grad_weight_embed: torch.Tensor, | ||
| 121 | +) -> None: | ||
| 122 | + """Accumulate into both FP32 gradient buffers in place and return None.""" | ||
| 123 | + if torch.compiler.is_compiling(): | ||
| 124 | + _compiled_engram_grad_w_reduce( | ||
| 125 | + grad_w_partial, weight_hidden, weight_embed, grad_weight_hidden, grad_weight_embed | ||
| 126 | + ) | ||
| 127 | + else: | ||
| 128 | + _engram_grad_w_reduce_impl(grad_w_partial, weight_hidden, weight_embed, grad_weight_hidden, grad_weight_embed) | ||
| @@ -110,3 +110,115 @@ def test_engram_grad_w_reduce_benchmark(benchmark_timer, benchmark_record, param | |||
| 110 | time_us=t_us, | 110 | time_us=t_us, |
| 111 | bandwidth_gbs=bandwidth_gbs, | 111 | bandwidth_gbs=bandwidth_gbs, |
| 112 | ) | 112 | ) |
| 113 | + | ||
| 114 | + | ||
| 115 | +def _compile_inputs_and_storage(partials=7, hidden=256): | ||
| 116 | + inputs = list(generate_test_data({'hidden': hidden, 'num_sms': partials})) | ||
| 117 | + storages = [] | ||
| 118 | + for _ in range(2): | ||
| 119 | + # Contiguous offset views also check that surrounding storage is untouched. | ||
| 120 | + storage = torch.randn(4 * hidden + 2, device=inputs[0].device) | ||
| 121 | + inputs.append(storage[1:-1].view(4, hidden)) | ||
| 122 | + storages.append(storage) | ||
| 123 | + return inputs, storages | ||
| 124 | + | ||
| 125 | + | ||
| 126 | +def _compile_inputs(partials=7, hidden=256): | ||
| 127 | + inputs, _ = _compile_inputs_and_storage(partials, hidden) | ||
| 128 | + return inputs | ||
| 129 | + | ||
| 130 | + | ||
| 131 | +def _assert_accumulation(call, inputs, expected, storages): | ||
| 132 | + readonly = [tensor.clone() for tensor in inputs[:3]] | ||
| 133 | + surrounding = [(storage, storage[[0, -1]].clone()) for storage in storages] | ||
| 134 | + pointers = [tensor.data_ptr() for tensor in inputs[3:]] | ||
| 135 | + grad_w_reduce_ref(*inputs[:3], *expected) | ||
| 136 | + assert call(*inputs) is None | ||
| 137 | + for actual, reference, pointer in zip(inputs[3:], expected, pointers, strict=True): | ||
| 138 | + assert actual.data_ptr() == pointer | ||
| 139 | + assert actual.dtype == torch.float32 and actual.is_contiguous() | ||
| 140 | + assert calc_diff(actual.cpu(), reference.cpu()) < 1e-10 | ||
| 141 | + for actual, before in zip(inputs[:3], readonly, strict=True): | ||
| 142 | + assert torch.equal(actual, before) | ||
| 143 | + for storage, before in surrounding: | ||
| 144 | + assert torch.equal(storage[[0, -1]], before) | ||
| 145 | + | ||
| 146 | + | ||
| 147 | +@pytest.mark.parametrize('partials,hidden', [(1, 256), (7, 256), (37, 2560), (38, 3072), (75, 6144), (7, 4096)]) | ||
| 148 | +def test_engram_grad_w_reduce_compile(partials, hidden): | ||
| 149 | + inputs, storages = _compile_inputs_and_storage(partials, hidden) | ||
| 150 | + expected = [tensor.clone() for tensor in inputs[3:]] | ||
| 151 | + compiled = torch.compile(engram_grad_w_reduce, backend='npugraph_ex', fullgraph=True, dynamic=False) | ||
| 152 | + for step in range(3): | ||
| 153 | + if step == 1: | ||
| 154 | + for tensor in inputs[:3]: | ||
| 155 | + tensor.copy_(torch.randn_like(tensor)) | ||
| 156 | + elif step == 2: | ||
| 157 | + replacement = _compile_inputs(partials, hidden) | ||
| 158 | + inputs[:3] = replacement[:3] | ||
| 159 | + with torch.inference_mode(): | ||
| 160 | + _assert_accumulation(compiled, inputs, expected, storages) | ||
| 161 | + # Switching persistent buffers must update the new caller-owned storage. | ||
| 162 | + replacement, storages = _compile_inputs_and_storage(partials, hidden) | ||
| 163 | + expected = [tensor.clone() for tensor in replacement[3:]] | ||
| 164 | + with torch.inference_mode(): | ||
| 165 | + _assert_accumulation(compiled, replacement, expected, storages) | ||
| 166 | + | ||
| 167 | + | ||
| 168 | +def test_engram_grad_w_reduce_opcheck(): | ||
| 169 | + from cann_ops_tilelang.engram.engram_grad_w_reduce import _compiled_engram_grad_w_reduce | ||
| 170 | + | ||
| 171 | + torch.library.opcheck( | ||
| 172 | + _compiled_engram_grad_w_reduce, | ||
| 173 | + tuple(_compile_inputs()), | ||
| 174 | + test_utils=('test_schema', 'test_faketensor'), | ||
| 175 | + ) | ||
| 176 | + | ||
| 177 | + | ||
| 178 | +def test_engram_grad_w_reduce_compile_composition(): | ||
| 179 | + def model(partial, wh, we, gwh, gwe): | ||
| 180 | + engram_grad_w_reduce(partial * 0.5, wh, we, gwh, gwe) | ||
| 181 | + return gwh + gwe | ||
| 182 | + | ||
| 183 | + inputs = _compile_inputs() | ||
| 184 | + expected = [tensor.clone() for tensor in inputs[3:]] | ||
| 185 | + compiled = torch.compile(model, backend='npugraph_ex', fullgraph=True, dynamic=False) | ||
| 186 | + for _ in range(3): | ||
| 187 | + inputs[0].copy_(torch.randn_like(inputs[0])) | ||
| 188 | + grad_w_reduce_ref(inputs[0] * 0.5, *inputs[1:3], *expected) | ||
| 189 | + with torch.inference_mode(): | ||
| 190 | + actual = compiled(*inputs) | ||
| 191 | + torch.testing.assert_close(actual, expected[0] + expected[1], rtol=1e-5, atol=1e-5) | ||
| 192 | + for tensor, reference in zip(inputs[3:], expected, strict=True): | ||
| 193 | + assert calc_diff(tensor.cpu(), reference.cpu()) < 1e-10 | ||
| 194 | + del actual | ||
| 195 | + | ||
| 196 | + | ||
| 197 | +@pytest.mark.parametrize('case', ['empty', 'hidden', 'hc', 'rank', 'dtype', 'shape', 'layout']) | ||
| 198 | +def test_engram_grad_w_reduce_invalid_contract(case): | ||
| 199 | + from torch._subclasses.fake_tensor import FakeTensorMode | ||
| 200 | + | ||
| 201 | + from cann_ops_tilelang.engram.engram_grad_w_reduce import _engram_grad_w_reduce_fake | ||
| 202 | + | ||
| 203 | + inputs = _compile_inputs() | ||
| 204 | + device = inputs[0].device | ||
| 205 | + if case == 'empty': | ||
| 206 | + inputs[0] = torch.empty(0, 4, 256, device=device) | ||
| 207 | + elif case == 'hidden': | ||
| 208 | + inputs[0] = torch.empty(7, 4, 128, device=device) | ||
| 209 | + elif case == 'hc': | ||
| 210 | + inputs[0] = torch.empty(7, 2, 256, device=device) | ||
| 211 | + elif case == 'rank': | ||
| 212 | + inputs[0] = inputs[0].flatten() | ||
| 213 | + elif case == 'dtype': | ||
| 214 | + inputs[1] = inputs[1].float() | ||
| 215 | + elif case == 'shape': | ||
| 216 | + inputs[4] = inputs[4][:2] | ||
| 217 | + else: | ||
| 218 | + inputs[0] = torch.empty(7, 4, 512, device=device)[..., ::2] | ||
| 219 | + with pytest.raises((ValueError, TypeError)): | ||
| 220 | + engram_grad_w_reduce(*inputs) | ||
| 221 | + with FakeTensorMode() as mode: | ||
| 222 | + fake = [mode.from_tensor(tensor) for tensor in inputs] | ||
| 223 | + with pytest.raises((ValueError, TypeError)): | ||
| 224 | + _engram_grad_w_reduce_fake(*fake) | ||