Star847
1325
代码
代码
Issues1503
Pull Requests2277
流水线
项目讨论
Wiki
分析
Star847
1325
已合并
【bug-fix】修复torch_npu2.9.0版本分布式用例未适配torch2.9.0的问题 #32159
xiaoqi-zhou创建于 3月21日
【bug-fix】修复torch_npu2.9.0版本分布式用例未适配torch2.9.0的问题 #32159
已合并
xiaoqi-zhou创建于 3月21日
从xiaoqi-zhou/pytorch_ascenddistributed_test_v2.9.0合入到Ascend/pytorchv2.9.0
共 4 个文件变更+1033-756
Mtest/distributed/pipelining/model_registry.py+83-13
@@ -1,19 +1,34 @@
1# Copyright (c) Meta Platforms, Inc. and affiliates1# Copyright (c) Meta Platforms, Inc. and affiliates
2# Owner(s): ["oncall: distributed"]2# Owner(s): ["oncall: distributed"]
3# This file is a model zoo for testing torch.distributed.pipelining.3# This file is a model zoo for testing torch.distributed.pipelining.
4+# Licensed under the BSD 3-Clause License (the "License");
5+# you may not use this file except in compliance with the License.
6+# You may obtain a copy of the License at
7+#
8+# https://github.com/pytorch/pytorch/blob/main/LICENSE
9+#
10+# Unless required by applicable law or agreed to in writing, software
11+# distributed under the License is distributed on an "AS IS" BASIS,
12+# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
13+# See the License for the specific language governing permissions and
14+# limitations under the License.
4import torch15import torch
5from torch.autograd import Function16from torch.autograd import Function
6from torch.distributed.pipelining import pipe_split, SplitPoint17from torch.distributed.pipelining import pipe_split, SplitPoint
7 18 
8 19 
9class ExampleCode(torch.nn.Module):20class ExampleCode(torch.nn.Module):
10- def __init__(self, d_hid):21+ def __init__(self, d_hid, splits=2):
22+ if splits > 4:
23+ raise ValueError(f"splits must be <= 4, got {splits}")
11 super().__init__()24 super().__init__()
25+ self.splits = splits
12 self.mm_param0 = torch.nn.Parameter(torch.randn(d_hid, d_hid))26 self.mm_param0 = torch.nn.Parameter(torch.randn(d_hid, d_hid))
13 self.mm_param1 = torch.nn.Parameter(torch.randn(d_hid, d_hid))27 self.mm_param1 = torch.nn.Parameter(torch.randn(d_hid, d_hid))
14 self.cval = torch.nn.Buffer(torch.randn((d_hid,), requires_grad=False))28 self.cval = torch.nn.Buffer(torch.randn((d_hid,), requires_grad=False))
15 self.lin0 = torch.nn.Linear(d_hid, d_hid)29 self.lin0 = torch.nn.Linear(d_hid, d_hid)
16 self.lin1 = torch.nn.Linear(d_hid, d_hid)30 self.lin1 = torch.nn.Linear(d_hid, d_hid)
31+ self.lin2 = torch.nn.Linear(d_hid, d_hid)
17 32 
18 def forward(self, x):33 def forward(self, x):
19 x = torch.mm(x, self.mm_param0)34 x = torch.mm(x, self.mm_param0)
@@ -24,8 +39,14 @@ class ExampleCode(torch.nn.Module):
24 pipe_split()39 pipe_split()
25 x = torch.relu(x) + a_constant40 x = torch.relu(x) + a_constant
26 x = torch.mm(x, self.mm_param1)41 x = torch.mm(x, self.mm_param1)
27- x = self.lin1(x)42+ if self.splits > 2:
28- x = torch.relu(x)43+ pipe_split()
44+ x = self.lin1(x)
45+ x = torch.relu(x)
46+ if self.splits > 3:
47+ pipe_split()
48+ x = self.lin2(x)
49+ x = torch.relu(x)
29 return x50 return x
30 51 
31 52 
@@ -33,12 +54,17 @@ class ModelWithKwargs(torch.nn.Module):
33 DEFAULT_DHID = 51254 DEFAULT_DHID = 512
34 DEFAULT_BATCH_SIZE = 25655 DEFAULT_BATCH_SIZE = 256
35 56 
36- def __init__(self, d_hid: int = DEFAULT_DHID):57+ def __init__(self, d_hid: int = DEFAULT_DHID, splits=2):
58+ if splits > 4:
59+ raise ValueError(f"splits must be <= 4, got {splits}")
37 super().__init__()60 super().__init__()
61+ self.splits = splits
38 self.mm_param0 = torch.nn.Parameter(torch.randn(d_hid, d_hid))62 self.mm_param0 = torch.nn.Parameter(torch.randn(d_hid, d_hid))
39 self.mm_param1 = torch.nn.Parameter(torch.randn(d_hid, d_hid))63 self.mm_param1 = torch.nn.Parameter(torch.randn(d_hid, d_hid))
40 self.lin0 = torch.nn.Linear(d_hid, d_hid)64 self.lin0 = torch.nn.Linear(d_hid, d_hid)
41 self.lin1 = torch.nn.Linear(d_hid, d_hid)65 self.lin1 = torch.nn.Linear(d_hid, d_hid)
66+ self.lin2 = torch.nn.Linear(d_hid, d_hid)
67+ self.lin3 = torch.nn.Linear(d_hid, d_hid)
42 68 
43 def forward(self, x, y=torch.zeros(DEFAULT_BATCH_SIZE, DEFAULT_DHID)):69 def forward(self, x, y=torch.zeros(DEFAULT_BATCH_SIZE, DEFAULT_DHID)):
44 x = torch.mm(x, self.mm_param0)70 x = torch.mm(x, self.mm_param0)
@@ -49,6 +75,14 @@ class ModelWithKwargs(torch.nn.Module):
49 x = torch.mm(x, self.mm_param1)75 x = torch.mm(x, self.mm_param1)
50 x = self.lin1(x)76 x = self.lin1(x)
51 x = torch.relu(x)77 x = torch.relu(x)
78+ if self.splits > 2:
79+ pipe_split()
80+ x = self.lin2(x)
81+ x = torch.relu(x)
82+ if self.splits > 3:
83+ pipe_split()
84+ x = self.lin3(x)
85+ x = torch.relu(x)
52 return x86 return x
53 87 
54 88 
@@ -88,13 +122,30 @@ class MLPModule(torch.nn.Module):
88 return x122 return x
89 123 
90 124 
125+class MLPKWargModule(torch.nn.Module):
126+ def __init__(self, d_hid: int, layer_num):
127+ super().__init__()
128+ self.net1 = torch.nn.Linear(d_hid, d_hid)
129+ self.relu = torch.nn.ReLU()
130+ self.net2 = torch.nn.Linear(d_hid, d_hid)
131+ self.layer_num = layer_num
132+ 
133+ def forward(self, x, unused_kwarg: torch.Tensor = torch.zeros(1)):
134+ x = self.net1(x)
135+ x = self.relu(x)
136+ x = self.net2(x)
137+ return x
138+ 
139+ 
91# Multi-MLP model140# Multi-MLP model
92class MultiMLP(torch.nn.Module):141class MultiMLP(torch.nn.Module):
93 def __init__(self, d_hid: int, n_layers: int = 2):142 def __init__(self, d_hid: int, n_layers: int = 2):
94 super().__init__()143 super().__init__()
95 self.layers = torch.nn.ModuleList([MLPModule(d_hid) for _ in range(n_layers)])144 self.layers = torch.nn.ModuleList([MLPModule(d_hid) for _ in range(n_layers)])
96 # For testing purpose only, this should be defined by user145 # For testing purpose only, this should be defined by user
97- self.split_spec = {f"layers.{i}": SplitPoint.BEGINNING for i in range(1, n_layers)}146+ self.split_spec = {
147+ f"layers.{i}": SplitPoint.BEGINNING for i in range(1, n_layers)
148+ }
98 149 
99 def forward(self, x):150 def forward(self, x):
100 for layer in self.layers:151 for layer in self.layers:
@@ -102,9 +153,26 @@ class MultiMLP(torch.nn.Module):
102 return x153 return x
103 154 
104 155 
156+# Multi-MLP with kwargs model
157+class MultiMLPKwargs(torch.nn.Module):
158+ def __init__(self, d_hid: int, n_layers: int = 2):
159+ super().__init__()
160+ self.layers = torch.nn.ModuleList(
161+ [MLPKWargModule(d_hid, i) for i in range(n_layers)]
162+ )
163+ # For testing purpose only, this should be defined by user
164+ self.split_spec = {
165+ f"layers.{i}": SplitPoint.BEGINNING for i in range(1, n_layers)
166+ }
167+ 
168+ def forward(self, x, unused_kwarg: torch.Tensor = torch.zeros(1)):
169+ for layer in self.layers:
170+ x = layer(x)
171+ return x
172+ 
173+ 
105class CustomLinearDx(Function):174class CustomLinearDx(Function):
106 @staticmethod175 @staticmethod
107- # pylint:disable=huawei-too-many-arguments
108 def forward(ctx, input_val, weight, bias, module, layer_idx):176 def forward(ctx, input_val, weight, bias, module, layer_idx):
109 ctx.save_for_backward(input_val, weight, bias)177 ctx.save_for_backward(input_val, weight, bias)
110 ctx.module = module178 ctx.module = module
@@ -113,7 +181,7 @@ class CustomLinearDx(Function):
113 181 
114 @staticmethod182 @staticmethod
115 def backward(ctx, grad_output):183 def backward(ctx, grad_output):
116- input_val, weight, bias = ctx.saved_tensors184+ input_val, weight, _ = ctx.saved_tensors
117 grad_input = grad_output.mm(weight)185 grad_input = grad_output.mm(weight)
118 ctx.module.cached_context[ctx.layer_idx].append(grad_output.clone())186 ctx.module.cached_context[ctx.layer_idx].append(grad_output.clone())
119 ctx.module.cached_context[str(ctx.layer_idx) + "_input"].append(187 ctx.module.cached_context[str(ctx.layer_idx) + "_input"].append(
@@ -130,7 +198,7 @@ class CustomLinearDxDw(Function):
130 198 
131 @staticmethod199 @staticmethod
132 def backward(ctx, grad_output):200 def backward(ctx, grad_output):
133- input_val, weight, bias = ctx.saved_tensors201+ input_val, weight, _ = ctx.saved_tensors
134 grad_input = grad_output.mm(weight)202 grad_input = grad_output.mm(weight)
135 grad_weight = grad_output.t().mm(input_val)203 grad_weight = grad_output.t().mm(input_val)
136 grad_bias = grad_output.sum(0)204 grad_bias = grad_output.sum(0)
@@ -145,10 +213,10 @@ class MLPModuleWithDw(torch.nn.Module):
145 self.fc2_weight = torch.nn.Parameter(torch.randn(d_hid, d_hid))213 self.fc2_weight = torch.nn.Parameter(torch.randn(d_hid, d_hid))
146 self.fc2_bias = torch.nn.Parameter(torch.randn(d_hid))214 self.fc2_bias = torch.nn.Parameter(torch.randn(d_hid))
147 215 
148- torch.nn.init.uniform_(self.fc1_weight, -0.01, 0.01)216+ torch.nn.init.uniform_(self.fc1_weight, -0.001, 0.001)
149- torch.nn.init.uniform_(self.fc2_weight, -0.01, 0.01)217+ torch.nn.init.uniform_(self.fc2_weight, -0.001, 0.001)
150- torch.nn.init.uniform_(self.fc1_bias, -0.01, 0.01)218+ torch.nn.init.uniform_(self.fc1_bias, -0.001, 0.001)
151- torch.nn.init.uniform_(self.fc2_bias, -0.01, 0.01)219+ torch.nn.init.uniform_(self.fc2_bias, -0.001, 0.001)
152 220 
153 self.cached_context = {}221 self.cached_context = {}
154 self.cached_context["fc1"] = []222 self.cached_context["fc1"] = []
@@ -209,7 +277,9 @@ class MultiMLPWithDw(torch.nn.Module):
209 [MLPModuleWithDw(d_hid) for _ in range(n_layers)]277 [MLPModuleWithDw(d_hid) for _ in range(n_layers)]
210 )278 )
211 # For testing purpose only, this should be defined by user279 # For testing purpose only, this should be defined by user
212- self.split_spec = {f"layers.{i}": SplitPoint.BEGINNING for i in range(1, n_layers)}280+ self.split_spec = {
281+ f"layers.{i}": SplitPoint.BEGINNING for i in range(1, n_layers)
282+ }
213 self.use_custom_logic = False283 self.use_custom_logic = False
214 284 
215 def forward(self, x):285 def forward(self, x):
Mtest/distributed/pipelining/schedule_registry.py+21-10
@@ -2,6 +2,17 @@
2# Owner(s): ["oncall: distributed"]2# Owner(s): ["oncall: distributed"]
3# This file is a Schedule zoo for testing torch.distributed.pipelining.3# This file is a Schedule zoo for testing torch.distributed.pipelining.
4# It includes schedules designed purely for testing purposes4# It includes schedules designed purely for testing purposes
5+# Licensed under the BSD 3-Clause License (the "License");
6+# you may not use this file except in compliance with the License.
7+# You may obtain a copy of the License at
8+#
9+# https://github.com/pytorch/pytorch/blob/main/LICENSE
10+#
11+# Unless required by applicable law or agreed to in writing, software
12+# distributed under the License is distributed on an "AS IS" BASIS,
13+# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
14+# See the License for the specific language governing permissions and
15+# limitations under the License.
5from typing import Callable, Optional16from typing import Callable, Optional
6 17 
7from torch.distributed.pipelining.schedules import (18from torch.distributed.pipelining.schedules import (
@@ -20,7 +31,7 @@ from torch.distributed.pipelining.stage import _PipelineStageBase
20F = _ComputationType.FORWARD31F = _ComputationType.FORWARD
21B = _ComputationType.FULL_BACKWARD32B = _ComputationType.FULL_BACKWARD
22W = _ComputationType.BACKWARD_WEIGHT33W = _ComputationType.BACKWARD_WEIGHT
23-INPUT = _ComputationType.BACKWARD_INPUT34+I = _ComputationType.BACKWARD_INPUT
24 35 
25 36 
26class ScheduleVShaped(PipelineScheduleMulti):37class ScheduleVShaped(PipelineScheduleMulti):
@@ -45,7 +56,7 @@ class ScheduleVShaped(PipelineScheduleMulti):
45 )56 )
46 57 
47 # Go through one microbatch58 # Go through one microbatch
48- # Note(whc) - it might be easier to work with thes schedules by writing them as a list of59+ # Note(whc) - it might be easier to work with this schedules by writing them as a list of
49 # ["0F0", ...] and then parsing them in the test infra to turn them into actions.60 # ["0F0", ...] and then parsing them in the test infra to turn them into actions.
50 self.pipeline_order = {61 self.pipeline_order = {
51 0: [62 0: [
@@ -156,12 +167,12 @@ class ScheduleWithW(PipelineScheduleMulti):
156 _Action(2, F, 0),167 _Action(2, F, 0),
157 _Action(2, F, 1),168 _Action(2, F, 1),
158 None,169 None,
159- _Action(2, INPUT, 0),170+ _Action(2, I, 0),
160 _Action(2, W, 0),171 _Action(2, W, 0),
161- _Action(0, INPUT, 0),172+ _Action(0, I, 0),
162- _Action(2, INPUT, 1),173+ _Action(2, I, 1),
163 _Action(0, W, 0),174 _Action(0, W, 0),
164- _Action(0, INPUT, 1),175+ _Action(0, I, 1),
165 _Action(2, W, 1),176 _Action(2, W, 1),
166 _Action(0, W, 1),177 _Action(0, W, 1),
167 ],178 ],
@@ -170,12 +181,12 @@ class ScheduleWithW(PipelineScheduleMulti):
170 _Action(1, F, 0),181 _Action(1, F, 0),
171 _Action(1, F, 1),182 _Action(1, F, 1),
172 _Action(3, F, 0),183 _Action(3, F, 0),
173- _Action(3, INPUT, 0),184+ _Action(3, I, 0),
174 _Action(3, F, 1),185 _Action(3, F, 1),
175- _Action(1, INPUT, 0),186+ _Action(1, I, 0),
176- _Action(3, INPUT, 1),187+ _Action(3, I, 1),
177 _Action(3, W, 0),188 _Action(3, W, 0),
178- _Action(1, INPUT, 1),189+ _Action(1, I, 1),
179 _Action(1, W, 0),190 _Action(1, W, 0),
180 _Action(3, W, 1),191 _Action(3, W, 1),
181 _Action(1, W, 1),192 _Action(1, W, 1),
Mtest/distributed/pipelining/test_schedule_multiproc.py+752-646
@@ -1,12 +1,22 @@
1# Copyright (c) Meta Platforms, Inc. and affiliates1# Copyright (c) Meta Platforms, Inc. and affiliates
2# Owner(s): ["oncall: distributed"]2# Owner(s): ["oncall: distributed"]
3+# Licensed under the BSD 3-Clause License (the "License");
4+# you may not use this file except in compliance with the License.
5+# You may obtain a copy of the License at
6+#
7+# https://github.com/pytorch/pytorch/blob/main/LICENSE
8+#
9+# Unless required by applicable law or agreed to in writing, software
10+# distributed under the License is distributed on an "AS IS" BASIS,
11+# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
12+# See the License for the specific language governing permissions and
13+# limitations under the License.
3import copy14import copy
4import logging15import logging
5-import os
6-import sys
7import tempfile16import tempfile
17+from dataclasses import dataclass
8 18 
9-from model_registry import ModelWithKwargs, MultiMLP, MultiMLPWithDw19+from model_registry import ModelWithKwargs, MultiMLP, MultiMLPKwargs, MultiMLPWithDw
10from schedule_registry import (20from schedule_registry import (
11 ScheduleUnbalanced,21 ScheduleUnbalanced,
12 ScheduleVShaped,22 ScheduleVShaped,
@@ -21,6 +31,7 @@ from torch.distributed.pipelining import (
21 pipeline,31 pipeline,
22 PipelineStage,32 PipelineStage,
23 Schedule1F1B,33 Schedule1F1B,
34+ ScheduleDualPipeV,
24 ScheduleGPipe,35 ScheduleGPipe,
25 ScheduleInterleaved1F1B,36 ScheduleInterleaved1F1B,
26 ScheduleInterleavedZeroBubble,37 ScheduleInterleavedZeroBubble,
@@ -28,73 +39,284 @@ from torch.distributed.pipelining import (
28 ScheduleZBVZeroBubble,39 ScheduleZBVZeroBubble,
29)40)
30from torch.distributed.pipelining.schedules import _PipelineScheduleRuntime41from torch.distributed.pipelining.schedules import _PipelineScheduleRuntime
31-from torch.testing._internal.common_cuda import TEST_MULTIGPU42+from torch.nn.modules.loss import MSELoss
32from torch.testing._internal.common_distributed import (43from torch.testing._internal.common_distributed import (
33- MultiProcContinousTest,44+ MultiProcContinuousTest,
34- requires_nccl,45+ requires_accelerator_dist_backend,
35)46)
36from torch.testing._internal.common_utils import (47from torch.testing._internal.common_utils import (
37 check_leaked_tensors,48 check_leaked_tensors,
38 instantiate_parametrized_tests,49 instantiate_parametrized_tests,
39 parametrize,50 parametrize,
51+ run_tests,
40 skip_but_pass_in_sandcastle_if,52 skip_but_pass_in_sandcastle_if,
41)53)
42 54 
43- 
44logger = logging.getLogger(__name__)55logger = logging.getLogger(__name__)
45 56 
46d_hid = 51257d_hid = 512
47-batch_size = 25658+batch_size = 64
48- 
49torch.manual_seed(0)59torch.manual_seed(0)
60+device_type = acc.type if (acc := torch.accelerator.current_accelerator()) else "cpu"
61+backend = dist.get_default_backend_for_device(device_type)
62+TEST_MULTIACCELERATOR = torch.accelerator.device_count() >= 2
50 63 
51 64 
52-class ScheduleTest(MultiProcContinousTest):65+@dataclass
66+class PipelineTestConfig:
67+ world_size: int
68+ device: torch.device
69+ rank: int
70+ 
71+ 
72+def setup_models_and_data(
73+ config: PipelineTestConfig, n_layers=None, model_class=MultiMLP
74+):
75+ """Setup models, input data, target data, and loss function."""
76+ if n_layers is None:
77+ n_layers = config.world_size
78+ 
79+ full_mod = model_class(d_hid, n_layers=n_layers)
80+ full_mod.to(config.device)
81+ ref_mod = copy.deepcopy(full_mod)
82+ 
83+ x = torch.randn(batch_size, d_hid, device=config.device)
84+ with torch.no_grad():
85+ y = ref_mod(x)
86+ target = y + torch.randn(batch_size, d_hid, device=config.device)
87+ 
88+ loss_fn = torch.nn.MSELoss(reduction="sum")
89+ return full_mod, ref_mod, x, target, loss_fn
90+ 
91+ 
92+def create_single_stage_pipeline(
93+ config: PipelineTestConfig, mod, x, chunks, use_tracer=True
94+):
95+ """Create a single-stage pipeline using either tracer or manual stage creation."""
96+ if use_tracer:
97+ x_mb = x.chunk(chunks)[0]
98+ split_spec = mod.split_spec if hasattr(mod, "split_spec") else None
99+ pipe = pipeline(mod, mb_args=(x_mb,), split_spec=split_spec)
100+ stage = pipe.build_stage(config.rank, config.device)
101+ stage_module = pipe.get_stage_module(config.rank)
102+ return stage, stage_module, [stage_module]
103+ else:
104+ # Manual stage creation
105+ submod_name = f"layers.{config.rank}"
106+ stage_module = mod.get_submodule(submod_name)
107+ stage = PipelineStage(
108+ stage_module, config.rank, config.world_size, config.device
109+ )
110+ return stage, stage_module, [stage_module]
111+ 
112+ 
113+def create_multi_stage_pipeline(
114+ config: PipelineTestConfig, mod, stages_per_rank, n_stages, stage_indices=None
115+):
116+ """Create multiple pipeline stages for interleaved schedules."""
117+ if stage_indices is None:
118+ stage_indices = [
119+ config.rank + i * config.world_size for i in range(stages_per_rank)
120+ ]
121+ 
122+ submod_names = [f"layers.{i}" for i in stage_indices]
123+ stage_modules = [mod.get_submodule(submod_name) for submod_name in submod_names]
124+ stages = [
125+ PipelineStage(stage_module, stage_idx, n_stages, config.device)
126+ for stage_module, stage_idx in zip(stage_modules, stage_indices, strict=True)
127+ ]
128+ return stages, stage_modules, submod_names
129+ 
130+ 
131+def run_reference_model(ref_mod, x, target, loss_fn, num_iterations=2, **kwargs):
132+ """Run reference model for specified iterations and return final output and loss."""
133+ ref_out = None
134+ ref_loss = None
135+ 
136+ for _ in range(num_iterations):
137+ ref_mod.zero_grad()
138+ ref_out = ref_mod(x, **kwargs)
139+ ref_loss = loss_fn(ref_out, target)
140+ ref_loss.backward()
141+ 
142+ return ref_out, ref_loss
143+ 
144+ 
145+def check_gradients(
146+ config: PipelineTestConfig,
147+ stage_modules,
148+ ref_mod,
149+ submod_names=None,
150+ rtol=1e-5,
151+ atol=4e-5,
152+):
153+ """Check that gradients match between pipeline stages and reference model using flexible comparison."""
154+ 
155+ def grad_check(grad1, grad2, param_name, rtol, atol, tolerance=0.05):
156+ if grad1 is None and grad2 is None:
157+ return
158+ if grad1 is None or grad2 is None:
159+ raise AssertionError(
160+ f"One gradient is None for {param_name}: {grad1} vs {grad2}"
161+ )
162+ try:
163+ torch.testing.assert_close(grad1, grad2, rtol=rtol, atol=atol)
164+ except AssertionError:
165+ print(
166+ f"Numerical issues detected for {param_name}: param grad {grad1} vs ref grad {grad2}"
167+ )
168+ raise
169+ 
170+ if submod_names is None:
171+ # Single stage case - need to detect tracer vs manual pipeline
172+ stage_modules = [stage_modules]
173+ 
174+ # Try to detect if this is a tracer-based pipeline by checking if parameter exists in ref_mod
175+ sample_param_name = next(iter(stage_modules[0].named_parameters()))[0]
176+ try:
177+ # Try to get parameter directly from reference model (tracer-based)
178+ ref_mod.get_parameter(sample_param_name)
179+ is_tracer_based = True
180+ except AttributeError:
181+ # Parameter doesn't exist at root level, must be manual pipeline
182+ is_tracer_based = False
183+ 
184+ if is_tracer_based:
185+ # Tracer-based pipeline: parameter names are full paths from root model
186+ for name, p in stage_modules[0].named_parameters():
187+ ref_p = ref_mod.get_parameter(name)
188+ grad_check(p.grad, ref_p.grad, name, rtol, atol)
189+ else:
190+ # Manual pipeline: parameter names are local to the submodule
191+ submod_name = f"layers.{config.rank}"
192+ ref_submod = ref_mod.get_submodule(submod_name)
193+ for name, p in stage_modules[0].named_parameters():
194+ ref_p = ref_submod.get_parameter(name)
195+ grad_check(p.grad, ref_p.grad, f"{submod_name}.{name}", rtol, atol)
196+ else:
197+ # Multi-stage case - always use submodule approach
198+ for stage_module, submod_name in zip(stage_modules, submod_names):
199+ ref_submod = ref_mod.get_submodule(submod_name)
200+ for name, p in stage_module.named_parameters():
201+ ref_p = ref_submod.get_parameter(name)
202+ grad_check(p.grad, ref_p.grad, f"{submod_name}.{name}", rtol, atol)
203+ 
204+ 
205+def zero_gradients(stage_modules):
206+ """Zero gradients for all stage modules."""
207+ if not isinstance(stage_modules, list):
208+ stage_modules = [stage_modules]
209+ for stage_module in stage_modules:
210+ stage_module.zero_grad()
211+ 
212+ 
213+from contextlib import contextmanager
214+ 
215+ 
216+@contextmanager
217+def patch_stage_init_method(stages):
218+ """Context manager to temporarily patch stage methods"""
219+ original_methods = []
220+ patched_methods = []
221+ 
222+ for stage in stages:
223+ original = stage._get_init_p2p_neighbors_ops
224+ 
225+ def create_patched_method(stage_self):
226+ ops = []
227+ next_stage_peer_rank = stage_self.stage_index_to_group_rank.get(stage_self.stage_index + 1)
228+ prev_stage_peer_rank = stage_self.stage_index_to_group_rank.get(stage_self.stage_index - 1)
229+ 
230+ recv_tensor = torch.zeros(1, dtype=torch.float32, device=stage_self.device)
231+ send_tensor = torch.tensor(stage_self.stage_index, dtype=torch.float32, device=stage_self.device)
232+ 
233+ # Forward
234+ if not stage_self.is_first:
235+ ops.append(
236+ dist.P2POp(
237+ dist.irecv,
238+ recv_tensor,
239+ group_peer=prev_stage_peer_rank,
240+ group=stage_self.group,
241+ )
242+ )
243+ if not stage_self.is_last:
244+ ops.append(
245+ dist.P2POp(
246+ dist.isend,
247+ send_tensor,
248+ group_peer=next_stage_peer_rank,
249+ group=stage_self.group,
250+ )
251+ )
252+ 
253+ # Backward
254+ if not stage_self.is_first:
255+ ops.append(
256+ dist.P2POp(
257+ dist.isend,
258+ send_tensor,
259+ group_peer=prev_stage_peer_rank,
260+ group=stage_self.group,
261+ )
262+ )
263+ if not stage_self.is_last:
264+ ops.append(
265+ dist.P2POp(
266+ dist.irecv,
267+ recv_tensor,
268+ group_peer=next_stage_peer_rank,
269+ group=stage_self.group,
270+ )
271+ )
272+ 
273+ return ops
274+ 
275+ patched = create_patched_method.__get__(stage, type(stage))
276+ stage._get_init_p2p_neighbors_ops = patched
277+ original_methods.append(original)
278+ patched_methods.append(patched)
279+ 
280+ try:
281+ yield
282+ finally:
283+ # Restore original methods
284+ for stage, original in zip(stages, original_methods):
285+ stage._get_init_p2p_neighbors_ops = original
286+ 
287+ 
288+class ScheduleTest(MultiProcContinuousTest):
289+ world_size = 4
290+ 
53 @classmethod291 @classmethod
54 def backend_str(cls) -> str:292 def backend_str(cls) -> str:
55 # Testing with HCCL backend293 # Testing with HCCL backend
56- return "hccl"294+ return backend
57 295 
58- @classmethod296+ @property
59- def setUpClass(cls):297+ def device(self) -> torch.device:
60- """298+ return torch.device(device_type, self.rank)
61- Class-scope test fixture. Run once for entire test class, before any test starts.299+ 
62- Set up the device.300+ @property
63- """301+ def config(self) -> PipelineTestConfig:
64- super().setUpClass()302+ """Lazily create and return the pipeline test configuration."""
65- dev_id = cls.rank % torch.npu.device_count()303+ return PipelineTestConfig(
66- cls.device = torch.device(f"npu:{dev_id}")304+ world_size=self.world_size, device=self.device, rank=self.rank
305+ )
67 306 
68 @parametrize("ScheduleClass", [_ScheduleForwardOnly])307 @parametrize("ScheduleClass", [_ScheduleForwardOnly])
69 def test_forward_only(self, ScheduleClass):308 def test_forward_only(self, ScheduleClass):
70- mod = MultiMLP(d_hid, n_layers=self.world_size)309+ mod, mod_ref, x, _, _ = setup_models_and_data(self.config)
71- mod.to(self.device)
72- 
73- mod_ref = copy.deepcopy(mod)
74- 
75- x = torch.randn(batch_size, d_hid, device=self.device)
76 x_clone = x.clone()310 x_clone = x.clone()
77 311 
78- num_microbatches = 4312+ num_microbatches = 2 * self.world_size
79- x_mb = x.chunk(num_microbatches)[0]313+ stage, _, _ = create_single_stage_pipeline(
80- 314+ self.config, mod, x, num_microbatches
81- # Create a pipeline
82- split_spec = mod.split_spec if hasattr(mod, "split_spec") else None
83- pipe = pipeline(
84- mod,
85- mb_args=(x_mb,),
86- split_spec=split_spec,
87 )315 )
88- 
89- stage = pipe.build_stage(
90- self.rank,
91- self.device,
92- )
93- 
94- # Attach to a schedule
95 schedule = ScheduleClass(stage, num_microbatches, scale_grads=False)316 schedule = ScheduleClass(stage, num_microbatches, scale_grads=False)
96 317 
97- # Run318+ # Run forward-only schedule
319+ out = None
98 num_iters = 20320 num_iters = 20
99 for _ in range(num_iters):321 for _ in range(num_iters):
100 if self.rank == 0:322 if self.rank == 0:
@@ -106,39 +328,90 @@ class ScheduleTest(MultiProcContinousTest):
106 else:328 else:
107 schedule.step()329 schedule.step()
108 330 
109- # Validate pipelined output is the same as reference model331+ # Validate pipelined output matches reference model
110 if self.rank == self.world_size - 1:332 if self.rank == self.world_size - 1:
111 for _ in range(num_iters):333 for _ in range(num_iters):
112 x_clone = mod_ref(x_clone)334 x_clone = mod_ref(x_clone)
113- 
114 torch.testing.assert_close(x_clone, out)335 torch.testing.assert_close(x_clone, out)
115 336 
337+ @parametrize(
338+ "ScheduleClass",
339+ [
340+ ScheduleGPipe,
341+ Schedule1F1B,
342+ ScheduleInterleaved1F1B,
343+ ScheduleLoopedBFS,
344+ ScheduleInterleavedZeroBubble
345+ ],
346+ )
347+ def test_eval_inference_mode(self, ScheduleClass):
348+ num_microbatches = 4
349+ if ScheduleClass in [
350+ ScheduleInterleaved1F1B,
351+ ScheduleLoopedBFS,
352+ ScheduleInterleavedZeroBubble,
353+ ]:
354+ # Multi-stage schedules
355+ stages_per_rank = 2
356+ n_stages = stages_per_rank * self.world_size
357+ mod, _, x, target, loss_fn = setup_models_and_data(
358+ self.config, n_layers=n_stages
359+ )
360+ 
361+ # Create multi-stage pipeline
362+ stages, stage_modules, _ = create_multi_stage_pipeline(
363+ self.config, mod, stages_per_rank, n_stages
364+ )
365+ schedule = ScheduleClass(
366+ stages, num_microbatches, loss_fn=loss_fn, scale_grads=False
367+ )
368+ else:
369+ # Single-stage schedules
370+ mod, _, x, target, loss_fn = setup_models_and_data(self.config)
371+ 
372+ # Create single-stage pipeline
373+ stage, stage_module, _ = create_single_stage_pipeline(
374+ self.config, mod, x, num_microbatches
375+ )
376+ stage_modules = [stage_module]
377+ schedule = ScheduleClass(
378+ stage, num_microbatches, loss_fn=loss_fn, scale_grads=False
379+ )
380+ 
381+ # Clear gradients and run eval
382+ zero_gradients(stage_modules)
383+ losses = []
384+ 
385+ if self.rank == 0:
386+ # Support with and without no_grad()
387+ with torch.no_grad():
388+ schedule.eval(x)
389+ elif self.rank == self.world_size - 1:
390+ schedule.eval(target=target, losses=losses)
391+ else:
392+ schedule.eval()
393+ 
394+ # Check that gradients were NOT computed during eval
395+ grad_computed_eval = any(
396+ param.grad is not None
397+ for stage_module in stage_modules
398+ for param in stage_module.parameters()
399+ )
400+ 
401+ # Verify that gradients were not computed during eval
402+ self.assertFalse(
403+ grad_computed_eval, "Gradients should not be computed during eval()"
404+ )
405+ 
406+ # Verify that losses are still computed during eval
407+ if self.rank == self.world_size - 1:
408+ self.assertTrue(len(losses) > 0, "Losses should be computed during eval()")
409+ 
116 @parametrize("ScheduleClass", [ScheduleGPipe, Schedule1F1B])410 @parametrize("ScheduleClass", [ScheduleGPipe, Schedule1F1B])
117 def test_multi_iter(self, ScheduleClass):411 def test_multi_iter(self, ScheduleClass):
118- mod = MultiMLP(d_hid, n_layers=self.world_size)412+ mod, _, x, target, loss_fn = setup_models_and_data(self.config)
119- mod.to(self.device)
120- 
121- x = torch.randn(batch_size, d_hid, device=self.device)
122- target = torch.randn(batch_size, d_hid, device=self.device)
123- loss_fn = torch.nn.MSELoss(reduction="sum")
124- 
125 chunks = 4413 chunks = 4
126- x_mb = x.chunk(chunks)[0]414+ stage, _, _ = create_single_stage_pipeline(self.config, mod, x, chunks)
127- 
128- # Create a pipeline
129- split_spec = mod.split_spec if hasattr(mod, "split_spec") else None
130- pipe = pipeline(
131- mod,
132- mb_args=(x_mb,),
133- split_spec=split_spec,
134- )
135- 
136- stage = pipe.build_stage(
137- self.rank,
138- self.device,
139- )
140- 
141- # Attach to a schedule
142 schedule = ScheduleClass(stage, chunks, loss_fn=loss_fn, scale_grads=False)415 schedule = ScheduleClass(stage, chunks, loss_fn=loss_fn, scale_grads=False)
143 416 
144 # Run417 # Run
@@ -151,9 +424,11 @@ class ScheduleTest(MultiProcContinousTest):
151 else:424 else:
152 schedule.step()425 schedule.step()
153 426 
427+ dist.barrier(device_ids=[self.rank])
428+ 
154 @parametrize("ScheduleClass", [ScheduleGPipe, Schedule1F1B])429 @parametrize("ScheduleClass", [ScheduleGPipe, Schedule1F1B])
155 def test_kwargs_with_tracer(self, ScheduleClass):430 def test_kwargs_with_tracer(self, ScheduleClass):
156- mod = ModelWithKwargs(d_hid)431+ mod = ModelWithKwargs(d_hid, splits=self.world_size)
157 mod.to(self.device)432 mod.to(self.device)
158 433 
159 x = torch.randn(batch_size, d_hid, device=self.device)434 x = torch.randn(batch_size, d_hid, device=self.device)
@@ -180,15 +455,16 @@ class ScheduleTest(MultiProcContinousTest):
180 schedule = ScheduleClass(stage, chunks, loss_fn=loss_fn, scale_grads=False)455 schedule = ScheduleClass(stage, chunks, loss_fn=loss_fn, scale_grads=False)
181 456 
182 # Run457 # Run
458+ out = None
459+ losses = []
183 if self.rank == 0:460 if self.rank == 0:
184 schedule.step(x, y=y)461 schedule.step(x, y=y)
185 elif self.rank == self.world_size - 1:462 elif self.rank == self.world_size - 1:
186- losses = []
187 out = schedule.step(target=target, losses=losses)463 out = schedule.step(target=target, losses=losses)
188 else:464 else:
189 schedule.step()465 schedule.step()
190 466 
191- dist.barrier()467+ dist.barrier(device_ids=[self.rank])
192 468 
193 # Last rank checks result469 # Last rank checks result
194 if self.rank == self.world_size - 1:470 if self.rank == self.world_size - 1:
@@ -199,160 +475,94 @@ class ScheduleTest(MultiProcContinousTest):
199 torch.testing.assert_close(pipe_loss, ref_loss)475 torch.testing.assert_close(pipe_loss, ref_loss)
200 476 
201 @parametrize("ScheduleClass", [ScheduleGPipe, Schedule1F1B])477 @parametrize("ScheduleClass", [ScheduleGPipe, Schedule1F1B])
202- @parametrize("ModelClass", [MultiMLP])478+ def test_grad_with_tracer(self, ScheduleClass):
203- def test_grad_with_tracer(self, ScheduleClass, ModelClass):479+ mod, ref_mod, x, target, loss_fn = setup_models_and_data(self.config)
204- mod = ModelClass(d_hid)
205- mod.to(self.device)
206- 
207- ref_mod = copy.deepcopy(mod)
208- x = torch.randn(batch_size, d_hid, device=self.device)
209- with torch.no_grad():
210- y = ref_mod(x)
211- # Add a small perturbation
212- target = y + torch.randn(batch_size, d_hid, device=self.device)
213- 
214- loss_fn = torch.nn.MSELoss(reduction="sum")
215 480 
216 # Run reference481 # Run reference
217- for _ in range(2):482+ ref_out, ref_loss = run_reference_model(ref_mod, x, target, loss_fn)
218- ref_mod.zero_grad()
219- ref_out = ref_mod(x)
220- ref_loss = loss_fn(ref_out, target)
221- ref_loss.backward()
222 483 
223- # Create a pipeline484+ # Create pipeline and schedule
224- chunks = 4485+ chunks = 2 * self.world_size
225- x_mb = x.chunk(chunks)[0]486+ stage, stage_module, stage_modules = create_single_stage_pipeline(
226- split_spec = mod.split_spec if hasattr(mod, "split_spec") else None487+ self.config, mod, x, chunks
227- pipe = pipeline(
228- mod,
229- mb_args=(x_mb,),
230- split_spec=split_spec,
231 )488 )
232- 
233- stage = pipe.build_stage(
234- self.rank,
235- self.device,
236- )
237- 
238- # Attach to a schedule
239 schedule = ScheduleClass(stage, chunks, loss_fn=loss_fn, scale_grads=False)489 schedule = ScheduleClass(stage, chunks, loss_fn=loss_fn, scale_grads=False)
240 490 
241- # Run491+ # Run pipeline
242- stage_module = pipe.get_stage_module(self.rank)492+ out = None
493+ losses = []
243 for _ in range(2):494 for _ in range(2):
244- # Zero gradients495+ zero_gradients(stage_module)
245- stage_module.zero_grad()
246 if self.rank == 0:496 if self.rank == 0:
247 schedule.step(x)497 schedule.step(x)
248 elif self.rank == self.world_size - 1:498 elif self.rank == self.world_size - 1:
249- losses = []
250 out = schedule.step(target=target, losses=losses)499 out = schedule.step(target=target, losses=losses)
251 else:500 else:
252 schedule.step()501 schedule.step()
253 502 
254- dist.barrier()503+ dist.barrier(device_ids=[self.rank])
255 504 
256 # Last rank checks result505 # Last rank checks result
257 if self.rank == self.world_size - 1:506 if self.rank == self.world_size - 1:
258- # Check output
259 torch.testing.assert_close(out, ref_out)507 torch.testing.assert_close(out, ref_out)
260- # Check loss
261- # Since the reduction used in the loss function above is "sum", we use
262- # "sum" here to reduce microbatch losses into a single value too.
263 pipe_loss = sum(losses)508 pipe_loss = sum(losses)
264 torch.testing.assert_close(pipe_loss, ref_loss)509 torch.testing.assert_close(pipe_loss, ref_loss)
265 510 
266- # Every rank checks gradients511+ # Check gradients using helper method
267- for name, p in stage_module.named_parameters():512+ check_gradients(self.config, stage_module, ref_mod)
268- ref_p = ref_mod.get_parameter(name)
269- try:
270- torch.testing.assert_close(p.grad, ref_p.grad, rtol=1e-5, atol=4e-5)
271- except AssertionError:
272- print(f"Gradient test failed for {name}: {p.grad} vs {ref_p.grad}")
273- raise
274 513 
275 @parametrize("ScheduleClass", [ScheduleGPipe, Schedule1F1B])514 @parametrize("ScheduleClass", [ScheduleGPipe, Schedule1F1B])
276 @parametrize("shape_inference", [True, False])515 @parametrize("shape_inference", [True, False])
277 def test_grad_with_manual(self, ScheduleClass, shape_inference):516 def test_grad_with_manual(self, ScheduleClass, shape_inference):
278- full_mod = MultiMLP(d_hid, n_layers=self.world_size)517+ mod, ref_mod, x, target, loss_fn = setup_models_and_data(self.config)
279- full_mod.to(self.device)
280- 
281- ref_mod = copy.deepcopy(full_mod)
282- x = torch.randn(batch_size, d_hid, device=self.device)
283- with torch.no_grad():
284- y = ref_mod(x)
285- # Add a small perturbation
286- target = y + torch.randn(batch_size, d_hid, device=self.device)
287- 
288- loss_fn = torch.nn.MSELoss(reduction="sum")
289 518 
290 # Run reference519 # Run reference
291- for _ in range(2):520+ ref_out, ref_loss = run_reference_model(ref_mod, x, target, loss_fn)
292- ref_mod.zero_grad()
293- ref_out = ref_mod(x)
294- ref_loss = loss_fn(ref_out, target)
295- ref_loss.backward()
296 521 
297- # Get a submodule, e.g. `layers.0` or `layers.1`522+ # Create manual pipeline stage
298- submod_name = f"layers.{self.rank}"523+ chunks = 2 * self.world_size
299- stage_module = full_mod.get_submodule(submod_name)524+ stage, stage_module, _ = create_single_stage_pipeline(
300- chunks = 4525+ self.config, mod, x, chunks, use_tracer=False
526+ )
301 527 
302- if shape_inference:528+ # Handle shape inference
303- input_args = None529+ if not shape_inference:
304- output_args = None
305- else:
306 input_args = (x.chunk(chunks)[0],)530 input_args = (x.chunk(chunks)[0],)
307 with torch.no_grad():531 with torch.no_grad():
308 output_args = stage_module(*input_args)532 output_args = stage_module(*input_args)
533+ stage = PipelineStage(
534+ stage_module,
535+ self.rank,
536+ self.world_size,
537+ self.device,
538+ input_args=input_args,
539+ output_args=output_args,
540+ )
309 541 
310- # Create a pipeline stage to wrap that submodule
311- stage = PipelineStage(
312- stage_module,
313- self.rank,
314- self.world_size,
315- self.device,
316- input_args=input_args,
317- output_args=output_args,
318- )
319- 
320- # Attach to a schedule
321 schedule = ScheduleClass(stage, chunks, loss_fn=loss_fn, scale_grads=False)542 schedule = ScheduleClass(stage, chunks, loss_fn=loss_fn, scale_grads=False)
322 543 
323- # Run544+ # Run pipeline
545+ out = None
546+ losses = []
324 for _ in range(2):547 for _ in range(2):
325- # Zero gradients548+ zero_gradients(stage_module)
326- stage_module.zero_grad()
327 if self.rank == 0:549 if self.rank == 0:
328 schedule.step(x)550 schedule.step(x)
329 elif self.rank == self.world_size - 1:551 elif self.rank == self.world_size - 1:
330- losses = []
331 out = schedule.step(target=target, losses=losses)552 out = schedule.step(target=target, losses=losses)
332 else:553 else:
333 schedule.step()554 schedule.step()
334 555 
335- dist.barrier()556+ dist.barrier(device_ids=[self.rank])
336 557 
337 # Last rank checks result558 # Last rank checks result
338 if self.rank == self.world_size - 1:559 if self.rank == self.world_size - 1:
339- # Check output
340 torch.testing.assert_close(out, ref_out)560 torch.testing.assert_close(out, ref_out)
341- # Check loss
342- # Since the reduction used in the loss function above is "sum", we use
343- # "sum" here to reduce microbatch losses into a single value too.
344 pipe_loss = sum(losses)561 pipe_loss = sum(losses)
345 torch.testing.assert_close(pipe_loss, ref_loss)562 torch.testing.assert_close(pipe_loss, ref_loss)
346 563 
347- # Every rank checks gradients564+ # Check gradients using helper method
348- ref_submod = ref_mod.get_submodule(submod_name)565+ check_gradients(self.config, stage_module, ref_mod)
349- for name, p in stage_module.named_parameters():
350- ref_p = ref_submod.get_parameter(name)
351- try:
352- torch.testing.assert_close(p.grad, ref_p.grad, rtol=1e-5, atol=4e-5)
353- except AssertionError:
354- print(f"Gradient test failed for {name}: {p.grad} vs {ref_p.grad}")
355- raise
356 566 
357 @parametrize(567 @parametrize(
358 "ScheduleClass",568 "ScheduleClass",
@@ -366,117 +576,83 @@ class ScheduleTest(MultiProcContinousTest):
366 def test_grad_with_manual_interleaved(self, ScheduleClass, use_new_runtime):576 def test_grad_with_manual_interleaved(self, ScheduleClass, use_new_runtime):
367 stages_per_rank = 2577 stages_per_rank = 2
368 n_stages = stages_per_rank * self.world_size578 n_stages = stages_per_rank * self.world_size
369- full_mod = MultiMLP(d_hid, n_layers=n_stages)579+ mod, ref_mod, x, target, loss_fn = setup_models_and_data(
370- full_mod.to(self.device)580+ self.config, n_layers=n_stages
371- 581+ )
372- ref_mod = copy.deepcopy(full_mod)
373- x = torch.randn(batch_size, d_hid, device=self.device)
374- with torch.no_grad():
375- y = ref_mod(x)
376- # Add a small perturbation
377- target = y + torch.randn(batch_size, d_hid, device=self.device)
378- 
379- loss_fn = torch.nn.MSELoss(reduction="sum")
380 582 
381 # Run reference583 # Run reference
382- for _ in range(2):584+ ref_out, ref_loss = run_reference_model(ref_mod, x, target, loss_fn)
383- ref_mod.zero_grad()585+ 
384- ref_out = ref_mod(x)586+ # Create multi-stage pipeline
385- ref_loss = loss_fn(ref_out, target)587+ stages, stage_modules, submod_names = create_multi_stage_pipeline(
386- ref_loss.backward()588+ self.config, mod, stages_per_rank, n_stages
589+ )
590+ print(f"Rank {self.rank} stages: {[stage.stage_index for stage in stages]}")
387 591 
388- # Get a submodule, e.g. `layers.0` or `layers.1`
389- stage_indices = [
390- self.rank + i * self.world_size
391- for i in range(stages_per_rank)
392- ]
393- print(f"Rank {self.rank} stages: {stage_indices}")
394- submod_names = [f"layers.{i}" for i in stage_indices]
395- stage_modules = [
396- full_mod.get_submodule(submod_name)
397- for submod_name in submod_names
398- ]
399- # Create a pipeline stage to wrap that submodule
400 num_microbatches = (592 num_microbatches = (
401 ScheduleClass.num_microbatches593 ScheduleClass.num_microbatches
402 if hasattr(ScheduleClass, "num_microbatches")594 if hasattr(ScheduleClass, "num_microbatches")
403- else 8595+ else 2 * self.world_size
404 )596 )
405- stages = [
406- PipelineStage(
407- stage_module,
408- stage_idx,
409- n_stages,
410- self.device,
411- )
412- for stage_module, stage_idx in zip(stage_modules, stage_indices)
413- ]
414 597 
415- # Attach to a schedule598+ # Create schedule
416 schedule = ScheduleClass(599 schedule = ScheduleClass(
417 stages, num_microbatches, loss_fn=loss_fn, scale_grads=False600 stages, num_microbatches, loss_fn=loss_fn, scale_grads=False
418 )601 )
602+ 
603+ # Handle new runtime testing
419 if use_new_runtime:604 if use_new_runtime:
420 old_schedule = schedule605 old_schedule = schedule
421 tmp_schedule = _PipelineScheduleRuntime(606 tmp_schedule = _PipelineScheduleRuntime(
422- stages,607+ stages, num_microbatches, loss_fn=loss_fn, scale_grads=False
423- num_microbatches,
424- loss_fn=loss_fn,
425- scale_grads=False,
426 )608 )
427- tmp_schedule._load_actions(old_schedule.pipeline_order)609+ tmp_schedule._prepare_schedule_with_comms(old_schedule.pipeline_order)
428- # test that csv round-trip works for compute_comms schedule610+ 
611+ # Test CSV round-trip for compute_comms schedule
429 schedule = _PipelineScheduleRuntime(612 schedule = _PipelineScheduleRuntime(
430- stages,613+ stages, num_microbatches, loss_fn=loss_fn, scale_grads=False
431- num_microbatches,
432- loss_fn=loss_fn,
433- scale_grads=False,
434 )614 )
435 with tempfile.NamedTemporaryFile() as f:615 with tempfile.NamedTemporaryFile() as f:
436 tmp_schedule._dump_csv(f.name)616 tmp_schedule._dump_csv(f.name)
437 f.seek(0)617 f.seek(0)
438 schedule._load_csv(f.name, format="compute_comms")618 schedule._load_csv(f.name, format="compute_comms")
619+ 
439 one_more_schedule = _PipelineScheduleRuntime(620 one_more_schedule = _PipelineScheduleRuntime(
440- stages,621+ stages, num_microbatches, loss_fn=loss_fn, scale_grads=False
441- num_microbatches,
442- loss_fn=loss_fn,
443- scale_grads=False,
444 )622 )
445- one_more_schedule._load_actions(623+ one_more_schedule._prepare_schedule_with_comms(
446 schedule.pipeline_order_with_comms, format="compute_comms"624 schedule.pipeline_order_with_comms, format="compute_comms"
447 )625 )
626+ 
627+ # Verify schedule consistency
448 self.assertEqual(628 self.assertEqual(
449 len(schedule.pipeline_order_with_comms),629 len(schedule.pipeline_order_with_comms),
450- len(630+ len(one_more_schedule.pipeline_order_with_comms),
451- one_more_schedule.pipeline_order_with_comms,
452- ),
453 )631 )
454 for rank in schedule.pipeline_order_with_comms:632 for rank in schedule.pipeline_order_with_comms:
455 self.assertEqual(633 self.assertEqual(
456 len(schedule.pipeline_order_with_comms[rank]),634 len(schedule.pipeline_order_with_comms[rank]),
457- len(635+ len(one_more_schedule.pipeline_order_with_comms[rank]),
458- one_more_schedule.pipeline_order_with_comms[rank],
459- ),
460 )636 )
461 for a, b in zip(637 for a, b in zip(
462- schedule.pipeline_order_with_comms[rank],638+ schedule.pipeline_order_with_comms[rank],
463- one_more_schedule.pipeline_order_with_comms[rank],639+ one_more_schedule.pipeline_order_with_comms[rank],
464 ):640 ):
465 self.assertEqual(a, b)641 self.assertEqual(a, b)
466 642 
467- # Run643+ # Run pipeline with tensor leak checking
644+ out = None
645+ losses = []
468 with check_leaked_tensors() as garbage_tensors:646 with check_leaked_tensors() as garbage_tensors:
469 for _ in range(2):647 for _ in range(2):
470- # Zero gradients648+ zero_gradients(stage_modules)
471- for stage_module in stage_modules:
472- stage_module.zero_grad()
473 if self.rank == 0:649 if self.rank == 0:
474 schedule.step(x)650 schedule.step(x)
475 elif self.rank == self.world_size - 1:651 elif self.rank == self.world_size - 1:
476- losses = []
477 out = schedule.step(target=target, losses=losses)652 out = schedule.step(target=target, losses=losses)
478 else:653 else:
479 schedule.step()654 schedule.step()
655+ 
480 self.assertEqual(656 self.assertEqual(
481 len(garbage_tensors),657 len(garbage_tensors),
482 0,658 0,
@@ -484,372 +660,35 @@ class ScheduleTest(MultiProcContinousTest):
484 )660 )
485 dist.barrier()661 dist.barrier()
486 662 
487- # Last rank checks result663+ # Verify results
488 if self.rank == self.world_size - 1:664 if self.rank == self.world_size - 1:
489- # Check output
490 torch.testing.assert_close(out, ref_out)665 torch.testing.assert_close(out, ref_out)
491- # Check loss
492- # Since the reduction used in the loss function above is "sum", we use
493- # "sum" here to reduce microbatch losses into a single value too.
494 pipe_loss = sum(losses)666 pipe_loss = sum(losses)
495 torch.testing.assert_close(pipe_loss, ref_loss)667 torch.testing.assert_close(pipe_loss, ref_loss)
496 668 
497- # Every rank checks gradients669+ # Check gradients - use relaxed tolerances for interleaved schedules
498- for stage_module, submod_name in zip(stage_modules, submod_names):670+ # since gradients are small
499- # Get corresponding submodule from reference model671+ check_gradients(
500- ref_submod = ref_mod.get_submodule(submod_name)672+ self.config, stage_modules, ref_mod, submod_names, rtol=5e-3, atol=5e-3
501- # Check gradients per parameter
502- for name, p in stage_module.named_parameters():
503- ref_p = ref_submod.get_parameter(name)
504- try:
505- torch.testing.assert_close(p.grad, ref_p.grad, rtol=1e-5, atol=4e-5)
506- except AssertionError:
507- print(f"Gradient test failed for {name}: {p.grad} vs {ref_p.grad}")
508- raise
509- 
510- @parametrize("ScheduleClass", [ScheduleWithW, ScheduleInterleavedZeroBubble])
511- def test_schedule_with_native_zero_bubble(self, ScheduleClass):
512- print(ScheduleClass)
513- if ScheduleClass is ScheduleInterleavedZeroBubble:
514- n_stages = 4
515- num_microbatches = 8
516- rank_stages = {
517- 0: [0, 2],
518- 1: [1, 3],
519- }
520- else:
521- n_stages = ScheduleClass.n_stages
522- num_microbatches = ScheduleClass.num_microbatches
523- rank_stages = ScheduleClass.rank_stages
524- 
525- num_steps = 4
526- full_mod = MultiMLP(d_hid, n_layers=n_stages)
527- full_mod.to(self.device)
528- 
529- ref_mod = copy.deepcopy(full_mod)
530- x = torch.randn(batch_size, d_hid, device=self.device)
531- with torch.no_grad():
532- y = ref_mod(x)
533- # Add a small perturbation
534- target = y + torch.randn(batch_size, d_hid, device=self.device)
535- 
536- loss_fn = torch.nn.MSELoss(reduction="sum")
537- 
538- # Create a pipeline stage to wrap that submodule
539- stage_indices = rank_stages.get(self.rank)
540- print(f"Rank {self.rank} stages: {stage_indices}")
541- submod_names = [f"layers.{i}" for i in stage_indices]
542- stage_modules = [
543- full_mod.get_submodule(submod_name)
544- for submod_name in submod_names
545- ]
546- stages = [
547- PipelineStage(
548- stage_module,
549- stage_idx,
550- n_stages,
551- self.device,
552- )
553- for stage_module, stage_idx in zip(stage_modules, rank_stages.get(self.rank))
554- ]
555- 
556- # We set scale_grads=False since we use a loss function that sums instead of mean-reduces
557- # (note: normally we recommend using mean-reduce loss functions, but we preserve at least one test case
558- # using sum scaling for completeness)
559- schedule = ScheduleClass(
560- stages, num_microbatches, loss_fn=loss_fn, scale_grads=False
561 )673 )
562 674 
563- # Run reference
564- ref_x = x.detach().clone().requires_grad_(x.requires_grad)
565- torch.testing.assert_close(x, ref_x)
566- for _ in range(num_steps):
567- ref_out = ref_mod(ref_x)
568- ref_loss = loss_fn(ref_out, target)
569- ref_loss.backward()
570- 
571- with check_leaked_tensors() as garbage_tensors:
572- # Run pipelined stages
573- for _ in range(num_steps):
574- if self.rank == 0:
575- schedule.step(x)
576- elif self.rank == self.world_size - 1:
577- losses = []
578- schedule.step(target=target, losses=losses)
579- else:
580- schedule.step()
581- self.assertEqual(
582- len(garbage_tensors),
583- 0,
584- "Found leaked tensors, check logs above for debug info",
585- )
586- 
587- # Every rank checks parameters compared with the reference model
588- for stage_module, submod_name in zip(stage_modules, submod_names):
589- # Get corresponding submodule from reference model
590- ref_submod = ref_mod.get_submodule(submod_name)
591- # Check gradients per parameter
592- for name, p in stage_module.named_parameters():
593- ref_p = ref_submod.get_parameter(name)
594- try:
595- torch.testing.assert_close(p.grad, ref_p.grad, rtol=1e-5, atol=4e-5)
596- except AssertionError:
597- print(
598- f"Parameter test failed for {submod_name}.{name}: {p.grad} vs {ref_p.grad}"
599- )
600- raise
601- 
602- @parametrize(
603- "ScheduleClass",
604- [
605- ScheduleWithReorderedB,
606- ],
607- )
608- def test_pipeline_schedule_runtime_custom_sched(self, ScheduleClass):
609- n_stages = 2
610- num_microbatches = 2
611- stages_per_rank = 1
612- full_mod = MultiMLP(d_hid, n_layers=n_stages)
613- full_mod.to(self.device)
614- 
615- ref_mod = copy.deepcopy(full_mod)
616- x = torch.randn(batch_size, d_hid, device=self.device)
617- with torch.no_grad():
618- y = ref_mod(x)
619- # Add a small perturbation
620- target = y + torch.randn(batch_size, d_hid, device=self.device)
621- 
622- loss_fn = torch.nn.MSELoss(reduction="sum")
623- 
624- # Run reference
625- for _ in range(2):
626- ref_mod.zero_grad()
627- ref_out = ref_mod(x)
628- ref_loss = loss_fn(ref_out, target)
629- ref_loss.backward()
630- 
631- # Get a submodule, e.g. `layers.0` or `layers.1`
632- stage_indices = [
633- self.rank + i * self.world_size
634- for i in range(stages_per_rank)
635- ]
636- print(f"Rank {self.rank} stages: {stage_indices}")
637- submod_names = [f"layers.{i}" for i in stage_indices]
638- stage_modules = [
639- full_mod.get_submodule(submod_name)
640- for submod_name in submod_names
641- ]
642- # Create a pipeline stage to wrap that submodule
643- num_microbatches = (
644- ScheduleClass.num_microbatches
645- if hasattr(ScheduleClass, "num_microbatches")
646- else 8
647- )
648- stages = [
649- PipelineStage(
650- stage_module,
651- stage_idx,
652- n_stages,
653- self.device,
654- )
655- for stage_module, stage_idx in zip(stage_modules, stage_indices)
656- ]
657- 
658- # Attach to a schedule
659- schedule = ScheduleClass(
660- stages, num_microbatches, loss_fn=loss_fn, scale_grads=False
661- )
662- assert isinstance(schedule, _PipelineScheduleRuntime)
663- 
664- # Run
665- with check_leaked_tensors() as garbage_tensors:
666- for _ in range(2):
667- # Zero gradients
668- for stage_module in stage_modules:
669- stage_module.zero_grad()
670- if self.rank == 0:
671- schedule.step(x)
672- elif self.rank == self.world_size - 1:
673- losses = []
674- out = schedule.step(target=target, losses=losses)
675- else:
676- schedule.step()
677- self.assertEqual(
678- len(garbage_tensors),
679- 0,
680- "Found leaked tensors, check logs above for debug info",
681- )
682- dist.barrier()
683- 
684- # Last rank checks result
685- if self.rank == self.world_size - 1:
686- # Check output
687- torch.testing.assert_close(out, ref_out)
688- # Check loss
689- # Since the reduction used in the loss function above is "sum", we use
690- # "sum" here to reduce microbatch losses into a single value too.
691- pipe_loss = sum(losses)
692- torch.testing.assert_close(pipe_loss, ref_loss)
693- 
694- # Every rank checks gradients
695- for stage_module, submod_name in zip(stage_modules, submod_names):
696- # Get corresponding submodule from reference model
697- ref_submod = ref_mod.get_submodule(submod_name)
698- # Check gradients per parameter
699- for name, p in stage_module.named_parameters():
700- ref_p = ref_submod.get_parameter(name)
701- try:
702- torch.testing.assert_close(p.grad, ref_p.grad, rtol=1e-5, atol=4e-5)
703- except AssertionError:
704- print(f"Gradient test failed for {name}: {p.grad} vs {ref_p.grad}")
705- raise
706- 
707- @parametrize(
708- "schedule_class", [ScheduleVShaped, ScheduleUnbalanced, ScheduleZBVZeroBubble]
709- )
710- @parametrize("use_new_runtime", [False, True])
711- def test_non_symmetric_stage_ids(self, schedule_class, use_new_runtime):
712- if schedule_class is ScheduleZBVZeroBubble:
713- n_stages = 4
714- rank_stages = {
715- 0: [0, 3],
716- 1: [1, 2],
717- }
718- else:
719- n_stages = schedule_class.n_stages
720- rank_stages = schedule_class.rank_stages
721- full_mod = MultiMLP(d_hid, n_layers=n_stages)
722- full_mod.to(self.device)
723- 
724- ref_mod = copy.deepcopy(full_mod)
725- x = torch.randn(batch_size, d_hid, device=self.device)
726- with torch.no_grad():
727- y = ref_mod(x)
728- # Add a small perturbation
729- target = y + torch.randn(batch_size, d_hid, device=self.device)
730- 
731- loss_fn = torch.nn.MSELoss(reduction="sum")
732- 
733- # Run reference
734- for _ in range(2):
735- ref_mod.zero_grad()
736- ref_out = ref_mod(x)
737- ref_loss = loss_fn(ref_out, target)
738- ref_loss.backward()
739- 
740- # Create a pipeline stage to wrap that submodule
741- num_microbatches = 1
742- stage_indices = rank_stages.get(self.rank)
743- print(f"Rank {self.rank} stages: {stage_indices}")
744- submod_names = [f"layers.{i}" for i in stage_indices]
745- stage_modules = [
746- full_mod.get_submodule(submod_name)
747- for submod_name in submod_names
748- ]
749- stages = [
750- PipelineStage(
751- stage_module,
752- stage_idx,
753- n_stages,
754- self.device,
755- )
756- for stage_module, stage_idx in zip(stage_modules, rank_stages.get(self.rank))
757- ]
758- 
759- schedule = schedule_class(
760- stages,
761- num_microbatches,
762- loss_fn=loss_fn,
763- scale_grads=False,
764- )
765- if use_new_runtime:
766- old_schedule = schedule
767- schedule = _PipelineScheduleRuntime(
768- stages,
769- num_microbatches,
770- loss_fn=loss_fn,
771- )
772- schedule._load_actions(old_schedule.pipeline_order)
773- 
774- # Run
775- for _ in range(2):
776- # Zero gradients
777- for stage_module in stage_modules:
778- stage_module.zero_grad()
779- if self.rank == 0:
780- losses = []
781- out = schedule.step(x, target=target, losses=losses)
782- else:
783- schedule.step()
784- 
785- dist.barrier()
786- 
787- # Last rank checks result
788- if self.rank == 0:
789- # Check output
790- torch.testing.assert_close(out, ref_out)
791- # Check loss
792- # Since the reduction used in the loss function above is "sum", we use
793- # "sum" here to reduce microbatch losses into a single value too.
794- pipe_loss = sum(losses)
795- torch.testing.assert_close(pipe_loss, ref_loss)
796- 
797- # Every rank checks gradients
798- for stage_module, submod_name in zip(stage_modules, submod_names):
799- # Get corresponding submodule from reference model
800- ref_submod = ref_mod.get_submodule(submod_name)
801- # Check gradients per parameter
802- for name, p in stage_module.named_parameters():
803- ref_p = ref_submod.get_parameter(name)
804- try:
805- torch.testing.assert_close(p.grad, ref_p.grad, rtol=1e-5, atol=4e-5)
806- except AssertionError:
807- print(f"Gradient test failed for {name}: {p.grad} vs {ref_p.grad}")
808- raise
809- 
810 @parametrize("ScheduleClass", [ScheduleInterleavedZeroBubble])675 @parametrize("ScheduleClass", [ScheduleInterleavedZeroBubble])
811 def test_schedule_with_weight_update_mlp_e2e(self, ScheduleClass):676 def test_schedule_with_weight_update_mlp_e2e(self, ScheduleClass):
812 stages_per_rank = 2677 stages_per_rank = 2
813 n_stages = stages_per_rank * self.world_size678 n_stages = stages_per_rank * self.world_size
814- full_mod = MultiMLPWithDw(d_hid, n_layers=n_stages)679+ full_mod, ref_mod, x, target, _ = setup_models_and_data(
815- full_mod.to(self.device)680+ self.config, n_layers=n_stages, model_class=MultiMLPWithDw
816- 681+ )
817- ref_mod = copy.deepcopy(full_mod)
818- x = torch.randn(batch_size, d_hid, device=self.device)
819- with torch.no_grad():
820- y = ref_mod(x)
821- # Add a small perturbation
822- target = y + torch.randn(batch_size, d_hid, device=self.device)
823- 
824- ref_loss_fn = torch.nn.MSELoss(reduction="sum")
825- full_loss_fn = torch.nn.MSELoss(reduction="sum")
826- 
827 full_mod.toggle()682 full_mod.toggle()
828- 683+ loss_fn = MSELoss()
829- # Get a submodule, e.g. `layers.0` or `layers.1`
830- stage_indices = [
831- self.rank + i * self.world_size
832- for i in range(stages_per_rank)
833- ]
834- submod_names = [f"layers.{i}" for i in stage_indices]
835- stage_modules = [
836- full_mod.get_submodule(submod_name)
837- for submod_name in submod_names
838- ]
839 684 
840 # Run reference685 # Run reference
841- for _ in range(2):686+ ref_out, ref_loss = run_reference_model(ref_mod, x, target, loss_fn)
842- ref_stage_modules = [
843- ref_mod.get_submodule(submod_name)
844- for submod_name in submod_names
845- ]
846- for stage_module in ref_stage_modules:
847- stage_module.zero_grad()
848 687 
849- ref_mod.zero_grad()688+ # Create multi-stage pipeline with custom dw_builder
850- ref_out = ref_mod(x)689+ stages, stage_modules, submod_names = create_multi_stage_pipeline(
851- ref_loss = ref_loss_fn(ref_out, target)690+ self.config, full_mod, stages_per_rank, n_stages
852- ref_loss.backward()691+ )
853 692 
854 class CustomState:693 class CustomState:
855 def __init__(self, stage_module, stage_idx, rank):694 def __init__(self, stage_module, stage_idx, rank):
@@ -860,7 +699,6 @@ class ScheduleTest(MultiProcContinousTest):
860 699 
861 def dw_builder(self):700 def dw_builder(self):
862 def dw_runner():701 def dw_runner():
863- # This inner function would be called by PipelineStage during `backward_weight_one_chunk`
864 self.i += 1702 self.i += 1
865 print(703 print(
866 f"[Rank {self.rank}] dw_count={self.i} stage={self.stage_idx}"704 f"[Rank {self.rank}] dw_count={self.i} stage={self.stage_idx}"
@@ -869,12 +707,15 @@ class ScheduleTest(MultiProcContinousTest):
869 707 
870 return dw_runner708 return dw_runner
871 709 
710+ # Create custom states and rebuild stages with dw_builder
872 cs = {}711 cs = {}
712+ stage_indices = [
713+ self.rank + i * self.world_size
714+ for i in range(stages_per_rank)
715+ ]
873 for stage_module, stage_idx in zip(stage_modules, stage_indices):716 for stage_module, stage_idx in zip(stage_modules, stage_indices):
874 cs[stage_idx] = CustomState(stage_module, stage_idx, self.rank)717 cs[stage_idx] = CustomState(stage_module, stage_idx, self.rank)
875 718 
876- # Create a pipeline stage to wrap that submodule
877- chunks = 2
878 stages = [719 stages = [
879 PipelineStage(720 PipelineStage(
880 stage_module,721 stage_module,
@@ -886,73 +727,338 @@ class ScheduleTest(MultiProcContinousTest):
886 for stage_module, stage_idx in zip(stage_modules, stage_indices)727 for stage_module, stage_idx in zip(stage_modules, stage_indices)
887 ]728 ]
888 729 
889- # Attach to a schedule730+ schedule = ScheduleClass(stages, 2, loss_fn=loss_fn)
890- schedule = ScheduleClass(
891- stages, chunks, loss_fn=full_loss_fn, scale_grads=False
892- )
893 731 
732+ # Run pipeline
733+ out = None
734+ losses = []
894 for _ in range(2):735 for _ in range(2):
895- # Zero gradients736+ zero_gradients(stage_modules)
896- for stage_module in stage_modules:
897- stage_module.zero_grad()
898 if self.rank == 0:737 if self.rank == 0:
899 schedule.step(x)738 schedule.step(x)
900 elif self.rank == self.world_size - 1:739 elif self.rank == self.world_size - 1:
901- losses = []740+ out = schedule.step(target=target, losses=losses)
741+ else:
742+ schedule.step()
743+ 
744+ dist.barrier(device_ids=[self.rank])
745+ 
746+ # Verify results
747+ if self.rank == self.world_size - 1:
748+ torch.testing.assert_close(out, ref_out)
749+ pipe_loss = sum(losses) / len(losses)
750+ torch.testing.assert_close(pipe_loss, ref_loss)
751+ 
752+ # Check gradients using helper method
753+ check_gradients(self.config, stage_modules, ref_mod, submod_names)
754+ 
755+ @parametrize(
756+ "schedule_class",
757+ [ScheduleZBVZeroBubble, ScheduleDualPipeV],
758+ )
759+ @parametrize("use_new_runtime", [False, True])
760+ def test_v_shape_schedules(self, schedule_class, use_new_runtime):
761+ n_stages = 8
762+ rank_stages = {0: [0, 7], 1: [1, 6], 2: [2, 5], 3: [3, 4]}
763+ mod, ref_mod, x, target, loss_fn = setup_models_and_data(
764+ self.config, n_layers=n_stages
765+ )
766+ 
767+ # Run reference
768+ ref_out, ref_loss = run_reference_model(ref_mod, x, target, loss_fn)
769+ 
770+ # Create multi-stage pipeline with custom stage indices
771+ num_microbatches = 8
772+ stage_indices = rank_stages[self.rank]
773+ stages, stage_modules, submod_names = create_multi_stage_pipeline(
774+ self.config, mod, len(stage_indices), n_stages, stage_indices
775+ )
776+ 
777+ with patch_stage_init_method(stages):
778+ schedule = schedule_class(
779+ stages, num_microbatches, loss_fn=loss_fn, scale_grads=False
780+ )
781+ 
782+ if schedule_class != ScheduleDualPipeV and use_new_runtime:
783+ old_schedule = schedule
784+ schedule = _PipelineScheduleRuntime(
785+ stages, num_microbatches, loss_fn=loss_fn
786+ )
787+ schedule._prepare_schedule_with_comms(old_schedule.pipeline_order)
788+ 
789+ # Run pipeline - special case where first and last stage are on rank 0
790+ out = None
791+ losses = []
792+ for _ in range(2):
793+ zero_gradients(stage_modules)
794+ if self.rank == 0:
795+ out = schedule.step(x, target=target, losses=losses)
796+ else:
797+ schedule.step()
798+ 
799+ # Verify results (rank 0 has both first and last stages)
800+ if self.rank == 0:
801+ torch.testing.assert_close(out, ref_out)
802+ pipe_loss = sum(losses)
803+ torch.testing.assert_close(pipe_loss, ref_loss)
804+ 
805+ # Check gradients using helper method
806+ check_gradients(self.config, stage_modules, ref_mod, submod_names)
807+ 
808+ @parametrize(
809+ "ScheduleClass",
810+ [ScheduleInterleavedZeroBubble, ScheduleInterleaved1F1B],
811+ )
812+ def test_zero_bubble_with_model_kwargs(self, ScheduleClass):
813+ stages_per_rank = 2
814+ n_stages = stages_per_rank * self.world_size
815+ mod, ref_mod, x, target, loss_fn = setup_models_and_data(
816+ self.config, n_layers=n_stages, model_class=MultiMLPKwargs
817+ )
818+ unused_kwarg = torch.tensor([1.0], device=self.device)
819+ 
820+ # Run reference with kwargs
821+ ref_out, ref_loss = run_reference_model(
822+ ref_mod, x, target, loss_fn, unused_kwarg=unused_kwarg
823+ )
824+ 
825+ # Create multi-stage pipeline
826+ stages, stage_modules, submod_names = create_multi_stage_pipeline(
827+ self.config, mod, stages_per_rank, n_stages
828+ )
829+ 
830+ num_microbatches = (
831+ ScheduleClass.num_microbatches
832+ if hasattr(ScheduleClass, "num_microbatches")
833+ else 2 * self.world_size
834+ )
835+ schedule = ScheduleClass(
836+ stages, num_microbatches, loss_fn=loss_fn, scale_grads=False
837+ )
838+ 
839+ # Run pipeline with kwargs
840+ out = None
841+ losses = []
842+ for _ in range(2):
843+ zero_gradients(stage_modules)
844+ if self.rank == 0:
845+ schedule.step(
846+ x,
847+ unused_kwarg=unused_kwarg.clone()
848+ .unsqueeze(0)
849+ .expand(num_microbatches, -1),
850+ )
851+ elif self.rank == self.world_size - 1:
902 out = schedule.step(target=target, losses=losses)852 out = schedule.step(target=target, losses=losses)
903 else:853 else:
904 schedule.step()854 schedule.step()
905 855 
906 dist.barrier()856 dist.barrier()
907- # Last rank checks result
908- if self.rank == self.world_size - 1:
909- # Check output
910- torch.testing.assert_close(out, ref_out)
911 857 
912- # Check loss858+ # Verify results
913- # Since the reduction used in the loss function above is "sum", we use859+ if self.rank == self.world_size - 1:
914- # "sum" here to reduce microbatch losses into a single value too.860+ torch.testing.assert_close(out, ref_out)
915 pipe_loss = sum(losses)861 pipe_loss = sum(losses)
916 torch.testing.assert_close(pipe_loss, ref_loss)862 torch.testing.assert_close(pipe_loss, ref_loss)
917 863 
918- # Every rank checks gradients864+ # Check gradients using helper method
919- for stage_module, submod_name in zip(stage_modules, submod_names):865+ check_gradients(
920- # Get corresponding submodule from reference model866+ self.config, stage_modules, ref_mod, submod_names, rtol=3e-5, atol=5e-3
921- ref_submod = ref_mod.get_submodule(submod_name)867+ )
922- # Check gradients per parameter
923- for name, p in stage_module.named_parameters():
924- ref_p = ref_submod.get_parameter(name)
925- torch.testing.assert_close(p.grad, ref_p.grad, rtol=1e-5, atol=4e-5)
926 868 
927 869 
928instantiate_parametrized_tests(ScheduleTest)870instantiate_parametrized_tests(ScheduleTest)
929 871 
930 872 
873+class CustomSchedulesTest(MultiProcContinuousTest):
874+ """
875+ These schedules are from the ScheduleRegistry and require world_size == 2
876+ The schedules test weird and unconventional schedules for edge cases
877+ """
878+ 
879+ world_size = 2
880+ 
881+ @classmethod
882+ def backend_str(cls) -> str:
883+ # Testing with HCCL backend
884+ return backend
885+ 
886+ @property
887+ def device(self) -> torch.device:
888+ return torch.device(device_type, self.rank)
889+ 
890+ @property
891+ def config(self) -> PipelineTestConfig:
892+ """Lazily create and return the pipeline test configuration."""
893+ return PipelineTestConfig(
894+ world_size=self.world_size, device=self.device, rank=self.rank
895+ )
896+ 
897+ @parametrize(
898+ "schedule_class",
899+ [ScheduleVShaped, ScheduleUnbalanced]
900+ )
901+ @parametrize("use_new_runtime", [False, True])
902+ def test_non_symmetric_stage_ids(self, schedule_class, use_new_runtime):
903+ n_stages = schedule_class.n_stages
904+ rank_stages = schedule_class.rank_stages
905+ 
906+ mod, ref_mod, x, target, loss_fn = setup_models_and_data(
907+ self.config, n_layers=n_stages
908+ )
909+ 
910+ # Run reference
911+ ref_out, ref_loss = run_reference_model(ref_mod, x, target, loss_fn)
912+ 
913+ # Create multi-stage pipeline with custom stage indices
914+ num_microbatches = 1
915+ stage_indices = rank_stages.get(self.rank)
916+ print(f"Rank {self.rank} stages: {stage_indices}")
917+ stages, stage_modules, submod_names = create_multi_stage_pipeline(
918+ self.config, mod, len(stage_indices), n_stages, stage_indices
919+ )
920+ 
921+ with patch_stage_init_method(stages):
922+ schedule = schedule_class(
923+ stages, num_microbatches, loss_fn=loss_fn, scale_grads=False
924+ )
925+ 
926+ if use_new_runtime:
927+ old_schedule = schedule
928+ schedule = _PipelineScheduleRuntime(
929+ stages, num_microbatches, loss_fn=loss_fn
930+ )
931+ schedule._prepare_schedule_with_comms(old_schedule.pipeline_order)
932+ 
933+ # Run pipeline - special case where first and last stage are on rank 0
934+ out = None
935+ losses = []
936+ for _ in range(2):
937+ zero_gradients(stage_modules)
938+ if self.rank == 0:
939+ out = schedule.step(x, target=target, losses=losses)
940+ else:
941+ schedule.step()
942+ 
943+ dist.barrier()
944+ 
945+ # Verify results (rank 0 has both first and last stages)
946+ if self.rank == 0:
947+ torch.testing.assert_close(out, ref_out)
948+ pipe_loss = sum(losses)
949+ torch.testing.assert_close(pipe_loss, ref_loss)
950+ 
951+ # Check gradients using helper method
952+ check_gradients(self.config, stage_modules, ref_mod, submod_names)
953+ 
954+ @parametrize("ScheduleClass", [ScheduleWithReorderedB])
955+ def test_pipeline_schedule_runtime_custom_sched(self, ScheduleClass):
956+ n_stages = 2
957+ stages_per_rank = 1
958+ mod, ref_mod, x, target, loss_fn = setup_models_and_data(
959+ self.config, n_layers=n_stages
960+ )
961+ 
962+ # Run reference
963+ ref_out, ref_loss = run_reference_model(ref_mod, x, target, loss_fn)
964+ 
965+ # Create pipeline stages
966+ stages, stage_modules, submod_names = create_multi_stage_pipeline(
967+ self.config, mod, stages_per_rank, n_stages
968+ )
969+ print(f"Rank {self.rank} stages: {[stage.stage_index for stage in stages]}")
970+ 
971+ num_microbatches = (
972+ ScheduleClass.num_microbatches
973+ if hasattr(ScheduleClass, "num_microbatches")
974+ else 8
975+ )
976+ 
977+ schedule = ScheduleClass(
978+ stages, num_microbatches, loss_fn=loss_fn, scale_grads=False
979+ )
980+ assert isinstance(schedule, _PipelineScheduleRuntime)
981+ 
982+ # Run pipeline with tensor leak checking
983+ with check_leaked_tensors() as garbage_tensors:
984+ for _ in range(2):
985+ zero_gradients(stage_modules)
986+ if self.rank == 0:
987+ schedule.step(x)
988+ elif self.rank == self.world_size - 1:
989+ losses = []
990+ out = schedule.step(target=target, losses=losses)
991+ else:
992+ schedule.step()
993+ 
994+ self.assertEqual(
995+ len(garbage_tensors),
996+ 0,
997+ "Found leaked tensors, check logs above for debug info",
998+ )
999+ dist.barrier()
1000+ 
1001+ # Verify results
1002+ if self.rank == self.world_size - 1:
1003+ torch.testing.assert_close(out, ref_out)
1004+ pipe_loss = sum(losses)
1005+ torch.testing.assert_close(pipe_loss, ref_loss)
1006+ 
1007+ # Check gradients using helper method
1008+ check_gradients(self.config, stage_modules, ref_mod, submod_names)
1009+ 
1010+ @parametrize("ScheduleClass", [ScheduleWithW])
1011+ def test_schedule_with_native_zero_bubble(self, ScheduleClass):
1012+ n_stages = ScheduleClass.n_stages
1013+ num_microbatches = ScheduleClass.num_microbatches
1014+ rank_stages = ScheduleClass.rank_stages
1015+ 
1016+ num_steps = 4
1017+ mod, ref_mod, x, target, loss_fn = setup_models_and_data(
1018+ self.config, n_layers=n_stages
1019+ )
1020+ 
1021+ # Create multi-stage pipeline with custom stage indices
1022+ stage_indices = rank_stages.get(self.rank)
1023+ print(f"Rank {self.rank} stages: {stage_indices}")
1024+ stages, stage_modules, submod_names = create_multi_stage_pipeline(
1025+ self.config, mod, len(stage_indices), n_stages, stage_indices
1026+ )
1027+ 
1028+ schedule = ScheduleClass(
1029+ stages, num_microbatches, loss_fn=loss_fn, scale_grads=False
1030+ )
1031+ 
1032+ # Run reference model
1033+ ref_x = x.detach().clone().requires_grad_(x.requires_grad)
1034+ torch.testing.assert_close(x, ref_x)
1035+ for _ in range(num_steps):
1036+ ref_out = ref_mod(ref_x)
1037+ ref_loss = loss_fn(ref_out, target)
1038+ ref_loss.backward()
1039+ 
1040+ # Run pipeline with tensor leak checking
1041+ losses = []
1042+ with check_leaked_tensors() as garbage_tensors:
1043+ for _ in range(num_steps):
1044+ if self.rank == 0:
1045+ schedule.step(x)
1046+ elif self.rank == self.world_size - 1:
1047+ schedule.step(target=target, losses=losses)
1048+ else:
1049+ schedule.step()
1050+ 
1051+ self.assertEqual(
1052+ len(garbage_tensors),
1053+ 0,
1054+ "Found leaked tensors, check logs above for debug info",
1055+ )
1056+ 
1057+ # Check gradients using helper method
1058+ check_gradients(self.config, stage_modules, ref_mod, submod_names)
1059+ 
1060+ 
1061+instantiate_parametrized_tests(CustomSchedulesTest)
1062+ 
931if __name__ == "__main__":1063if __name__ == "__main__":
932- # Check if NPU and HCCL are available1064+ run_tests()
933- if not (
934- dist.is_available()
935- and dist.is_hccl_available()
936- and torch.npu.device_count() > 1
937- ):
938- print(
939- "c10d HCCL not available or not enough NPUs, skipping tests",
940- file=sys.stderr,
941- )
942- sys.exit(0)
943- 
944- rank = int(os.getenv("RANK", -1))
945- world_size = int(os.getenv("WORLD_SIZE", 2))
946- 
947- if rank != -1:
948- # Launched with torchrun or other multi-proc launchers. Directly run the test.
949- ScheduleTest.run_rank(rank, world_size)
950- else:
951- # Launched as a single process. Spawn subprocess to run the tests.
952- # Also need a rendezvous file for `init_process_group` purpose.
953- rdvz_file = tempfile.NamedTemporaryFile(delete=False).name
954- torch.multiprocessing.spawn(
955- ScheduleTest.run_rank,
956- nprocs=world_size,
957- args=(world_size, rdvz_file),
958- )
Mtest/distributed/pipelining/test_stage.py+177-87
@@ -1,8 +1,17 @@
1# Copyright (c) Meta Platforms, Inc. and affiliates1# Copyright (c) Meta Platforms, Inc. and affiliates
2# Owner(s): ["oncall: distributed"]2# Owner(s): ["oncall: distributed"]
3+# Licensed under the BSD 3-Clause License (the "License");
4+# you may not use this file except in compliance with the License.
5+# You may obtain a copy of the License at
6+#
7+# https://github.com/pytorch/pytorch/blob/main/LICENSE
8+#
9+# Unless required by applicable law or agreed to in writing, software
10+# distributed under the License is distributed on an "AS IS" BASIS,
11+# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
12+# See the License for the specific language governing permissions and
13+# limitations under the License.
3import os14import os
4-import sys
5-import tempfile
6 15 
7from model_registry import ExampleCode, ModelWithKwargs, MultiMLP16from model_registry import ExampleCode, ModelWithKwargs, MultiMLP
8 17 
@@ -15,23 +24,27 @@ from torch.distributed.pipelining import (
15 ScheduleGPipe,24 ScheduleGPipe,
16)25)
17from torch.distributed.pipelining._utils import PipeliningShapeError26from torch.distributed.pipelining._utils import PipeliningShapeError
18-from torch.testing._internal.common_cuda import TEST_MULTIGPU
19from torch.testing._internal.common_distributed import (27from torch.testing._internal.common_distributed import (
20- MultiProcContinousTest,28+ MultiProcContinuousTest,
21- requires_nccl,29+ MultiProcessTestCase,
30+ requires_accelerator_dist_backend,
22)31)
23from torch.testing._internal.common_utils import (32from torch.testing._internal.common_utils import (
24 instantiate_parametrized_tests,33 instantiate_parametrized_tests,
25 parametrize,34 parametrize,
35+ run_tests,
26 skip_but_pass_in_sandcastle_if,36 skip_but_pass_in_sandcastle_if,
27)37)
28from torch.utils._pytree import tree_map_only38from torch.utils._pytree import tree_map_only
29 39 
30- 
31d_hid = 51240d_hid = 512
32batch_size = 25641batch_size = 256
33chunks = 442chunks = 4
34 43 
44+device_type = acc.type if (acc := torch.accelerator.current_accelerator()) else "cpu"
45+backend = dist.get_default_backend_for_device(device_type)
46+TEST_MULTIACCELERATOR = torch.accelerator.device_count() >= 2
47+ 
35torch.manual_seed(0)48torch.manual_seed(0)
36 49 
37 50 
@@ -59,25 +72,29 @@ def get_flatten_hook():
59 return flatten_hook72 return flatten_hook
60 73 
61 74 
62-class StageTest(MultiProcContinousTest):75+class StageTest(MultiProcContinuousTest):
76+ world_size = int(os.getenv("WORLD_SIZE", 2))
77+ 
63 @classmethod78 @classmethod
64 def backend_str(cls) -> str:79 def backend_str(cls) -> str:
65 # Testing with HCCL backend80 # Testing with HCCL backend
66- return "hccl"81+ return backend
67 82 
68 @classmethod83 @classmethod
69- def setUpClass(cls):84+ def device_type(cls) -> str:
70- """85+ return device_type
71- Class-scope test fixture. Run once for entire test class, before any test starts.86+ 
72- Set up the device.87+ @property
73- """88+ def device(self) -> torch.device:
74- super().setUpClass()89+ return torch.device(device_type, self.rank)
75- dev_id = cls.rank % torch.npu.device_count()
76- cls.device = torch.device(f"npu:{dev_id}")
77 90 
78 @parametrize("ModelClass", [ExampleCode, MultiMLP])91 @parametrize("ModelClass", [ExampleCode, MultiMLP])
79 def test_tracer(self, ModelClass):92 def test_tracer(self, ModelClass):
80- mod = ModelClass(d_hid)93+ # mod = ModelClass(d_hid, self.world_size)
94+ if ModelClass == ExampleCode:
95+ mod = ModelClass(d_hid)
96+ else:
97+ mod = ModelClass(d_hid, n_layers=self.world_size)
81 mod.to(self.device)98 mod.to(self.device)
82 99 
83 x = torch.randn(batch_size, d_hid, device=self.device)100 x = torch.randn(batch_size, d_hid, device=self.device)
@@ -117,25 +134,9 @@ class StageTest(MultiProcContinousTest):
117 old_keys = mod.state_dict().keys()134 old_keys = mod.state_dict().keys()
118 assert all(k in old_keys for k in submod_keys)135 assert all(k in old_keys for k in submod_keys)
119 136 
120- if self.rank == 0:
121- # intended to run this code on all ranks, but the problem is if rank0 throws,
122- # it won't perform the send that unblocks rank 1.
123- 
124- with self.assertRaisesRegex(PipeliningShapeError, "dtype mismatch"):
125- _run_step(x.to(torch.int32))
126- 
127- # output of stage's mlp layer will be flattened by this hook, the stage should err
128- handle = stage.submod.register_forward_hook(get_flatten_hook())
129- with self.assertRaisesRegex(PipeliningShapeError, "shape mismatch"):
130- _run_step(x)
131- handle.remove()
132- 
133- stage.submod.register_forward_hook(get_dtype_change_hook(torch.bfloat16))
134- with self.assertRaisesRegex(PipeliningShapeError, "dtype mismatch"):
135- _run_step(x)
136- 
137 @parametrize("ModelClass", [ModelWithKwargs])137 @parametrize("ModelClass", [ModelWithKwargs])
138 def test_tracer_kwargs(self, ModelClass):138 def test_tracer_kwargs(self, ModelClass):
139+ # mod = ModelClass(d_hid, self.world_size)
139 mod = ModelClass(d_hid)140 mod = ModelClass(d_hid)
140 mod.to(self.device)141 mod.to(self.device)
141 142 
@@ -211,23 +212,6 @@ class StageTest(MultiProcContinousTest):
211 ref_out = full_mod(x)212 ref_out = full_mod(x)
212 torch.testing.assert_close(out, ref_out)213 torch.testing.assert_close(out, ref_out)
213 214 
214- if self.rank == 0:
215- with self.assertRaisesRegex(PipeliningShapeError, "shape mismatch"):
216- _run_step(torch.randn(batch_size + 1, d_hid, device=self.device))
217- 
218- with self.assertRaisesRegex(PipeliningShapeError, "dtype mismatch"):
219- _run_step(x.to(torch.int32))
220- 
221- # output of stage's mlp layer will be flattened by this hook, the stage should err
222- handle = stage_mod.register_forward_hook(get_flatten_hook())
223- with self.assertRaisesRegex(PipeliningShapeError, "shape mismatch"):
224- _run_step(x)
225- handle.remove()
226- 
227- stage_mod.register_forward_hook(get_dtype_change_hook(torch.bfloat16))
228- with self.assertRaisesRegex(PipeliningShapeError, "dtype mismatch"):
229- _run_step(x)
230- 
231 def test_custom_dw_with_fb_schedule(self):215 def test_custom_dw_with_fb_schedule(self):
232 """Tests that separate weight grad function 'dw_runner' gets run under a schedule that's only aware of F/B."""216 """Tests that separate weight grad function 'dw_runner' gets run under a schedule that's only aware of F/B."""
233 full_mod = MultiMLP(d_hid, n_layers=self.world_size)217 full_mod = MultiMLP(d_hid, n_layers=self.world_size)
@@ -286,18 +270,151 @@ class StageTest(MultiProcContinousTest):
286 ref_out = full_mod(x)270 ref_out = full_mod(x)
287 torch.testing.assert_close(out, ref_out)271 torch.testing.assert_close(out, ref_out)
288 272 
289- if self.rank == 0:273+ def test_output_chunks_memory_usage(self):
290- with self.assertRaisesRegex(PipeliningShapeError, "shape mismatch"):274+ """Test that output_chunks doesn't store memory for non-first stages."""
291- _run_step(torch.randn(batch_size + 1, d_hid, device=self.device))
292- 
293- def test_custom_dw_errors(self):
294- """Tests expected errors are raised"""
295 full_mod = MultiMLP(d_hid, n_layers=self.world_size)275 full_mod = MultiMLP(d_hid, n_layers=self.world_size)
296 full_mod.to(self.device)276 full_mod.to(self.device)
297 stage_mod = full_mod.get_submodule(f"layers.{self.rank}")277 stage_mod = full_mod.get_submodule(f"layers.{self.rank}")
298- 
299 x = torch.randn(batch_size, d_hid, device=self.device)278 x = torch.randn(batch_size, d_hid, device=self.device)
300 target = torch.randn(batch_size, d_hid, device=self.device)279 target = torch.randn(batch_size, d_hid, device=self.device)
280+ stage = PipelineStage(
281+ stage_mod,
282+ self.rank,
283+ self.world_size,
284+ self.device,
285+ )
286+ self.assertEqual(
287+ len(stage.output_chunks), 0, "output_chunks should be empty initially"
288+ )
289+ 
290+ schedule = ScheduleGPipe(
291+ stage, chunks, loss_fn=torch.nn.MSELoss(reduction="sum")
292+ )
293+ 
294+ def _run_step(x):
295+ if self.rank == 0:
296+ return schedule.step(x)
297+ elif self.rank == self.world_size - 1:
298+ return schedule.step(target=target)
299+ else:
300+ return schedule.step()
301+ 
302+ _run_step(x)
303+ 
304+ # Verify fwd_cache is empty
305+ self.assertEqual(len(stage.fwd_cache), 0, "fwd_cache should be cleared")
306+ 
307+ # Check output_chunks state after step
308+ if self.rank == self.world_size - 1:
309+ self.assertEqual(
310+ len(stage.output_chunks),
311+ chunks,
312+ "Last stage should store output chunks",
313+ )
314+ else:
315+ self.assertEqual(
316+ len(stage.output_chunks),
317+ 0,
318+ f"Non-last stage (rank {self.rank}) should not store output chunks",
319+ )
320+ 
321+ # Clear the schedule and stage caches
322+ stage.clear_runtime_states()
323+ if self.rank == self.world_size - 1:
324+ # Last stage should have output_chunks populated
325+ self.assertEqual(
326+ len(stage.output_chunks), 0, "Last stage should store output chunks"
327+ )
328+ 
329+ 
330+instantiate_parametrized_tests(StageTest)
331+ 
332+ 
333+class StageNegativeTest(MultiProcessTestCase):
334+ @property
335+ def world_size(self) -> int:
336+ # return torch.get_device_module(device_type).device_count()
337+ return int(os.getenv("WORLD_SIZE", 2))
338+ 
339+ @property
340+ def device(self) -> torch.device:
341+ device = torch.device(device_type, self.rank)
342+ return torch.device(device_type, self.rank)
343+ 
344+ def setUp(self):
345+ super().setUp()
346+ self._spawn_processes()
347+ 
348+ def tearDown(self):
349+ super().tearDown()
350+ try:
351+ os.remove(self.file_name)
352+ except OSError:
353+ pass
354+ 
355+ def init_pg(self):
356+ store = dist.FileStore(self.file_name, self.world_size)
357+ dist.init_process_group(
358+ backend=backend,
359+ store=store,
360+ rank=self.rank,
361+ world_size=self.world_size,
362+ device_id=self.device,
363+ )
364+ 
365+ # def test_shape_prop_mismatch(self):
366+ # """Tests shape prop errors are raised"""
367+ # self.init_pg()
368+ 
369+ # full_mod = MultiMLP(d_hid, n_layers=self.world_size)
370+ # full_mod.to(self.device)
371+ # stage_mod = full_mod.get_submodule(f"layers.{self.rank}")
372+ 
373+ # x = torch.randn(batch_size, d_hid, device=self.device)
374+ 
375+ # stage = PipelineStage(
376+ # stage_mod,
377+ # self.rank,
378+ # self.world_size,
379+ # self.device,
380+ # )
381+ 
382+ # # Attach to a schedule
383+ # schedule = ScheduleGPipe(stage, chunks)
384+ 
385+ # # Run
386+ # def _run_step(x):
387+ # if self.rank == 0:
388+ # return schedule.step(x)
389+ # else:
390+ # return schedule.step()
391+ 
392+ # _run_step(x)
393+ 
394+ # if self.rank == 0:
395+ # with self.assertRaisesRegex(PipeliningShapeError, "shape mismatch"):
396+ # _run_step(torch.randn(batch_size + 1, d_hid, device=self.device))
397+ 
398+ # with self.assertRaisesRegex(PipeliningShapeError, "dtype mismatch"):
399+ # _run_step(x.to(torch.int32))
400+ 
401+ # # output of stage's mlp layer will be flattened by this hook, the stage should err
402+ # handle = stage_mod.register_forward_hook(get_flatten_hook())
403+ # with self.assertRaisesRegex(PipeliningShapeError, "shape mismatch"):
404+ # _run_step(x)
405+ # handle.remove()
406+ 
407+ # stage_mod.register_forward_hook(get_dtype_change_hook(torch.bfloat16))
408+ # with self.assertRaisesRegex(PipeliningShapeError, "dtype mismatch"):
409+ # _run_step(x)
410+ 
411+ def test_custom_dw_errors(self):
412+ """Tests expected errors are raised"""
413+ self.init_pg()
414+ 
415+ full_mod = MultiMLP(d_hid, n_layers=self.world_size)
416+ full_mod.to(self.device)
417+ stage_mod = full_mod.get_submodule(f"layers.{self.rank}")
301 418 
302 stage_with_dw_builder = PipelineStage(419 stage_with_dw_builder = PipelineStage(
303 stage_mod,420 stage_mod,
@@ -306,37 +423,10 @@ class StageTest(MultiProcContinousTest):
306 self.device,423 self.device,
307 dw_builder=lambda: None,424 dw_builder=lambda: None,
308 )425 )
426+ stage_with_dw_builder._has_backward = True
309 with self.assertRaisesRegex(AssertionError, "backward_one_chunk"):427 with self.assertRaisesRegex(AssertionError, "backward_one_chunk"):
310 stage_with_dw_builder.backward_weight_one_chunk(bwd_chunk_id=0)428 stage_with_dw_builder.backward_weight_one_chunk(bwd_chunk_id=0)
311 429 
312 430 
313-instantiate_parametrized_tests(StageTest)
314- 
315if __name__ == "__main__":431if __name__ == "__main__":
316- # Check if NPU and HCCL are available432+ run_tests()
317- if not (
318- dist.is_available()
319- and dist.is_hccl_available()
320- and torch.npu.device_count() > 1
321- ):
322- print(
323- "c10d HCCL not available or not enough GPUs, skipping tests",
324- file=sys.stderr,
325- )
326- sys.exit(0)
327- 
328- rank = int(os.getenv("RANK", -1))
329- world_size = int(os.getenv("WORLD_SIZE", 2))
330- 
331- if rank != -1:
332- # Launched with torchrun or other multi-proc launchers. Directly run the test.
333- StageTest.run_rank(rank, world_size)
334- else:
335- # Launched as a single process. Spawn subprocess to run the tests.
336- # Also need a rendezvous file for `init_process_group` purpose.
337- rdvz_file = tempfile.NamedTemporaryFile(delete=False).name
338- torch.multiprocessing.spawn(
339- StageTest.run_rank,
340- nprocs=world_size,
341- args=(world_size, rdvz_file),
342- )