已合并
[fix] Fix NPU pipeline None placeholder Error #40004
yuheng_wang创建于 7月3日
[fix] Fix NPU pipeline None placeholder Error #40004
已合并
共 2 个文件变更+130-10
| @@ -59,6 +59,28 @@ class TinyTransformer(torch.nn.Module): | |||
| 59 | return self.output(self.norm(x)) | 59 | return self.output(self.norm(x)) |
| 60 | 60 | ||
| 61 | 61 | ||
| 62 | +class NoneMaskBlock(torch.nn.Module): | ||
| 63 | + def __init__(self) -> None: | ||
| 64 | + super().__init__() | ||
| 65 | + self.linear = torch.nn.Linear(d_hid, d_hid) | ||
| 66 | + | ||
| 67 | + def forward(self, x, mask=None): | ||
| 68 | + assert mask is None | ||
| 69 | + return self.linear(x) | ||
| 70 | + | ||
| 71 | + | ||
| 72 | +class NoneInputPipelineModel(torch.nn.Module): | ||
| 73 | + def __init__(self) -> None: | ||
| 74 | + super().__init__() | ||
| 75 | + self.block0 = NoneMaskBlock() | ||
| 76 | + self.block1 = NoneMaskBlock() | ||
| 77 | + self.split_spec = {"block1": SplitPoint.BEGINNING} | ||
| 78 | + | ||
| 79 | + def forward(self, x, mask=None): | ||
| 80 | + x = self.block0(x, mask) | ||
| 81 | + return self.block1(x, mask) | ||
| 82 | + | ||
| 83 | + | ||
| 62 | class FXTracerPublicApiTest(TestCase): | 84 | class FXTracerPublicApiTest(TestCase): |
| 63 | def test_public_pipeline_is_patched(self): | 85 | def test_public_pipeline_is_patched(self): |
| 64 | import torch.distributed.pipelining as pipelining | 86 | import torch.distributed.pipelining as pipelining |
| @@ -70,7 +92,6 @@ class FXTracerPublicApiTest(TestCase): | |||
| 70 | self.assertIn("atomic_units", signature.parameters) | 92 | self.assertIn("atomic_units", signature.parameters) |
| 71 | self.assertEqual(pipeline.__module__, "torch_npu.distributed.pipelining") | 93 | self.assertEqual(pipeline.__module__, "torch_npu.distributed.pipelining") |
| 72 | 94 | ||
| 73 | - | ||
| 74 | class FXTracerScheduleTest(TestCase): | 95 | class FXTracerScheduleTest(TestCase): |
| 75 | MAIN_PROCESS_RANK = -1 | 96 | MAIN_PROCESS_RANK = -1 |
| 76 | world_size = 2 | 97 | world_size = 2 |
| @@ -183,6 +204,23 @@ class FXTracerScheduleTest(TestCase): | |||
| 183 | stage = pipe.build_stage(self.rank, self.device) | 204 | stage = pipe.build_stage(self.rank, self.device) |
| 184 | return model, pipe, stage, x | 205 | return model, pipe, stage, x |
| 185 | 206 | ||
| 207 | + def _build_none_input_pipe_and_stage(self): | ||
| 208 | + self._seed_all(2028) | ||
| 209 | + self._set_device() | ||
| 210 | + model = NoneInputPipelineModel().to(self.device) | ||
| 211 | + x = torch.randn(batch_size, d_hid, device=self.device) | ||
| 212 | + x_mb = x.chunk(chunks)[0] | ||
| 213 | + | ||
| 214 | + pipe = pipeline( | ||
| 215 | + module=model, | ||
| 216 | + mb_args=(x_mb, None), | ||
| 217 | + split_spec=model.split_spec, | ||
| 218 | + mode="fx", | ||
| 219 | + atomic_units=["block0", "block1"], | ||
| 220 | + ) | ||
| 221 | + stage = pipe.build_stage(self.rank, self.device) | ||
| 222 | + return pipe, stage, x | ||
| 223 | + | ||
| 186 | def _schedule_step(self, schedule, *args, **kwargs): | 224 | def _schedule_step(self, schedule, *args, **kwargs): |
| 187 | original_fork_rng = torch.random.fork_rng | 225 | original_fork_rng = torch.random.fork_rng |
| 188 | 226 | ||
| @@ -266,6 +304,35 @@ class FXTracerScheduleTest(TestCase): | |||
| 266 | 304 | ||
| 267 | self.assertLessEqual(fx_elapsed, export_elapsed * 1.2) | 305 | self.assertLessEqual(fx_elapsed, export_elapsed * 1.2) |
| 268 | 306 | ||
| 307 | + | ||
| 308 | + def test_fx_schedule_accepts_none_input(self): | ||
| 309 | + self.dist_init() | ||
| 310 | + pipe, stage, x = self._build_none_input_pipe_and_stage() | ||
| 311 | + self.assertEqual(pipe.num_stages, self.world_size) | ||
| 312 | + | ||
| 313 | + for node in pipe.split_gm.graph.nodes: | ||
| 314 | + if node.op != "call_module": | ||
| 315 | + continue | ||
| 316 | + stage_module = pipe.split_gm.get_submodule(node.target) | ||
| 317 | + placeholders = [ | ||
| 318 | + stage_node | ||
| 319 | + for stage_node in stage_module.graph.nodes | ||
| 320 | + if stage_node.op == "placeholder" | ||
| 321 | + ] | ||
| 322 | + self.assertEqual(len(node.args), len(placeholders)) | ||
| 323 | + self.assertTrue( | ||
| 324 | + all( | ||
| 325 | + placeholder.meta.get("val", "__missing__") is not None | ||
| 326 | + for placeholder in placeholders | ||
| 327 | + ) | ||
| 328 | + ) | ||
| 329 | + | ||
| 330 | + loss_fn = torch.nn.MSELoss(reduction="sum") | ||
| 331 | + schedule = Schedule1F1B(stage, chunks, loss_fn=loss_fn, scale_grads=False) | ||
| 332 | + target = torch.randn(batch_size, d_hid, device=self.device) | ||
| 333 | + losses = [] | ||
| 334 | + self._schedule_step(schedule, x, target=target, losses=losses) | ||
| 335 | + | ||
| 269 | 336 | ||
| 270 | def test_fx_tracer_builds_pipe_and_stage(self): | 337 | def test_fx_tracer_builds_pipe_and_stage(self): |
| 271 | self.dist_init() | 338 | self.dist_init() |
| @@ -3,6 +3,7 @@ from __future__ import annotations | |||
| 3 | 3 | ||
| 4 | import contextlib | 4 | import contextlib |
| 5 | import importlib | 5 | import importlib |
| 6 | +import warnings | ||
| 6 | from typing import Any, Optional | 7 | from typing import Any, Optional |
| 7 | 8 | ||
| 8 | import torch | 9 | import torch |
| @@ -96,24 +97,76 @@ def _prune_none_root_placeholders(pipe) -> None: | |||
| 96 | 97 | ||
| 97 | if not hasattr(pipe, "get_stage_module") or not hasattr(pipe, "num_stages"): | 98 | if not hasattr(pipe, "get_stage_module") or not hasattr(pipe, "num_stages"): |
| 98 | return | 99 | return |
| 100 | + if not hasattr(pipe, "split_gm") or not isinstance(pipe.split_gm, fx.GraphModule): | ||
| 101 | + return | ||
| 99 | 102 | ||
| 103 | + root_graph_changed = False | ||
| 100 | for stage_index in range(pipe.num_stages): | 104 | for stage_index in range(pipe.num_stages): |
| 101 | stage_module = pipe.get_stage_module(stage_index) | 105 | stage_module = pipe.get_stage_module(stage_index) |
| 102 | if not isinstance(stage_module, fx.GraphModule): | 106 | if not isinstance(stage_module, fx.GraphModule): |
| 103 | continue | 107 | continue |
| 104 | 108 | ||
| 105 | graph = stage_module.graph | 109 | graph = stage_module.graph |
| 106 | - changed = False | ||
| 107 | placeholders = [node for node in graph.nodes if node.op == "placeholder"] | 110 | placeholders = [node for node in graph.nodes if node.op == "placeholder"] |
| 108 | - for placeholder in placeholders: | 111 | + remove_indices = [ |
| 109 | - if placeholder.meta.get("val", "__missing__") is None: | 112 | + placeholder_index |
| 110 | - placeholder.replace_all_uses_with(None) | 113 | + for placeholder_index, placeholder in enumerate(placeholders) |
| 111 | - graph.erase_node(placeholder) | 114 | + if placeholder.meta.get("val", "__missing__") is None |
| 112 | - changed = True | 115 | + ] |
| 116 | + if not remove_indices: | ||
| 117 | + continue | ||
| 113 | 118 | ||
| 114 | - if changed: | 119 | + root_call_node = None |
| 115 | - graph.lint() | 120 | + for node in pipe.split_gm.graph.nodes: |
| 116 | - stage_module.recompile() | 121 | + if node.op != "call_module": |
| 122 | + continue | ||
| 123 | + | ||
| 124 | + try: | ||
| 125 | + called_module = pipe.split_gm.get_submodule(node.target) | ||
| 126 | + except AttributeError: | ||
| 127 | + continue | ||
| 128 | + | ||
| 129 | + if called_module is stage_module: | ||
| 130 | + root_call_node = node | ||
| 131 | + break | ||
| 132 | + | ||
| 133 | + if root_call_node is None: | ||
| 134 | + warnings.warn( | ||
| 135 | + "Failed to find root graph call_module node for pipeline " | ||
| 136 | + f"stage {stage_index}. Skip pruning None placeholders for " | ||
| 137 | + "this stage to avoid graph inconsistency.", | ||
| 138 | + RuntimeWarning, | ||
| 139 | + stacklevel=2, | ||
| 140 | + ) | ||
| 141 | + continue | ||
| 142 | + if max(remove_indices) >= len(root_call_node.args): | ||
| 143 | + warnings.warn( | ||
| 144 | + "Root graph call_module args do not match pipeline stage " | ||
| 145 | + f"placeholders for stage {stage_index}. Skip pruning None " | ||
| 146 | + "placeholders for this stage to avoid graph inconsistency.", | ||
| 147 | + RuntimeWarning, | ||
| 148 | + stacklevel=2, | ||
| 149 | + ) | ||
| 150 | + continue | ||
| 151 | + for placeholder_index in remove_indices: | ||
| 152 | + placeholder = placeholders[placeholder_index] | ||
| 153 | + placeholder.replace_all_uses_with(None) | ||
| 154 | + graph.erase_node(placeholder) | ||
| 155 | + | ||
| 156 | + graph.lint() | ||
| 157 | + stage_module.recompile() | ||
| 158 | + | ||
| 159 | + remove_index_set = set(remove_indices) | ||
| 160 | + root_call_node.args = tuple( | ||
| 161 | + arg | ||
| 162 | + for arg_index, arg in enumerate(root_call_node.args) | ||
| 163 | + if arg_index not in remove_index_set | ||
| 164 | + ) | ||
| 165 | + root_graph_changed = True | ||
| 166 | + | ||
| 167 | + if root_graph_changed: | ||
| 168 | + pipe.split_gm.graph.lint() | ||
| 169 | + pipe.split_gm.recompile() | ||
| 117 | 170 | ||
| 118 | 171 | ||
| 119 | def _quiet_modify_graph_op_device( | 172 | def _quiet_modify_graph_op_device( |