add pp ut-pytorch-AtomGit

已合并

add pp ut #21486

zhanhao创建于 2025年6月3日从refs/pull/21486/head 合入到 master

add pp ut #21486

已合并

zhanhao创建于 2025年6月3日

从refs/pull/21486/head合入到master

共 13 个文件变更+3308-0

Atest/distributed/pipelining/__init__.py+0-0

The file is empty

Atest/distributed/pipelining/artifacts/zb1p_2rank_2stagep_comms.csv+2-0

@@ -0,0 +1,2 @@
	1	+0F0,0SEND_F0,2RECV_F0,0F1,0SEND_F1,2RECV_F1,2F0,2SEND_F0,2F1,2SEND_F1,2RECV_B0,2B0,2SEND_B0,0F2,0SEND_F2,2RECV_B1,2B1,2SEND_B1,0F3,0SEND_F3,2RECV_F2,0RECV_B0,0B0,2F2,2SEND_F2,2RECV_F3,0RECV_B1,0B1,2F3,2SEND_F3,2RECV_B2,2B2,2SEND_B2,0F4,0SEND_F4,2RECV_B3,2B3,2SEND_B3,0F5,0SEND_F5,2RECV_F4,0RECV_B2,0B2,2F4,2SEND_F4,2RECV_F5,0RECV_B3,0B3,2F5,2SEND_F5,2RECV_B4,2B4,2SEND_B4,0F6,0SEND_F6,2RECV_B5,2B5,2SEND_B5,0F7,0SEND_F7,2RECV_F6,0RECV_B4,0B4,2F6,2SEND_F6,2RECV_F7,0RECV_B5,0B5,2F7,2SEND_F7,2RECV_B6,2B6,2SEND_B6,2RECV_B7,2B7,2SEND_B7,0RECV_B6,0B6,0RECV_B7,0B7
	2	+1RECV_F0,1F0,1SEND_F0,1RECV_F1,1F1,1SEND_F1,3RECV_F0,3F0,3RECV_F1,3I0,3SEND_B0,1RECV_B0,3F1,1RECV_F2,3I1,3SEND_B1,1RECV_B1,3W0,1RECV_F3,1F2,1SEND_F2,1I0,1SEND_B0,3W1,3RECV_F2,1F3,1SEND_F3,1I1,1SEND_B1,1W0,3RECV_F3,3F2,3I2,3SEND_B2,1RECV_B2,1W1,1RECV_F4,3F3,3I3,3SEND_B3,1RECV_B3,3W2,1RECV_F5,1F4,1SEND_F4,1I2,1SEND_B2,3W3,3RECV_F4,1F5,1SEND_F5,1I3,1SEND_B3,1W2,3RECV_F5,3F4,3I4,3SEND_B4,1RECV_B4,1W3,1RECV_F6,3F5,3I5,3SEND_B5,1RECV_B5,3W4,1RECV_F7,1F6,1SEND_F6,1I4,1SEND_B4,3W5,3RECV_F6,1F7,1SEND_F7,1I5,1SEND_B5,1W4,3RECV_F7,3F6,3I6,3SEND_B6,1RECV_B6,1W5,3F7,3I7,3SEND_B7,1RECV_B7,3W6,1I6,1SEND_B6,3W7,1I7,1SEND_B7,1W6,1W7

Atest/distributed/pipelining/artifacts/zb1p_2rank_2stagep_compute.csv+2-0

@@ -0,0 +1,2 @@
	1	+0F0,0F1,2F0,,2F1,2I0,2W0,0F2,2I1,2W1,0F3,0I0,0W0,2F2,0I1,0W1,2F3,2I2,2W2,0F4,2I3,2W3,0F5,0I2,0W2,2F4,0I3,0W3,2F5,2I4,2W4,0F6,2I5,2W5,0F7,0I4,0W4,2F6,0I5,0W5,2F7,2I6,2W6,2I7,2W7,0I6,0W6,0I7,0W7
	2	+,1F0,1F1,3F0,3I0,3F1,3I1,3W0,1F2,1I0,3W1,1F3,1I1,1W0,3F2,3I2,1W1,3F3,3I3,3W2,1F4,1I2,3W3,1F5,1I3,1W2,3F4,3I4,1W3,3F5,3I5,3W4,1F6,1I4,3W5,1F7,1I5,1W4,3F6,3I6,1W5,3F7,3I7,3W6,1I6,3W7,1I7,1W6,1W7

Atest/distributed/pipelining/model_registry.py+230-0

@@ -0,0 +1,230 @@
	1	+# Copyright (c) Meta Platforms, Inc. and affiliates
	2	+# Owner(s): ["oncall: distributed"]
	3	+# This file is a model zoo for testing torch.distributed.pipelining.
	4	+import torch
	5	+from torch.autograd import Function
	6	+from torch.distributed.pipelining import pipe_split, SplitPoint
	7	+
	8	+
	9	+class ExampleCode(torch.nn.Module):
	10	+ def __init__(self, d_hid):
	11	+ super().__init__()
	12	+ 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))
	14	+ self.cval = torch.nn.Buffer(torch.randn((d_hid,), requires_grad=False))
	15	+ self.lin0 = torch.nn.Linear(d_hid, d_hid)
	16	+ self.lin1 = torch.nn.Linear(d_hid, d_hid)
	17	+
	18	+ def forward(self, x):
	19	+ x = torch.mm(x, self.mm_param0)
	20	+ x = torch.relu(x)
	21	+ # try passing a value that doesn't require_grad across skip boundaries
	22	+ a_constant = self.cval.clone()
	23	+ x = self.lin0(x)
	24	+ pipe_split()
	25	+ x = torch.relu(x) + a_constant
	26	+ x = torch.mm(x, self.mm_param1)
	27	+ x = self.lin1(x)
	28	+ x = torch.relu(x)
	29	+ return x
	30	+
	31	+
	32	+class ModelWithKwargs(torch.nn.Module):
	33	+ DEFAULT_DHID = 512
	34	+ DEFAULT_BATCH_SIZE = 256
	35	+
	36	+ def __init__(self, d_hid: int = DEFAULT_DHID):
	37	+ super().__init__()
	38	+ 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))
	40	+ self.lin0 = torch.nn.Linear(d_hid, d_hid)
	41	+ self.lin1 = torch.nn.Linear(d_hid, d_hid)
	42	+
	43	+ def forward(self, x, y=torch.zeros(DEFAULT_BATCH_SIZE, DEFAULT_DHID)):
	44	+ x = torch.mm(x, self.mm_param0)
	45	+ x = x + y
	46	+ x = self.lin0(x)
	47	+ x = torch.relu(x)
	48	+ pipe_split()
	49	+ x = torch.mm(x, self.mm_param1)
	50	+ x = self.lin1(x)
	51	+ x = torch.relu(x)
	52	+ return x
	53	+
	54	+
	55	+class ModelWithParamAlias(torch.nn.Module):
	56	+ default_dhid = 512
	57	+ default_batch_size = 256
	58	+
	59	+ def __init__(self, d_hid: int = default_dhid):
	60	+ super().__init__()
	61	+ self.mm_param1 = self.mm_param0 = torch.nn.Parameter(torch.randn(d_hid, d_hid))
	62	+ self.lin1 = self.lin0 = torch.nn.Linear(d_hid, d_hid)
	63	+
	64	+ def forward(self, x, y):
	65	+ x = torch.mm(x, self.mm_param0)
	66	+ x = x + y
	67	+ x = self.lin0(x)
	68	+ x = torch.relu(x)
	69	+ pipe_split()
	70	+ x = torch.mm(x, self.mm_param1)
	71	+ x = self.lin1(x)
	72	+ x = torch.relu(x)
	73	+ return x
	74	+
	75	+
	76	+# MLP Layer
	77	+class MLPModule(torch.nn.Module):
	78	+ def __init__(self, d_hid: int):
	79	+ super().__init__()
	80	+ self.net1 = torch.nn.Linear(d_hid, d_hid)
	81	+ self.relu = torch.nn.ReLU()
	82	+ self.net2 = torch.nn.Linear(d_hid, d_hid)
	83	+
	84	+ def forward(self, x):
	85	+ x = self.net1(x)
	86	+ x = self.relu(x)
	87	+ x = self.net2(x)
	88	+ return x
	89	+
	90	+
	91	+# Multi-MLP model
	92	+class MultiMLP(torch.nn.Module):
	93	+ def __init__(self, d_hid: int, n_layers: int = 2):
	94	+ super().__init__()
	95	+ self.layers = torch.nn.ModuleList([MLPModule(d_hid) for _ in range(n_layers)])
	96	+ # 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)}
	98	+
	99	+ def forward(self, x):
	100	+ for layer in self.layers:
	101	+ x = layer(x)
	102	+ return x
	103	+
	104	+
	105	+class CustomLinearDx(Function):
	106	+ @staticmethod
	107	+ # pylint:disable=huawei-too-many-arguments
	108	+ def forward(ctx, input_val, weight, bias, module, layer_idx):
	109	+ ctx.save_for_backward(input_val, weight, bias)
	110	+ ctx.module = module
	111	+ ctx.layer_idx = layer_idx
	112	+ return input_val.mm(weight.t()) + bias
	113	+
	114	+ @staticmethod
	115	+ def backward(ctx, grad_output):
	116	+ input_val, weight, bias = ctx.saved_tensors
	117	+ grad_input = grad_output.mm(weight)
	118	+ ctx.module.cached_context[ctx.layer_idx].append(grad_output.clone())
	119	+ ctx.module.cached_context[str(ctx.layer_idx) + "_input"].append(
	120	+ input_val.clone()
	121	+ )
	122	+ return grad_input, None, None, None, None
	123	+
	124	+
	125	+class CustomLinearDxDw(Function):
	126	+ @staticmethod
	127	+ def forward(ctx, input_val, weight, bias):
	128	+ ctx.save_for_backward(input_val, weight, bias)
	129	+ return input_val.mm(weight.t()) + bias
	130	+
	131	+ @staticmethod
	132	+ def backward(ctx, grad_output):
	133	+ input_val, weight, bias = ctx.saved_tensors
	134	+ grad_input = grad_output.mm(weight)
	135	+ grad_weight = grad_output.t().mm(input_val)
	136	+ grad_bias = grad_output.sum(0)
	137	+ return grad_input, grad_weight, grad_bias
	138	+
	139	+
	140	+class MLPModuleWithDw(torch.nn.Module):
	141	+ def __init__(self, d_hid: int):
	142	+ super().__init__()
	143	+ self.fc1_weight = torch.nn.Parameter(torch.randn(d_hid, d_hid))
	144	+ self.fc1_bias = torch.nn.Parameter(torch.randn(d_hid))
	145	+ self.fc2_weight = torch.nn.Parameter(torch.randn(d_hid, d_hid))
	146	+ self.fc2_bias = torch.nn.Parameter(torch.randn(d_hid))
	147	+
	148	+ torch.nn.init.uniform_(self.fc1_weight, -0.01, 0.01)
	149	+ torch.nn.init.uniform_(self.fc2_weight, -0.01, 0.01)
	150	+ torch.nn.init.uniform_(self.fc1_bias, -0.01, 0.01)
	151	+ torch.nn.init.uniform_(self.fc2_bias, -0.01, 0.01)
	152	+
	153	+ self.cached_context = {}
	154	+ self.cached_context["fc1"] = []
	155	+ self.cached_context["fc2"] = []
	156	+ self.cached_context["fc1_input"] = []
	157	+ self.cached_context["fc2_input"] = []
	158	+
	159	+ self.use_custom_logic = False
	160	+
	161	+ def forward(self, x):
	162	+ if not self.use_custom_logic:
	163	+ self.hidden = CustomLinearDxDw.apply(x, self.fc1_weight, self.fc1_bias)
	164	+ self.hidden = torch.nn.functional.relu(self.hidden)
	165	+ output = CustomLinearDxDw.apply(self.hidden, self.fc2_weight, self.fc2_bias)
	166	+ return output
	167	+
	168	+ self.hidden = CustomLinearDx.apply(
	169	+ x, self.fc1_weight, self.fc1_bias, self, "fc1"
	170	+ )
	171	+ self.hidden = torch.nn.functional.relu(self.hidden)
	172	+ output = CustomLinearDx.apply(
	173	+ self.hidden, self.fc2_weight, self.fc2_bias, self, "fc2"
	174	+ )
	175	+ return output
	176	+
	177	+ def compute_dW(self):
	178	+ grad_output_fc1 = self.cached_context["fc1"].pop(0)
	179	+ grad_output_fc2 = self.cached_context["fc2"].pop(0)
	180	+ cached_input_fc1 = self.cached_context["fc1_input"].pop(0)
	181	+ cached_input_fc2 = self.cached_context["fc2_input"].pop(0)
	182	+
	183	+ dW2 = grad_output_fc2.t().mm(cached_input_fc2)
	184	+ db2 = grad_output_fc2.sum(0)
	185	+
	186	+ dW1 = grad_output_fc1.t().mm(cached_input_fc1)
	187	+ db1 = grad_output_fc1.sum(0)
	188	+
	189	+ if self.fc1_weight.grad is not None:
	190	+ self.fc1_weight.grad += dW1
	191	+ self.fc1_bias.grad += db1
	192	+ self.fc2_weight.grad += dW2
	193	+ self.fc2_bias.grad += db2
	194	+ else:
	195	+ self.fc1_weight.grad = dW1
	196	+ self.fc1_bias.grad = db1
	197	+ self.fc2_weight.grad = dW2
	198	+ self.fc2_bias.grad = db2
	199	+
	200	+ def toggle(self):
	201	+ self.use_custom_logic = not self.use_custom_logic
	202	+
	203	+
	204	+# Multi-MLP model With Dw
	205	+class MultiMLPWithDw(torch.nn.Module):
	206	+ def __init__(self, d_hid: int, n_layers: int = 2):
	207	+ super().__init__()
	208	+ self.layers = torch.nn.ModuleList(
	209	+ [MLPModuleWithDw(d_hid) for _ in range(n_layers)]
	210	+ )
	211	+ # 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)}
	213	+ self.use_custom_logic = False
	214	+
	215	+ def forward(self, x):
	216	+ for layer in self.layers:
	217	+ x = layer(x)
	218	+ return x
	219	+
	220	+ def toggle(self):
	221	+ self.use_custom_logic = not self.use_custom_logic
	222	+ for layer in self.layers:
	223	+ layer.toggle()
	224	+
	225	+ def compute_dW(self):
	226	+ if not self.use_custom_logic:
	227	+ raise RuntimeError("Need to call toggle() to enable custom backward and dW")
	228	+
	229	+ for i in reversed(range(len(self.layers))):
	230	+ self.layers[i].compute_dW()

Atest/distributed/pipelining/schedule_registry.py+230-0

@@ -0,0 +1,230 @@
	1	+# Copyright (c) Meta Platforms, Inc. and affiliates
	2	+# Owner(s): ["oncall: distributed"]
	3	+# This file is a Schedule zoo for testing torch.distributed.pipelining.
	4	+# It includes schedules designed purely for testing purposes
	5	+from typing import Callable, Optional
	6	+
	7	+from torch.distributed.pipelining.schedules import (
	8	+ _Action,
	9	+ _ComputationType,
	10	+ _PipelineScheduleRuntime,
	11	+ PipelineScheduleMulti,
	12	+ RECV_B,
	13	+ RECV_F,
	14	+ SEND_B,
	15	+ SEND_F,
	16	+)
	17	+from torch.distributed.pipelining.stage import _PipelineStageBase
	18	+
	19	+
	20	+F = _ComputationType.FORWARD
	21	+B = _ComputationType.FULL_BACKWARD
	22	+W = _ComputationType.BACKWARD_WEIGHT
	23	+INPUT = _ComputationType.BACKWARD_INPUT
	24	+
	25	+
	26	+class ScheduleVShaped(PipelineScheduleMulti):
	27	+ n_stages = 4
	28	+ rank_stages = {
	29	+ 0: [0, 3],
	30	+ 1: [1, 2],
	31	+ }
	32	+
	33	+ def __init__(
	34	+ self,
	35	+ stages: list[_PipelineStageBase],
	36	+ n_microbatches: int,
	37	+ loss_fn: Optional[Callable] = None,
	38	+ scale_grads: bool = True,
	39	+ ):
	40	+ super().__init__(
	41	+ stages=stages,
	42	+ n_microbatches=n_microbatches,
	43	+ loss_fn=loss_fn,
	44	+ scale_grads=scale_grads,
	45	+ )
	46	+
	47	+ # Go through one microbatch
	48	+ # Note(whc) - it might be easier to work with thes schedules by writing them as a list of
	49	+ # ["0F0", ...] and then parsing them in the test infra to turn them into actions.
	50	+ self.pipeline_order = {
	51	+ 0: [
	52	+ _Action(0, F, 0),
	53	+ None,
	54	+ None,
	55	+ _Action(3, F, 0),
	56	+ _Action(3, B, 0),
	57	+ None,
	58	+ None,
	59	+ _Action(0, B, 0),
	60	+ ],
	61	+ 1: [
	62	+ None,
	63	+ _Action(1, F, 0),
	64	+ _Action(2, F, 0),
	65	+ None,
	66	+ None,
	67	+ _Action(2, B, 0),
	68	+ _Action(1, B, 0),
	69	+ None,
	70	+ ],
	71	+ }
	72	+ self._validate_and_set_stage_mapping(self.pipeline_order)
	73	+
	74	+
	75	+class ScheduleUnbalanced(PipelineScheduleMulti):
	76	+ n_stages = 5
	77	+ rank_stages = {
	78	+ 0: [0, 1, 4],
	79	+ 1: [2, 3],
	80	+ }
	81	+
	82	+ def __init__(
	83	+ self,
	84	+ stages: list[_PipelineStageBase],
	85	+ n_microbatches: int,
	86	+ loss_fn: Optional[Callable] = None,
	87	+ scale_grads: bool = True,
	88	+ ):
	89	+ super().__init__(
	90	+ stages=stages,
	91	+ n_microbatches=n_microbatches,
	92	+ loss_fn=loss_fn,
	93	+ scale_grads=scale_grads,
	94	+ )
	95	+
	96	+ self.pipeline_order = {
	97	+ 0: [
	98	+ _Action(0, F, 0),
	99	+ _Action(1, F, 0),
	100	+ None,
	101	+ None,
	102	+ _Action(4, F, 0),
	103	+ _Action(4, B, 0),
	104	+ None,
	105	+ None,
	106	+ _Action(1, B, 0),
	107	+ _Action(0, B, 0),
	108	+ ],
	109	+ 1: [
	110	+ None,
	111	+ None,
	112	+ _Action(2, F, 0),
	113	+ _Action(3, F, 0),
	114	+ None,
	115	+ None,
	116	+ _Action(3, B, 0),
	117	+ _Action(2, B, 0),
	118	+ None,
	119	+ None,
	120	+ ],
	121	+ }
	122	+ self._validate_and_set_stage_mapping(self.pipeline_order)
	123	+
	124	+
	125	+class ScheduleWithW(PipelineScheduleMulti):
	126	+ n_stages = 4
	127	+ num_microbatches = 2
	128	+ rank_stages = {
	129	+ 0: [0, 2],
	130	+ 1: [1, 3],
	131	+ }
	132	+
	133	+ def __init__(
	134	+ self,
	135	+ stages: list[_PipelineStageBase],
	136	+ n_microbatches: int,
	137	+ loss_fn: Optional[Callable] = None,
	138	+ enable_zero_bubble: bool = True,
	139	+ scale_grads: bool = True,
	140	+ ):
	141	+ super().__init__(
	142	+ stages=stages,
	143	+ n_microbatches=n_microbatches,
	144	+ loss_fn=loss_fn,
	145	+ scale_grads=scale_grads,
	146	+ )
	147	+
	148	+ # Needs to be updated as part of all schedules using "W"
	149	+ self.use_full_backward = False
	150	+
	151	+ # Go through two microbatches
	152	+ self.pipeline_order = {
	153	+ 0: [
	154	+ _Action(0, F, 0),
	155	+ _Action(0, F, 1),
	156	+ _Action(2, F, 0),
	157	+ _Action(2, F, 1),
	158	+ None,
	159	+ _Action(2, INPUT, 0),
	160	+ _Action(2, W, 0),
	161	+ _Action(0, INPUT, 0),
	162	+ _Action(2, INPUT, 1),
	163	+ _Action(0, W, 0),
	164	+ _Action(0, INPUT, 1),
	165	+ _Action(2, W, 1),
	166	+ _Action(0, W, 1),
	167	+ ],
	168	+ 1: [
	169	+ None,
	170	+ _Action(1, F, 0),
	171	+ _Action(1, F, 1),
	172	+ _Action(3, F, 0),
	173	+ _Action(3, INPUT, 0),
	174	+ _Action(3, F, 1),
	175	+ _Action(1, INPUT, 0),
	176	+ _Action(3, INPUT, 1),
	177	+ _Action(3, W, 0),
	178	+ _Action(1, INPUT, 1),
	179	+ _Action(1, W, 0),
	180	+ _Action(3, W, 1),
	181	+ _Action(1, W, 1),
	182	+ ],
	183	+ }
	184	+ self._validate_and_set_stage_mapping(self.pipeline_order)
	185	+
	186	+
	187	+class ScheduleWithReorderedB(_PipelineScheduleRuntime):
	188	+ n_stages = 2
	189	+ num_microbatches = 2
	190	+ rank_stages = {
	191	+ 0: [0],
	192	+ 1: [1],
	193	+ }
	194	+
	195	+ def __init__(
	196	+ self,
	197	+ stages: list[_PipelineStageBase],
	198	+ n_microbatches: int,
	199	+ loss_fn: Optional[Callable] = None,
	200	+ scale_grads: bool = True,
	201	+ ):
	202	+ super().__init__(
	203	+ stages=stages,
	204	+ n_microbatches=n_microbatches,
	205	+ loss_fn=loss_fn,
	206	+ scale_grads=scale_grads,
	207	+ )
	208	+ # Go through two microbatches
	209	+ self.pipeline_order_with_comms = {
	210	+ 0: [
	211	+ _Action(0, F, 0),
	212	+ _Action(0, F, 1),
	213	+ _Action(0, SEND_F, 0),
	214	+ _Action(0, SEND_F, 1),
	215	+ _Action(0, RECV_B, 0),
	216	+ _Action(0, RECV_B, 1),
	217	+ _Action(0, B, 0),
	218	+ _Action(0, B, 1),
	219	+ ],
	220	+ 1: [
	221	+ _Action(1, RECV_F, 0),
	222	+ _Action(1, RECV_F, 1),
	223	+ _Action(1, F, 0),
	224	+ _Action(1, F, 1),
	225	+ _Action(1, B, 0),
	226	+ _Action(1, B, 1),
	227	+ _Action(1, SEND_B, 0),
	228	+ _Action(1, SEND_B, 1),
	229	+ ],
	230	+ }

Atest/distributed/pipelining/test_backward.py+187-0

@@ -0,0 +1,187 @@
	1	+# Copyright (c) Meta Platforms, Inc. and affiliates
	2	+# Owner(s): ["oncall: distributed"]
	3	+import copy
	4	+
	5	+from model_registry import MLPModule
	6	+
	7	+import torch
	8	+from torch.distributed.pipelining._backward import (
	9	+ stage_backward,
	10	+ stage_backward_input,
	11	+ stage_backward_weight,
	12	+)
	13	+from torch.testing._internal.common_utils import run_tests, TestCase
	14	+
	15	+
	16	+d_hid = 512
	17	+batch_size = 256
	18	+
	19	+
	20	+class StageBackwardTests(TestCase):
	21	+ def test_stage_backward(self):
	22	+ # MLP as a stage module
	23	+ mod = MLPModule(d_hid)
	24	+ x = torch.randn(batch_size, d_hid)
	25	+ # As in a pipeline stage, the inputs to this stage requires gradients
	26	+ x.requires_grad_(True)
	27	+ target = torch.randn(batch_size, d_hid)
	28	+ loss_fn = torch.nn.MSELoss(reduction="sum")
	29	+
	30	+ # Make a copy
	31	+ ref_mod = copy.deepcopy(mod)
	32	+ ref_x = x.detach().requires_grad_(x.requires_grad)
	33	+ ref_target = target.detach()
	34	+
	35	+ # Forward and backward in stage manner
	36	+ out = mod(x)
	37	+ loss = loss_fn(out, target)
	38	+ grad_inputs = stage_backward(
	39	+ stage_output=loss,
	40	+ output_grads=None,
	41	+ input_values=(x,),
	42	+ )
	43	+
	44	+ # Run reference
	45	+ ref_out = ref_mod(ref_x)
	46	+ ref_loss = loss_fn(ref_out, ref_target)
	47	+ ref_loss.backward()
	48	+
	49	+ torch.testing.assert_close(grad_inputs[0], ref_x.grad)
	50	+
	51	+ # Every rank checks gradients
	52	+ for name, p in mod.named_parameters():
	53	+ ref_p = ref_mod.get_parameter(name)
	54	+ try:
	55	+ torch.testing.assert_close(p.grad, ref_p.grad)
	56	+ except AssertionError:
	57	+ print(f"Gradient test failed for {name}: {p.grad} vs {ref_p.grad}")
	58	+ raise
	59	+
	60	+ def test_stage_backward_input(self):
	61	+ # MLP as a stage module
	62	+ mod = MLPModule(d_hid)
	63	+ x = torch.randn(batch_size, d_hid)
	64	+ # As in a pipeline stage, the inputs to this stage requires gradients
	65	+ x.requires_grad_(True)
	66	+ target = torch.randn(batch_size, d_hid)
	67	+ loss_fn = torch.nn.MSELoss(reduction="sum")
	68	+
	69	+ # Make a copy
	70	+ ref_mod = copy.deepcopy(mod)
	71	+ ref_x = x.detach().requires_grad_(x.requires_grad)
	72	+ ref_target = target.detach()
	73	+
	74	+ # Forward, then backward of loss with respect to inputs
	75	+ out = mod(x)
	76	+ loss = loss_fn(out, target)
	77	+ dinputs, param_groups = stage_backward_input(
	78	+ stage_outputs_or_loss=(loss,),
	79	+ output_grads=None,
	80	+ input_values=[x],
	81	+ weights=mod.parameters(),
	82	+ )
	83	+
	84	+ # Run reference
	85	+ ref_out = ref_mod(ref_x)
	86	+ ref_loss = loss_fn(ref_out, ref_target)
	87	+ ref_loss.backward()
	88	+
	89	+ torch.testing.assert_close(x.grad, ref_x.grad)
	90	+ torch.testing.assert_close(dinputs[0], ref_x.grad)
	91	+ for _, p in mod.named_parameters():
	92	+ # Check that the weight gradients were not updated
	93	+ self.assertEqual(p.grad, None)
	94	+
	95	+ def test_stage_backward_weight(self):
	96	+ # MLP as a stage module
	97	+ mod = MLPModule(d_hid)
	98	+ x = torch.randn(batch_size, d_hid)
	99	+ # As in a pipeline stage, the inputs to this stage requires gradients
	100	+ x.requires_grad_(True)
	101	+ target = torch.randn(batch_size, d_hid)
	102	+ loss_fn = torch.nn.MSELoss(reduction="sum")
	103	+
	104	+ # Make a copy
	105	+ ref_mod = copy.deepcopy(mod)
	106	+ ref_x = x.detach().requires_grad_(x.requires_grad)
	107	+ ref_target = target.detach()
	108	+
	109	+ # Forward, then backward of loss with respect to inputs
	110	+ out = mod(x)
	111	+ loss = loss_fn(out, target)
	112	+ dinputs, param_groups = stage_backward_input(
	113	+ stage_outputs_or_loss=(loss,),
	114	+ output_grads=None,
	115	+ input_values=[x],
	116	+ weights=mod.parameters(),
	117	+ )
	118	+
	119	+ # backward of loss with respect to weights
	120	+ stage_backward_weight(mod.parameters(), param_groups, retain_graph=True)
	121	+
	122	+ # Run reference
	123	+ ref_out = ref_mod(ref_x)
	124	+ ref_loss = loss_fn(ref_out, ref_target)
	125	+ ref_loss.backward()
	126	+
	127	+ # Every rank checks gradients
	128	+ for name, p in mod.named_parameters():
	129	+ ref_p = ref_mod.get_parameter(name)
	130	+ try:
	131	+ torch.testing.assert_close(p.grad, ref_p.grad)
	132	+ except AssertionError:
	133	+ print(f"Gradient test failed for {name}: {p.grad} vs {ref_p.grad}")
	134	+ raise
	135	+
	136	+ def test_stage_backward_weight_multiple_iters(self):
	137	+ # MLP as a stage module
	138	+ mod = MLPModule(d_hid)
	139	+ inputs = []
	140	+ for _ in range(10):
	141	+ x = torch.randn(batch_size, d_hid)
	142	+ inputs.append(x)
	143	+ # As in a pipeline stage, the inputs to this stage requires gradients
	144	+ x.requires_grad_(True)
	145	+
	146	+ target = torch.randn(batch_size, d_hid)
	147	+ loss_fn = torch.nn.MSELoss(reduction="sum")
	148	+
	149	+ # Make a copy
	150	+ ref_mod = copy.deepcopy(mod)
	151	+ ref_inputs = []
	152	+ for x in inputs:
	153	+ ref_inputs.append(x.detach().requires_grad_(x.requires_grad))
	154	+ ref_target = target.detach()
	155	+
	156	+ # Forward, then backward of loss with respect to inputs
	157	+ for x in inputs:
	158	+ out = mod(x)
	159	+ loss = loss_fn(out, target)
	160	+ dinputs, param_groups = stage_backward_input(
	161	+ stage_outputs_or_loss=(loss,),
	162	+ output_grads=None,
	163	+ input_values=[x],
	164	+ weights=mod.parameters(),
	165	+ )
	166	+
	167	+ # backward of loss with respect to weights
	168	+ stage_backward_weight(mod.parameters(), param_groups)
	169	+
	170	+ # Run reference
	171	+ for ref_x in ref_inputs:
	172	+ ref_out = ref_mod(ref_x)
	173	+ ref_loss = loss_fn(ref_out, ref_target)
	174	+ ref_loss.backward()
	175	+
	176	+ # Every rank checks gradients
	177	+ for name, p in mod.named_parameters():
	178	+ ref_p = ref_mod.get_parameter(name)
	179	+ try:
	180	+ torch.testing.assert_close(p.grad, ref_p.grad)
	181	+ except AssertionError:
	182	+ print(f"Gradient test failed for {name}: {p.grad} vs {ref_p.grad}")
	183	+ raise
	184	+
	185	+
	186	+if __name__ == "__main__":
	187	+ run_tests()

Atest/distributed/pipelining/test_microbatch.py+91-0

@@ -0,0 +1,91 @@
	1	+# Copyright (c) Meta Platforms, Inc. and affiliates
	2	+# Owner(s): ["oncall: distributed"]
	3	+from model_registry import ModelWithKwargs
	4	+
	5	+import torch
	6	+from torch.distributed.pipelining import pipeline
	7	+from torch.distributed.pipelining.microbatch import (
	8	+ merge_chunks,
	9	+ split_args_kwargs_into_chunks,
	10	+ TensorChunkSpec,
	11	+)
	12	+from torch.testing._internal.common_utils import run_tests, TestCase
	13	+
	14	+
	15	+d_hid = 512
	16	+torch.manual_seed(0)
	17	+
	18	+
	19	+class MicrobatchTests(TestCase):
	20	+ def test_split_and_merge(self):
	21	+ x0 = torch.randn(128, d_hid)
	22	+ x1 = torch.randn(256, d_hid)
	23	+ x2 = torch.randn(512, d_hid)
	24	+
	25	+ args = (x0, x1, x2)
	26	+ kwargs = {"x0": x0, "x1": x1, "x2": x2}
	27	+
	28	+ # Default chunking: dim 0
	29	+ arg_chunks, kwarg_chunks = split_args_kwargs_into_chunks(args, kwargs, 2)
	30	+ assert len(arg_chunks) == 2
	31	+ assert len(kwarg_chunks) == 2
	32	+ assert arg_chunks[0][0].shape == torch.Size([64, d_hid])
	33	+ assert arg_chunks[1][0].shape == torch.Size([64, d_hid])
	34	+ assert arg_chunks[0][1].shape == torch.Size([128, d_hid])
	35	+ assert arg_chunks[0][2].shape == torch.Size([256, d_hid])
	36	+ assert kwarg_chunks[0]["x0"].shape == torch.Size([64, d_hid])
	37	+ assert kwarg_chunks[0]["x1"].shape == torch.Size([128, d_hid])
	38	+ assert kwarg_chunks[1]["x2"].shape == torch.Size([256, d_hid])
	39	+
	40	+ # Merge chunks back together
	41	+ merged_args = merge_chunks(
	42	+ arg_chunks,
	43	+ (TensorChunkSpec(0), TensorChunkSpec(0), TensorChunkSpec(0)),
	44	+ )
	45	+ torch.testing.assert_close(merged_args, args)
	46	+
	47	+ merged_kwargs = merge_chunks(
	48	+ kwarg_chunks,
	49	+ {
	50	+ "x0": TensorChunkSpec(0),
	51	+ "x1": TensorChunkSpec(0),
	52	+ "x2": TensorChunkSpec(0),
	53	+ },
	54	+ )
	55	+ torch.testing.assert_close(merged_kwargs, kwargs)
	56	+ print("Microbatch test passed")
	57	+
	58	+ def test_chunk_spec(self):
	59	+ mod = ModelWithKwargs()
	60	+ batch_size = ModelWithKwargs.DEFAULT_BATCH_SIZE
	61	+
	62	+ x = torch.randn(batch_size, d_hid)
	63	+ y = torch.randn(batch_size, d_hid)
	64	+
	65	+ num_chunks = 4
	66	+
	67	+ args_chunk_spec = TensorChunkSpec.from_tuple((0,))
	68	+ kwargs_chunk_spec = TensorChunkSpec.from_dict({"y": 0})
	69	+
	70	+ args_split, kwargs_split = split_args_kwargs_into_chunks(
	71	+ (x,),
	72	+ {"y": y},
	73	+ num_chunks,
	74	+ args_chunk_spec,
	75	+ kwargs_chunk_spec,
	76	+ )
	77	+
	78	+ pipe = pipeline(
	79	+ mod,
	80	+ mb_args=args_split[0],
	81	+ mb_kwargs=kwargs_split[0],
	82	+ )
	83	+
	84	+ ref = mod(x, y)
	85	+ out = pipe(x, y)[0]
	86	+ torch.testing.assert_close(out, ref)
	87	+ print(f"equivalence test passed {torch.sum(out)} ref {torch.sum(ref)}")
	88	+
	89	+
	90	+if __name__ == "__main__":
	91	+ run_tests()

Atest/distributed/pipelining/test_pipe.py+123-0

@@ -0,0 +1,123 @@
	1	+# Copyright (c) Meta Platforms, Inc. and affiliates
	2	+# Owner(s): ["oncall: distributed"]
	3	+from model_registry import MLPModule, ModelWithParamAlias
	4	+
	5	+import torch
	6	+from torch.distributed.pipelining import pipe_split, pipeline
	7	+from torch.testing._internal.common_utils import (
	8	+ instantiate_parametrized_tests,
	9	+ parametrize,
	10	+ run_tests,
	11	+ TestCase,
	12	+)
	13	+
	14	+
	15	+d_hid = 512
	16	+microbatch_size = 16
	17	+
	18	+torch.manual_seed(0)
	19	+
	20	+
	21	+# Basic example
	22	+class ExampleCode(torch.nn.Module):
	23	+ def __init__(self) -> None:
	24	+ super().__init__()
	25	+ self.mm_param1 = torch.nn.Parameter(torch.randn(d_hid, d_hid))
	26	+ self.mm_param2 = torch.nn.Parameter(torch.randn(d_hid, d_hid))
	27	+ self.lin1 = torch.nn.Linear(d_hid, d_hid)
	28	+ self.lin2 = torch.nn.Linear(d_hid, d_hid)
	29	+
	30	+ def forward(self, x, y):
	31	+ x = torch.mm(x, self.mm_param1) # mutli-use param
	32	+ skip_connection = x
	33	+ x = x + y
	34	+ x = torch.relu(x)
	35	+ pipe_split()
	36	+ x = torch.mm(x, self.mm_param1) # mutli-use param
	37	+ x = self.lin1(x)
	38	+ pipe_split()
	39	+ x = torch.relu(x)
	40	+ x = x + skip_connection
	41	+ x = torch.mm(x, self.mm_param2)
	42	+ pipe_split()
	43	+ x = self.lin2(x)
	44	+ x = torch.relu(x)
	45	+ return x
	46	+
	47	+
	48	+class MultiMLP(torch.nn.Module):
	49	+ def __init__(self) -> None:
	50	+ super().__init__()
	51	+ self.mlp0 = MLPModule(d_hid)
	52	+ self.mlp1 = MLPModule(d_hid)
	53	+ self.mlp2 = MLPModule(d_hid)
	54	+ self.mlp3 = MLPModule(d_hid)
	55	+
	56	+ def forward(self, x, y):
	57	+ x = self.mlp0(x)
	58	+ pipe_split()
	59	+ x = self.mlp1(x)
	60	+ pipe_split()
	61	+ x = self.mlp2(x)
	62	+ pipe_split()
	63	+ x = self.mlp3(x)
	64	+ return x - y
	65	+
	66	+
	67	+EXPECTED_N_STAGES = {
	68	+ ExampleCode: 4,
	69	+ MultiMLP: 4,
	70	+ ModelWithParamAlias: 2,
	71	+}
	72	+
	73	+# Currently, we don't enforce full set equality on the FQNs between the original
	74	+# and pipelined models, because in the multi-use param case, PP will deduplicate
	75	+# the FQNs from the state_dict.
	76	+CHECK_FQN_SET_EQUALITY = False
	77	+
	78	+
	79	+class PipeTests(TestCase):
	80	+ @parametrize("ModelClass", [ExampleCode, MultiMLP, ModelWithParamAlias])
	81	+ def test_model_split(self, ModelClass):
	82	+ mod = ModelClass()
	83	+ x = torch.randn(microbatch_size, d_hid)
	84	+ y = torch.randn(microbatch_size, d_hid)
	85	+
	86	+ pipe = pipeline(
	87	+ mod,
	88	+ mb_args=(x, y),
	89	+ )
	90	+
	91	+ assert (
	92	+ pipe.num_stages == EXPECTED_N_STAGES[ModelClass]
	93	+ ), f"nstages = {pipe.num_stages}, expect {EXPECTED_N_STAGES[ModelClass]}"
	94	+
	95	+ ref_out = mod(x, y)
	96	+ out = pipe(x, y)[0]
	97	+ torch.testing.assert_close(out, ref_out)
	98	+ print(f"equivalence test passed {torch.sum(out)} ref {torch.sum(ref_out)}")
	99	+
	100	+ # Check qualname
	101	+ # state_dict.keys include both parameters and persistent buffers
	102	+ old_names = set(mod.state_dict().keys())
	103	+ new_names = set()
	104	+ for idx in range(pipe.num_stages):
	105	+ stage_mod = pipe.get_stage_module(idx)
	106	+ stage_fqns = set(stage_mod.state_dict().keys())
	107	+ assert stage_fqns.issubset(old_names)
	108	+ new_names.update(stage_fqns)
	109	+
	110	+ if CHECK_FQN_SET_EQUALITY:
	111	+ assert (
	112	+ old_names == new_names
	113	+ ), f"""
	114	+ old names {old_names}
	115	+ new names {new_names}
	116	+ """
	117	+ print("Qualname check passed")
	118	+
	119	+
	120	+instantiate_parametrized_tests(PipeTests)
	121	+
	122	+if __name__ == "__main__":
	123	+ run_tests()

Atest/distributed/pipelining/test_schedule.py+993-0

@@ -0,0 +1,993 @@
	1	+# Copyright (c) Meta Platforms, Inc. and affiliates
	2	+# Owner(s): ["oncall: distributed"]
	3	+import copy
	4	+import csv
	5	+import logging
	6	+import os
	7	+
	8	+from model_registry import MultiMLP
	9	+
	10	+import torch
	11	+from torch.distributed.pipelining import (
	12	+ Schedule1F1B,
	13	+ ScheduleGPipe,
	14	+ ScheduleInterleaved1F1B,
	15	+ ScheduleInterleavedZeroBubble,
	16	+ ScheduleLoopedBFS,
	17	+)
	18	+from torch.distributed.pipelining._utils import generate_stage_to_rank_mapping
	19	+from torch.distributed.pipelining.schedules import (
	20	+ _Action,
	21	+ _add_send_recv,
	22	+ _add_unshard_reshard,
	23	+ _format_pipeline_order,
	24	+ _merge_bw,
	25	+ _PipelineSchedule,
	26	+ _PipelineScheduleRuntime,
	27	+ _simulate_comms_compute,
	28	+ _validate_schedule,
	29	+ B,
	30	+ F,
	31	+ get_schedule_class,
	32	+ I,
	33	+ PipelineScheduleSingle,
	34	+ RECV_F,
	35	+ RESHARD,
	36	+ SEND_B,
	37	+ UNSHARD,
	38	+ W,
	39	+)
	40	+from torch.distributed.pipelining.stage import _PipelineStageBase, PipelineStage
	41	+from torch.testing._internal.common_utils import (
	42	+ check_leaked_tensors,
	43	+ instantiate_parametrized_tests,
	44	+ parametrize,
	45	+ run_tests,
	46	+ TestCase,
	47	+)
	48	+from torch.testing._internal.distributed.fake_pg import FakeStore
	49	+
	50	+
	51	+ARTIFACTS_DIR = os.path.join(os.path.dirname(os.path.abspath(__file__)), "artifacts")
	52	+
	53	+logger = logging.getLogger(__name__)
	54	+torch.manual_seed(0)
	55	+
	56	+
	57	+class MockPipelineStage(_PipelineStageBase):
	58	+ def __init__(self, args, *kwargs):
	59	+ # Mock the necessary attributes
	60	+ self.submod = None
	61	+ self.num_stages = kwargs.get("num_stages", 1)
	62	+ self.group_size = kwargs.get("group_size", 1)
	63	+ self.group_rank = kwargs.get("group_rank", 0)
	64	+ self.group = kwargs.get("group", None)
	65	+
	66	+ def _create_grad_recv_info(self, args, *kwargs):
	67	+ return None
	68	+
	69	+ def _prepare_forward_infra(self, n_microbatches):
	70	+ pass
	71	+
	72	+ def _prepare_backward_infra(self, n_microbatches):
	73	+ pass
	74	+
	75	+
	76	+class ScheduleTest(TestCase):
	77	+ def test_get_schedule_class(self):
	78	+ # List of all expected schedule names
	79	+ schedule_names = [
	80	+ "1F1B",
	81	+ "1f1b",
	82	+ "Interleaved1F1B",
	83	+ "INTERLEAVED1F1B",
	84	+ "GPipe",
	85	+ "LoopedBFS",
	86	+ "PipelineScheduleSingle",
	87	+ "PipelineScheduleMulti",
	88	+ ]
	89	+
	90	+ # Test each schedule name
	91	+ for name in schedule_names:
	92	+ with self.subTest(name=name):
	93	+ schedule_class = get_schedule_class(name)
	94	+ self.assertIsNotNone(
	95	+ schedule_class, f"Class for {name} should not be None"
	96	+ )
	97	+ self.assertTrue(
	98	+ issubclass(schedule_class, _PipelineSchedule),
	99	+ f"{name} should be a subclass of _PipelineSchedule",
	100	+ )
	101	+
	102	+ error_case = ["ScheduleThatDoesNotExist"]
	103	+ for name in error_case:
	104	+ # Test that the original name is included in the error message
	105	+ with self.assertRaisesRegex(ValueError, f"{name}"):
	106	+ get_schedule_class(name)
	107	+
	108	+ @parametrize(
	109	+ "ScheduleClass",
	110	+ [
	111	+ Schedule1F1B,
	112	+ ScheduleGPipe,
	113	+ ScheduleInterleaved1F1B,
	114	+ ScheduleInterleavedZeroBubble,
	115	+ ScheduleLoopedBFS,
	116	+ ],
	117	+ )
	118	+ def test_schedule_with_single_stage(self, ScheduleClass):
	119	+ """
	120	+ Test that schedules with only a single stage work as expected for all schedules.
	121	+ """
	122	+ store = FakeStore()
	123	+ torch.distributed.init_process_group(
	124	+ backend="fake", rank=0, world_size=1, store=store
	125	+ )
	126	+ d_hid, batch_size = 512, 256
	127	+ n_stages = 1
	128	+ device = "cpu"
	129	+ full_mod = MultiMLP(d_hid, n_layers=n_stages)
	130	+ full_mod.to(device)
	131	+
	132	+ x = torch.randn(batch_size, d_hid, device=device)
	133	+ ref_mod = copy.deepcopy(full_mod)
	134	+ with torch.no_grad():
	135	+ y = ref_mod(x)
	136	+ # Add a small perturbation
	137	+ target = y + torch.randn(batch_size, d_hid, device=device)
	138	+
	139	+ def loss_fn(y, target):
	140	+ return torch.nn.functional.cross_entropy(y, target)
	141	+
	142	+ # Run reference
	143	+ for _ in range(2):
	144	+ ref_mod.zero_grad()
	145	+ ref_out = ref_mod(x)
	146	+ ref_loss = loss_fn(ref_out, target)
	147	+ ref_loss.backward()
	148	+
	149	+ submod_name = "layers.0"
	150	+ stage_module = full_mod.get_submodule(submod_name)
	151	+
	152	+ # Create a pipeline stage to wrap that submodule
	153	+ num_microbatches = 2
	154	+ stages = [
	155	+ PipelineStage(
	156	+ stage_module,
	157	+ 0,
	158	+ n_stages,
	159	+ device,
	160	+ )
	161	+ ]
	162	+
	163	+ if issubclass(ScheduleClass, PipelineScheduleSingle):
	164	+ stages = stages[0]
	165	+
	166	+ # Attach to a schedule
	167	+ schedule = ScheduleClass(
	168	+ stages,
	169	+ num_microbatches,
	170	+ loss_fn=loss_fn,
	171	+ )
	172	+ # Run
	173	+ for _ in range(2):
	174	+ # Zero gradients
	175	+ stage_module.zero_grad()
	176	+ losses = []
	177	+ out = schedule.step(x, target=target, losses=losses)
	178	+
	179	+ # Check output
	180	+ torch.testing.assert_close(out, ref_out)
	181	+ # Check loss
	182	+ # Since the reduction used in the loss function above is "mean", we use
	183	+ # "mean" here to reduce microbatch losses into a single value too.
	184	+ pipe_loss = torch.stack(losses).mean()
	185	+ torch.testing.assert_close(pipe_loss, ref_loss)
	186	+
	187	+ # Check gradients
	188	+ # Get corresponding submodule from reference model
	189	+ ref_submod = ref_mod.get_submodule(submod_name)
	190	+ # Check gradients per parameter
	191	+ for name, p in stage_module.named_parameters():
	192	+ ref_p = ref_submod.get_parameter(name)
	193	+ try:
	194	+ torch.testing.assert_close(p.grad, ref_p.grad, rtol=1e-5, atol=4e-5)
	195	+ except AssertionError:
	196	+ print(f"Gradient test failed for {name}: {p.grad} vs {ref_p.grad}")
	197	+ raise
	198	+
	199	+ torch.distributed.destroy_process_group()
	200	+
	201	+ def test_zero_bubble_schedule_errors_with_compile(self):
	202	+ """
	203	+ Test that zero bubble schedules raise an error when used with torch.compile.
	204	+ """
	205	+ store = FakeStore()
	206	+ torch.distributed.init_process_group(
	207	+ backend="fake", rank=0, world_size=1, store=store
	208	+ )
	209	+ n_stages = 1
	210	+ device = torch.device("cpu")
	211	+ model = MultiMLP(8, n_layers=n_stages)
	212	+ # full_mod
	213	+ compiled_model = torch.compile(model)
	214	+ stage = PipelineStage(
	215	+ compiled_model,
	216	+ 0,
	217	+ n_stages,
	218	+ device,
	219	+ )
	220	+ with self.assertRaises(RuntimeError):
	221	+ ScheduleInterleavedZeroBubble([stage], 2)
	222	+
	223	+ torch.distributed.destroy_process_group()
	224	+
	225	+
	226	+instantiate_parametrized_tests(ScheduleTest)
	227	+
	228	+
	229	+class TestSchedulePlan(TestCase):
	230	+ def setUp(self):
	231	+ # Define a list of test cases with varying num_local_stages, num_microbatches, and group_size
	232	+ # These should succeed since num_microbatches % group_size == 0
	233	+ self.test_cases = [
	234	+ # small number of stages
	235	+ (2, 2, 2),
	236	+ (2, 4, 4),
	237	+ (2, 8, 2),
	238	+ (2, 8, 4),
	239	+ (2, 8, 8),
	240	+ (4, 4, 4),
	241	+ (4, 8, 4),
	242	+ (4, 8, 8),
	243	+ # large microbatches
	244	+ (4, 16, 4),
	245	+ (4, 32, 4),
	246	+ (4, 64, 4),
	247	+ # large groups
	248	+ (4, 16, 16),
	249	+ (4, 32, 32),
	250	+ (4, 128, 64),
	251	+ # odd num pipeline stages
	252	+ (3, 2, 2),
	253	+ (3, 8, 2),
	254	+ (3, 12, 4),
	255	+ # odd group_sizes
	256	+ (4, 6, 3),
	257	+ (4, 10, 5),
	258	+ # n_mb non divisible by group_size
	259	+ (2, 3, 4),
	260	+ (2, 4, 4),
	261	+ (2, 10, 4),
	262	+ (2, 15, 4),
	263	+ ]
	264	+
	265	+ @parametrize(
	266	+ "ScheduleClass",
	267	+ [ScheduleInterleaved1F1B, ScheduleLoopedBFS],
	268	+ )
	269	+ def test_pipeline_order(self, ScheduleClass):
	270	+ for num_local_stages, num_microbatches, group_size in self.test_cases:
	271	+ with self.subTest(
	272	+ num_local_stages=num_local_stages,
	273	+ num_microbatches=num_microbatches,
	274	+ group_size=group_size,
	275	+ ):
	276	+ if num_microbatches % group_size != 0:
	277	+ continue
	278	+
	279	+ logger.info(
	280	+ "num_local_stages=%d num_microbatches=%d group_size=%d",
	281	+ num_local_stages,
	282	+ num_microbatches,
	283	+ group_size,
	284	+ )
	285	+ num_stages = num_local_stages * group_size
	286	+ stages = [
	287	+ MockPipelineStage(group_size=group_size, num_stages=num_stages)
	288	+ for i in range(num_local_stages)
	289	+ ]
	290	+
	291	+ schedule = ScheduleClass(stages, num_microbatches)
	292	+ _formatted_pipeline_order = _format_pipeline_order(
	293	+ schedule.pipeline_order
	294	+ )
	295	+
	296	+ def stage_to_rank(stage):
	297	+ return stage % group_size
	298	+
	299	+ comms_sch = _add_send_recv(
	300	+ schedule.pipeline_order,
	301	+ stage_to_rank=stage_to_rank,
	302	+ num_stages=num_stages,
	303	+ )
	304	+ _simulate_comms_compute(
	305	+ comms_sch,
	306	+ stage_to_rank=stage_to_rank,
	307	+ num_stages=num_stages,
	308	+ )
	309	+
	310	+ @parametrize(
	311	+ "ScheduleClass",
	312	+ [ScheduleInterleaved1F1B, ScheduleInterleavedZeroBubble],
	313	+ )
	314	+ def test_pipeline_order_flex_and_zero_bubble(self, ScheduleClass):
	315	+ for num_local_stages, num_microbatches, group_size in self.test_cases:
	316	+ with self.subTest(
	317	+ num_local_stages=num_local_stages,
	318	+ num_microbatches=num_microbatches,
	319	+ group_size=group_size,
	320	+ ):
	321	+ warmups_ops_last_stage = (num_local_stages - 1) * (
	322	+ num_microbatches // max(1, num_microbatches // group_size)
	323	+ )
	324	+ warmup_ops = warmups_ops_last_stage + 2 * (group_size - 1)
	325	+ warmup_ops = min(warmup_ops, num_microbatches * num_local_stages)
	326	+
	327	+ num_stages = num_local_stages * group_size
	328	+ stages = [
	329	+ MockPipelineStage(group_size=group_size, num_stages=num_stages)
	330	+ for i in range(num_local_stages)
	331	+ ]
	332	+ schedule = ScheduleClass(stages, num_microbatches)
	333	+ _format_pipeline_order(schedule.pipeline_order)
	334	+
	335	+ def stage_to_rank(stage):
	336	+ return stage % group_size
	337	+
	338	+ comms_sch = _add_send_recv(
	339	+ schedule.pipeline_order,
	340	+ stage_to_rank=stage_to_rank,
	341	+ num_stages=num_stages,
	342	+ )
	343	+ _simulate_comms_compute(
	344	+ comms_sch,
	345	+ stage_to_rank=stage_to_rank,
	346	+ num_stages=num_stages,
	347	+ )
	348	+
	349	+
	350	+instantiate_parametrized_tests(TestSchedulePlan)
	351	+
	352	+
	353	+class TestScheduleLowering(TestCase):
	354	+ """Tests lowering passes that convert simple compute-only (FBW) schedules into compute+comms schedules"""
	355	+
	356	+ def _parse_actions(self, actions: list[str]) -> list[_Action]:
	357	+ return [_Action.from_str(s) for s in actions]
	358	+
	359	+ @parametrize(
	360	+ "action_str_and_ref",
	361	+ [
	362	+ ("1F0", _Action(1, F, 0)),
	363	+ ("2I1", _Action(2, I, 1)),
	364	+ ("0W3", _Action(0, W, 3)),
	365	+ ("0B3", _Action(0, B, 3)),
	366	+ ("1UNSHARD", _Action(1, UNSHARD, None)),
	367	+ ("3RESHARD", _Action(3, RESHARD, None)),
	368	+ ("2SEND_B2", _Action(2, SEND_B, 2)),
	369	+ ("1RECV_F1", _Action(1, RECV_F, 1)),
	370	+ ],
	371	+ )
	372	+ def test_action_parse(self, action_str_and_ref):
	373	+ """Test that actions can be parsed from strings and round-tripped back to the same strings."""
	374	+ act_str, ref = action_str_and_ref
	375	+ act = _Action.from_str(act_str)
	376	+ self.assertEqual(act, ref)
	377	+ self.assertEqual(act_str, act.__repr__())
	378	+
	379	+ @parametrize(
	380	+ "test_info",
	381	+ [
	382	+ {
	383	+ "compute": ["0F0", "0F1", " ", "0B0", "0B1"],
	384	+ "comms": ["0UNSHARD", "0F0", "0F1", "0B0", "0B1", "0RESHARD"],
	385	+ },
	386	+ ],
	387	+ )
	388	+ def test_unshard_reshard(self, test_info):
	389	+ """Test the lowering pass that takes a 'compute only' schedule (with only F,B,W ops) and adds
	390	+ FSDP unshard/reshard operations to the schedule. This is just part of the process of adding communication
	391	+ ops and producing a complete schedule.
	392	+ """
	393	+ compute_sch = self._parse_actions(test_info["compute"])
	394	+ expected_comms_sch = self._parse_actions(test_info["comms"])
	395	+
	396	+ comms_sch = _add_unshard_reshard(compute_sch)
	397	+ for expected, actual in zip(expected_comms_sch, comms_sch):
	398	+ self.assertEqual(
	399	+ expected,
	400	+ actual,
	401	+ (
	402	+ f"Mismatch: expected action {expected} but found {actual}."
	403	+ f"\nWhole Schedule: {comms_sch}"
	404	+ ),
	405	+ )
	406	+
	407	+ @parametrize(
	408	+ "test_info",
	409	+ [
	410	+ {
	411	+ "compute": [
	412	+ "0F0",
	413	+ "0F1",
	414	+ "0F2",
	415	+ "0I0",
	416	+ "0I1",
	417	+ "0W0",
	418	+ "0I2",
	419	+ "0W2",
	420	+ "0W1",
	421	+ ],
	422	+ "comms": ["0F0", "0F1", "0F2", "0I0", "0I1", "0W0", "0B2", "0W1"],
	423	+ },
	424	+ ],
	425	+ )
	426	+ def test_merge_bw(self, test_info):
	427	+ """Test the pass that merges adjacent I and W operations into a B operation."""
	428	+ compute_sch = self._parse_actions(test_info["compute"])
	429	+ expected_merged_sch = self._parse_actions(test_info["comms"])
	430	+
	431	+ merged_sch = _merge_bw(compute_sch)
	432	+ for expected, actual in zip(expected_merged_sch, merged_sch):
	433	+ self.assertEqual(
	434	+ expected,
	435	+ actual,
	436	+ (
	437	+ f"Mismatch: expected action {expected} but found {actual}."
	438	+ f"\nWhole Schedule: {merged_sch}"
	439	+ ),
	440	+ )
	441	+
	442	+ @parametrize(
	443	+ "test_info",
	444	+ [
	445	+ {
	446	+ "schedule": "simple_2_rank_2_stage",
	447	+ "compute": {
	448	+ 0: ["0F0", "0F1", " ", "0B0", " ", "0B1"],
	449	+ 1: [" ", "1F0", "1B0", "1F1", "1B1", " "],
	450	+ },
	451	+ "comms": {
	452	+ 0: [
	453	+ "0F0",
	454	+ "0SEND_F0",
	455	+ "0F1",
	456	+ "0SEND_F1",
	457	+ "0RECV_B0",
	458	+ "0B0",
	459	+ "0RECV_B1",
	460	+ "0B1",
	461	+ ],
	462	+ 1: [
	463	+ "1RECV_F0",
	464	+ "1RECV_F1",
	465	+ "1F0",
	466	+ "1B0",
	467	+ "1SEND_B0",
	468	+ "1F1",
	469	+ "1B1",
	470	+ "1SEND_B1",
	471	+ ],
	472	+ },
	473	+ "stage_to_rank": lambda stage_idx: stage_idx,
	474	+ "num_stages": 2,
	475	+ "simulated_steps": 11,
	476	+ },
	477	+ {
	478	+ "schedule": "v_2_rank_4_stage",
	479	+ "compute": {
	480	+ 0: [
	481	+ "0F0",
	482	+ "0F1",
	483	+ " ",
	484	+ "3F0",
	485	+ "3B0",
	486	+ "3F1",
	487	+ "3B1",
	488	+ "0B0",
	489	+ "3W0",
	490	+ "0B1",
	491	+ "3W1",
	492	+ "0W0",
	493	+ "0W1",
	494	+ ],
	495	+ 1: [
	496	+ " ",
	497	+ "1F0",
	498	+ "2F0",
	499	+ "1F1",
	500	+ "2F1",
	501	+ "2B0",
	502	+ "1B0",
	503	+ "2B1",
	504	+ "1B1",
	505	+ "2W0",
	506	+ "2W1",
	507	+ "1W0",
	508	+ "1W1",
	509	+ ],
	510	+ },
	511	+ "comms": {
	512	+ 0: [
	513	+ "0F0",
	514	+ "0SEND_F0",
	515	+ "0F1",
	516	+ "0SEND_F1",
	517	+ "3RECV_F0",
	518	+ "3F0",
	519	+ "3B0",
	520	+ "3SEND_B0",
	521	+ "3RECV_F1",
	522	+ "3F1",
	523	+ "3B1",
	524	+ "3SEND_B1",
	525	+ "0RECV_B0",
	526	+ "0B0",
	527	+ "3W0",
	528	+ "0RECV_B1",
	529	+ "0B1",
	530	+ "3W1",
	531	+ "0W0",
	532	+ "0W1",
	533	+ ],
	534	+ 1: [
	535	+ "1RECV_F0",
	536	+ # interesting that this gets scheduled up front, is that expected?
	537	+ "1RECV_F1",
	538	+ "1F0",
	539	+ "2F0",
	540	+ "2SEND_F0",
	541	+ "1F1",
	542	+ # ditto
	543	+ "2RECV_B0",
	544	+ "2F1",
	545	+ "2SEND_F1",
	546	+ "2B0",
	547	+ # ditto
	548	+ "2RECV_B1",
	549	+ "1B0",
	550	+ "1SEND_B0",
	551	+ "2B1",
	552	+ "1B1",
	553	+ "1SEND_B1",
	554	+ "2W0",
	555	+ "2W1",
	556	+ "1W0",
	557	+ "1W1",
	558	+ ],
	559	+ },
	560	+ "stage_to_rank": lambda stage_idx: [0, 1, 1, 0][stage_idx],
	561	+ "num_stages": 4,
	562	+ "simulated_steps": 24,
	563	+ },
	564	+ ],
	565	+ )
	566	+ def test_send_recv(self, test_info):
	567	+ """Tests the lowering pass that adds send/recv ops to a compute-only schedule."""
	568	+ compute_sch = {
	569	+ rank: self._parse_actions(test_info["compute"][rank])
	570	+ for rank in test_info["compute"]
	571	+ }
	572	+ expected_comms_sch = {
	573	+ rank: self._parse_actions(test_info["comms"][rank])
	574	+ for rank in test_info["comms"]
	575	+ }
	576	+
	577	+ comms_sch = _add_send_recv(
	578	+ compute_sch, test_info["stage_to_rank"], test_info["num_stages"]
	579	+ )
	580	+ for rank in expected_comms_sch:
	581	+ for i, (expected, actual) in enumerate(
	582	+ zip(expected_comms_sch[rank], comms_sch[rank])
	583	+ ):
	584	+ self.assertEqual(
	585	+ expected,
	586	+ actual,
	587	+ (
	588	+ f"Mismatch on rank {rank} at position {i}."
	589	+ f"\nExpected: {expected_comms_sch[rank]}"
	590	+ f"\nActual: {comms_sch[rank]}"
	591	+ ),
	592	+ )
	593	+ self.assertEqual(len(comms_sch[rank]), len(expected_comms_sch[rank]))
	594	+
	595	+ simulated_schedule = _simulate_comms_compute(
	596	+ comms_sch,
	597	+ stage_to_rank=test_info["stage_to_rank"],
	598	+ num_stages=test_info["num_stages"],
	599	+ )
	600	+ num_steps = max([len(simulated_schedule[rank]) for rank in simulated_schedule])
	601	+ self.assertEqual(num_steps, test_info["simulated_steps"])
	602	+
	603	+ @parametrize("csv_name", ["zb1p_2rank_2stagep"])
	604	+ def test_csv(self, csv_name):
	605	+ def _dump_csv(pipeline_order_with_comms, filename: str):
	606	+ """Dump a CSV representation of the compute + comms schedule into a file with the provided filename."""
	607	+ with open(filename, "w", newline="") as csvfile:
	608	+ writer = csv.writer(csvfile)
	609	+ for rank in pipeline_order_with_comms:
	610	+ writer.writerow(pipeline_order_with_comms[rank])
	611	+
	612	+ compute_sch = {}
	613	+ with open(
	614	+ os.path.join(ARTIFACTS_DIR, f"{csv_name}_compute.csv"), newline=""
	615	+ ) as csvfile:
	616	+ for rank, row in enumerate(csv.reader(csvfile)):
	617	+ compute_sch[rank] = [_Action.from_str(s) for s in row]
	618	+ num_model_chunks = 2
	619	+ pipeline_parallel_size = 2
	620	+ num_stages = num_model_chunks * pipeline_parallel_size
	621	+
	622	+ for rank in compute_sch:
	623	+ compute_sch[rank] = _merge_bw(compute_sch[rank])
	624	+
	625	+ comms_sch = _add_send_recv(
	626	+ compute_sch,
	627	+ stage_to_rank=lambda chunk_index: chunk_index % pipeline_parallel_size,
	628	+ num_stages=num_stages,
	629	+ )
	630	+
	631	+ comms_csv = os.path.join(ARTIFACTS_DIR, f"{csv_name}_comms.csv")
	632	+
	633	+ # Uncomment to regenerate reference output
	634	+
	635	+ sch_ref = {}
	636	+ with open(comms_csv, newline="") as ref:
	637	+ for rank, row in enumerate(csv.reader(ref)):
	638	+ sch_ref[rank] = [_Action.from_str(s) for s in row]
	639	+
	640	+ for rank in sch_ref:
	641	+ for timestep, (a, b) in enumerate(zip(comms_sch[rank], sch_ref[rank])):
	642	+ self.assertEqual(a, b, f"Mismatch at {timestep=}, {a=}, expected {b}")
	643	+
	644	+ simulated_schedule = _simulate_comms_compute(
	645	+ comms_sch,
	646	+ stage_to_rank=lambda s: s % pipeline_parallel_size,
	647	+ num_stages=num_stages,
	648	+ )
	649	+
	650	+ num_steps = max([len(simulated_schedule[rank]) for rank in simulated_schedule])
	651	+ self.assertEqual(num_steps, 113)
	652	+
	653	+ def test_grad_with_v_schedule(self):
	654	+ """
	655	+ We have a special case for V schedules where 2 adjacent stages are on the same rank.
	656	+ E.g.
	657	+ rank0: stage 0, stage3
	658	+ rank1: stage 1, stage 2,
	659	+
	660	+ The special case involves not using send/recv ops but directly passing tensors between colocated stages.
	661	+
	662	+ This test runs on a single rank and just tests the 'stage1, stage2' portion for both F and B, comparing
	663	+ gradients to a reference model with 2 layers.
	664	+ """
	665	+ store = FakeStore()
	666	+ torch.distributed.init_process_group(
	667	+ backend="fake", rank=0, world_size=1, store=store
	668	+ )
	669	+ d_hid = 512
	670	+ batch_size = 256
	671	+ n_stages = 2
	672	+ device = "npu"
	673	+ full_mod = MultiMLP(d_hid, n_layers=n_stages)
	674	+ full_mod.to(device)
	675	+
	676	+ ref_mod = copy.deepcopy(full_mod)
	677	+ x = torch.randn(batch_size, d_hid, device=device)
	678	+ with torch.no_grad():
	679	+ y = ref_mod(x)
	680	+ # Add a small perturbation
	681	+ target = y + torch.randn(batch_size, d_hid, device=device)
	682	+
	683	+ loss_fn = torch.nn.MSELoss(reduction="sum")
	684	+
	685	+ # Run reference
	686	+ for _ in range(2):
	687	+ ref_mod.zero_grad()
	688	+ ref_out = ref_mod(x)
	689	+ ref_loss = loss_fn(ref_out, target)
	690	+ ref_loss.backward()
	691	+
	692	+ stage_indices = [0, 1]
	693	+ submod_names = [f"layers.{i}" for i in stage_indices]
	694	+ stage_modules = [
	695	+ full_mod.get_submodule(submod_name)
	696	+ for submod_name in submod_names
	697	+ ]
	698	+ # Create a pipeline stage to wrap that submodule
	699	+ num_microbatches = 2
	700	+ stages = [
	701	+ PipelineStage(
	702	+ stage_module,
	703	+ stage_idx,
	704	+ n_stages,
	705	+ device,
	706	+ )
	707	+ for stage_module, stage_idx in zip(stage_modules, stage_indices)
	708	+ ]
	709	+
	710	+ # Attach to a schedule
	711	+ schedule = _PipelineScheduleRuntime(
	712	+ stages,
	713	+ num_microbatches,
	714	+ loss_fn=loss_fn,
	715	+ scale_grads=False,
	716	+ )
	717	+ schedule._load_actions(
	718	+ {
	719	+ 0: self._parse_actions(
	720	+ [
	721	+ "0F0",
	722	+ "0F1",
	723	+ "1F0",
	724	+ "1F1",
	725	+ "1B0",
	726	+ "1B1",
	727	+ "0B0",
	728	+ "0B1",
	729	+ ]
	730	+ ),
	731	+ },
	732	+ format="compute_comms",
	733	+ )
	734	+
	735	+ # Run
	736	+ with check_leaked_tensors() as garbage_tensors:
	737	+ for _ in range(2):
	738	+ # Zero gradients
	739	+ for stage_module in stage_modules:
	740	+ stage_module.zero_grad()
	741	+ losses = []
	742	+ out = schedule.step(x, target=target, losses=losses)
	743	+ self.assertEqual(
	744	+ len(garbage_tensors),
	745	+ 0,
	746	+ "Found leaked tensors, check logs above for debug info",
	747	+ )
	748	+
	749	+ # Check output
	750	+ torch.testing.assert_close(out, ref_out)
	751	+ # Check loss
	752	+ # Since the reduction used in the loss function above is "sum", we use
	753	+ # "sum" here to reduce microbatch losses into a single value too.
	754	+ pipe_loss = sum(losses)
	755	+ torch.testing.assert_close(pipe_loss, ref_loss)
	756	+
	757	+ # Check gradients
	758	+ for stage_module, submod_name in zip(stage_modules, submod_names):
	759	+ # Get corresponding submodule from reference model
	760	+ ref_submod = ref_mod.get_submodule(submod_name)
	761	+ # Check gradients per parameter
	762	+ for name, p in stage_module.named_parameters():
	763	+ ref_p = ref_submod.get_parameter(name)
	764	+ try:
	765	+ torch.testing.assert_close(p.grad, ref_p.grad, rtol=1e-5, atol=4e-5)
	766	+ except AssertionError:
	767	+ print(f"Gradient test failed for {name}: {p.grad} vs {ref_p.grad}")
	768	+ raise
	769	+
	770	+ torch.distributed.destroy_process_group()
	771	+
	772	+ def test_grad_with_split_b_w(self):
	773	+ """
	774	+ Ensure that separate dInput and dWeight computations are correctly executed.
	775	+ This test runs on a single rank and just tests a single stage with 2 microbatches with separate B, W operations.
	776	+ """
	777	+ store = FakeStore()
	778	+ torch.distributed.init_process_group(
	779	+ backend="fake", rank=0, world_size=1, store=store
	780	+ )
	781	+ d_hid = 512
	782	+ batch_size = 256
	783	+ n_stages = 1
	784	+ device = "npu"
	785	+ full_mod = MultiMLP(d_hid, n_layers=n_stages)
	786	+ full_mod.to(device)
	787	+
	788	+ ref_mod = copy.deepcopy(full_mod)
	789	+ x = torch.randn(batch_size, d_hid, device=device)
	790	+ with torch.no_grad():
	791	+ y = ref_mod(x)
	792	+ # Add a small perturbation
	793	+ target = y + torch.randn(batch_size, d_hid, device=device)
	794	+
	795	+ loss_fn = torch.nn.MSELoss(reduction="sum")
	796	+
	797	+ # Run reference
	798	+ for _ in range(2):
	799	+ ref_mod.zero_grad()
	800	+ ref_out = ref_mod(x)
	801	+ ref_loss = loss_fn(ref_out, target)
	802	+ ref_loss.backward()
	803	+
	804	+ stage_indices = [0]
	805	+ submod_names = [f"layers.{i}" for i in stage_indices]
	806	+ stage_modules = [
	807	+ full_mod.get_submodule(submod_name)
	808	+ for submod_name in submod_names
	809	+ ]
	810	+ # Create a pipeline stage to wrap that submodule
	811	+ num_microbatches = 2
	812	+ stages = [
	813	+ PipelineStage(
	814	+ stage_module,
	815	+ stage_idx,
	816	+ n_stages,
	817	+ device,
	818	+ )
	819	+ for stage_module, stage_idx in zip(stage_modules, stage_indices)
	820	+ ]
	821	+
	822	+ # Attach to a schedule
	823	+ schedule = _PipelineScheduleRuntime(
	824	+ stages,
	825	+ num_microbatches,
	826	+ loss_fn=loss_fn,
	827	+ )
	828	+ schedule._load_actions(
	829	+ {
	830	+ 0: self._parse_actions(
	831	+ [
	832	+ "0F0",
	833	+ "0F1",
	834	+ "0I0",
	835	+ "0I1",
	836	+ "0W0",
	837	+ "0W1",
	838	+ ]
	839	+ ),
	840	+ },
	841	+ format="compute_comms",
	842	+ )
	843	+
	844	+ # Run
	845	+ with check_leaked_tensors() as garbage_tensors:
	846	+ for _ in range(2):
	847	+ # Zero gradients
	848	+ for stage_module in stage_modules:
	849	+ stage_module.zero_grad()
	850	+ losses = []
	851	+ out = schedule.step(x, target=target, losses=losses)
	852	+ self.assertEqual(
	853	+ len(garbage_tensors),
	854	+ 0,
	855	+ "Found leaked tensors, check logs above for debug info",
	856	+ )
	857	+
	858	+ # Check output
	859	+ torch.testing.assert_close(out, ref_out)
	860	+ # Check loss
	861	+ # Since the reduction used in the loss function above is "sum", we use
	862	+ # "sum" here to reduce microbatch losses into a single value too.
	863	+ pipe_loss = sum(losses)
	864	+ torch.testing.assert_close(pipe_loss, ref_loss)
	865	+
	866	+ # Check gradients
	867	+ for stage_module, submod_name in zip(stage_modules, submod_names):
	868	+ # Get corresponding submodule from reference model
	869	+ ref_submod = ref_mod.get_submodule(submod_name)
	870	+ # Check gradients per parameter
	871	+ for name, p in stage_module.named_parameters():
	872	+ ref_p = ref_submod.get_parameter(name)
	873	+ try:
	874	+ torch.testing.assert_close(p.grad, ref_p.grad, rtol=1e-5, atol=4e-5)
	875	+ except AssertionError:
	876	+ print(f"Gradient test failed for {name}: {p.grad} vs {ref_p.grad}")
	877	+ raise
	878	+
	879	+ torch.distributed.destroy_process_group()
	880	+
	881	+
	882	+class TestValidateSchedule(TestCase):
	883	+ def test_valid_schedule(self):
	884	+ schedule_actions = [
	885	+ {
	886	+ 0: [_Action(0, F, 0), _Action(0, B, 0)],
	887	+ 1: [_Action(1, F, 0), _Action(1, B, 0)],
	888	+ },
	889	+ {
	890	+ 0: [_Action(0, F, 0), _Action(0, I, 0), _Action(0, W, 0)],
	891	+ 1: [_Action(1, F, 0), _Action(1, I, 0), _Action(1, W, 0)],
	892	+ },
	893	+ ]
	894	+ pp_group_size = 2
	895	+ num_stages = 2
	896	+ num_microbatches = 1
	897	+ for actions in schedule_actions:
	898	+ _validate_schedule(actions, pp_group_size, num_stages, num_microbatches)
	899	+
	900	+ def test_invalid_schedule_missing_rank(self):
	901	+ actions = {
	902	+ 0: [_Action(0, F, 0), _Action(0, B, 0)],
	903	+ }
	904	+ pp_group_size = 2
	905	+ num_stages = 2
	906	+ num_microbatches = 1
	907	+ with self.assertRaises(AssertionError):
	908	+ _validate_schedule(actions, pp_group_size, num_stages, num_microbatches)
	909	+
	910	+ def test_invalid_schedule_missing_action(self):
	911	+ actions = {
	912	+ 0: [_Action(0, F, 0)],
	913	+ 1: [_Action(1, F, 0)],
	914	+ }
	915	+ pp_group_size = 2
	916	+ num_stages = 2
	917	+ num_microbatches = 1
	918	+ with self.assertRaises(AssertionError):
	919	+ _validate_schedule(actions, pp_group_size, num_stages, num_microbatches)
	920	+
	921	+
	922	+class ScheduleUtilTests(TestCase):
	923	+ def test_generate_stage_to_rank_mapping(self):
	924	+ stage_to_rank = generate_stage_to_rank_mapping(2, 2)
	925	+ self.assertEqual(
	926	+ stage_to_rank,
	927	+ {
	928	+ 0: 0,
	929	+ 1: 1,
	930	+ },
	931	+ )
	932	+ stage_to_rank = generate_stage_to_rank_mapping(2, 4)
	933	+ self.assertEqual(stage_to_rank, {0: 0, 1: 1, 2: 0, 3: 1})
	934	+ stage_to_rank = generate_stage_to_rank_mapping(4, 8)
	935	+ self.assertEqual(
	936	+ stage_to_rank, {0: 0, 1: 1, 2: 2, 3: 3, 4: 0, 5: 1, 6: 2, 7: 3}
	937	+ )
	938	+ stage_to_rank = generate_stage_to_rank_mapping(2, 4, style="v")
	939	+ self.assertEqual(
	940	+ stage_to_rank,
	941	+ {
	942	+ 0: 0,
	943	+ 1: 1,
	944	+ 2: 1,
	945	+ 3: 0,
	946	+ },
	947	+ )
	948	+ stage_to_rank = generate_stage_to_rank_mapping(4, 12, style="v")
	949	+ self.assertEqual(
	950	+ stage_to_rank,
	951	+ {
	952	+ 0: 0,
	953	+ 1: 1,
	954	+ 2: 2,
	955	+ 3: 3,
	956	+ 4: 3,
	957	+ 5: 2,
	958	+ 6: 1,
	959	+ 7: 0,
	960	+ 8: 0,
	961	+ 9: 1,
	962	+ 10: 2,
	963	+ 11: 3,
	964	+ },
	965	+ )
	966	+ stage_to_rank = generate_stage_to_rank_mapping(4, 16, style="v")
	967	+ self.assertEqual(
	968	+ stage_to_rank,
	969	+ {
	970	+ 0: 0,
	971	+ 1: 1,
	972	+ 2: 2,
	973	+ 3: 3,
	974	+ 4: 3,
	975	+ 5: 2,
	976	+ 6: 1,
	977	+ 7: 0,
	978	+ 8: 0,
	979	+ 9: 1,
	980	+ 10: 2,
	981	+ 11: 3,
	982	+ 12: 3,
	983	+ 13: 2,
	984	+ 14: 1,
	985	+ 15: 0,
	986	+ },
	987	+ )
	988	+
	989	+
	990	+instantiate_parametrized_tests(TestScheduleLowering)
	991	+
	992	+if __name__ == "__main__":
	993	+ run_tests()

Atest/distributed/pipelining/test_schedule_multiproc.py+958-0

@@ -0,0 +1,958 @@
	1	+# Copyright (c) Meta Platforms, Inc. and affiliates
	2	+# Owner(s): ["oncall: distributed"]
	3	+import copy
	4	+import logging
	5	+import os
	6	+import sys
	7	+import tempfile
	8	+
	9	+from model_registry import ModelWithKwargs, MultiMLP, MultiMLPWithDw
	10	+from schedule_registry import (
	11	+ ScheduleUnbalanced,
	12	+ ScheduleVShaped,
	13	+ ScheduleWithReorderedB,
	14	+ ScheduleWithW,
	15	+)
	16	+
	17	+import torch
	18	+import torch.distributed as dist
	19	+from torch.distributed.pipelining import (
	20	+ _ScheduleForwardOnly,
	21	+ pipeline,
	22	+ PipelineStage,
	23	+ Schedule1F1B,
	24	+ ScheduleGPipe,
	25	+ ScheduleInterleaved1F1B,
	26	+ ScheduleInterleavedZeroBubble,
	27	+ ScheduleLoopedBFS,
	28	+ ScheduleZBVZeroBubble,
	29	+)
	30	+from torch.distributed.pipelining.schedules import _PipelineScheduleRuntime
	31	+from torch.testing._internal.common_cuda import TEST_MULTIGPU
	32	+from torch.testing._internal.common_distributed import (
	33	+ MultiProcContinousTest,
	34	+ requires_nccl,
	35	+)
	36	+from torch.testing._internal.common_utils import (
	37	+ check_leaked_tensors,
	38	+ instantiate_parametrized_tests,
	39	+ parametrize,
	40	+ skip_but_pass_in_sandcastle_if,
	41	+)
	42	+
	43	+
	44	+logger = logging.getLogger(__name__)
	45	+
	46	+d_hid = 512
	47	+batch_size = 256
	48	+
	49	+torch.manual_seed(0)
	50	+
	51	+
	52	+class ScheduleTest(MultiProcContinousTest):
	53	+ @classmethod
	54	+ def backend_str(cls) -> str:
	55	+ # Testing with HCCL backend
	56	+ return "hccl"
	57	+
	58	+ @classmethod
	59	+ def setUpClass(cls):
	60	+ """
	61	+ Class-scope test fixture. Run once for entire test class, before any test starts.
	62	+ Set up the device.
	63	+ """
	64	+ super().setUpClass()
	65	+ dev_id = cls.rank % torch.npu.device_count()
	66	+ cls.device = torch.device(f"npu:{dev_id}")
	67	+
	68	+ @parametrize("ScheduleClass", [_ScheduleForwardOnly])
	69	+ def test_forward_only(self, ScheduleClass):
	70	+ mod = MultiMLP(d_hid, n_layers=self.world_size)
	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()
	77	+
	78	+ num_microbatches = 4
	79	+ x_mb = x.chunk(num_microbatches)[0]
	80	+
	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	+ )
	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)
	96	+
	97	+ # Run
	98	+ num_iters = 20
	99	+ for _ in range(num_iters):
	100	+ if self.rank == 0:
	101	+ schedule.step(x)
	102	+ dist.recv(x, src=self.world_size - 1)
	103	+ elif self.rank == self.world_size - 1:
	104	+ out = schedule.step()
	105	+ dist.send(out, dst=0)
	106	+ else:
	107	+ schedule.step()
	108	+
	109	+ # Validate pipelined output is the same as reference model
	110	+ if self.rank == self.world_size - 1:
	111	+ for _ in range(num_iters):
	112	+ x_clone = mod_ref(x_clone)
	113	+
	114	+ torch.testing.assert_close(x_clone, out)
	115	+
	116	+ @parametrize("ScheduleClass", [ScheduleGPipe, Schedule1F1B])
	117	+ def test_multi_iter(self, ScheduleClass):
	118	+ mod = MultiMLP(d_hid, n_layers=self.world_size)
	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 = 4
	126	+ x_mb = x.chunk(chunks)[0]
	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)
	143	+
	144	+ # Run
	145	+ for _ in range(20):
	146	+ if self.rank == 0:
	147	+ schedule.step(x)
	148	+ elif self.rank == self.world_size - 1:
	149	+ losses = []
	150	+ schedule.step(target=target, losses=losses)
	151	+ else:
	152	+ schedule.step()
	153	+
	154	+ @parametrize("ScheduleClass", [ScheduleGPipe, Schedule1F1B])
	155	+ def test_kwargs_with_tracer(self, ScheduleClass):
	156	+ mod = ModelWithKwargs(d_hid)
	157	+ mod.to(self.device)
	158	+
	159	+ x = torch.randn(batch_size, d_hid, device=self.device)
	160	+ y = torch.randn(batch_size, d_hid, device=self.device)
	161	+ target = torch.randn(batch_size, d_hid, device=self.device)
	162	+ loss_fn = torch.nn.MSELoss(reduction="sum")
	163	+
	164	+ chunks = 4
	165	+ x_mb = x.chunk(chunks)[0]
	166	+ y_mb = y.chunk(chunks)[0]
	167	+
	168	+ pipe = pipeline(
	169	+ mod,
	170	+ mb_args=(x_mb,),
	171	+ mb_kwargs={"y": y_mb},
	172	+ )
	173	+
	174	+ stage = pipe.build_stage(
	175	+ self.rank,
	176	+ self.device,
	177	+ )
	178	+
	179	+ # Attach to a schedule
	180	+ schedule = ScheduleClass(stage, chunks, loss_fn=loss_fn, scale_grads=False)
	181	+
	182	+ # Run
	183	+ if self.rank == 0:
	184	+ schedule.step(x, y=y)
	185	+ elif self.rank == self.world_size - 1:
	186	+ losses = []
	187	+ out = schedule.step(target=target, losses=losses)
	188	+ else:
	189	+ schedule.step()
	190	+
	191	+ dist.barrier()
	192	+
	193	+ # Last rank checks result
	194	+ if self.rank == self.world_size - 1:
	195	+ ref_out = mod(x, y=y)
	196	+ ref_loss = loss_fn(ref_out, target)
	197	+ pipe_loss = sum(losses)
	198	+ torch.testing.assert_close(out, ref_out, rtol=1e-2, atol=5e-3)
	199	+ torch.testing.assert_close(pipe_loss, ref_loss)
	200	+
	201	+ @parametrize("ScheduleClass", [ScheduleGPipe, Schedule1F1B])
	202	+ @parametrize("ModelClass", [MultiMLP])
	203	+ def test_grad_with_tracer(self, ScheduleClass, ModelClass):
	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	+
	216	+ # Run reference
	217	+ for _ in range(2):
	218	+ ref_mod.zero_grad()
	219	+ ref_out = ref_mod(x)
	220	+ ref_loss = loss_fn(ref_out, target)
	221	+ ref_loss.backward()
	222	+
	223	+ # Create a pipeline
	224	+ chunks = 4
	225	+ x_mb = x.chunk(chunks)[0]
	226	+ split_spec = mod.split_spec if hasattr(mod, "split_spec") else None
	227	+ pipe = pipeline(
	228	+ mod,
	229	+ mb_args=(x_mb,),
	230	+ split_spec=split_spec,
	231	+ )
	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)
	240	+
	241	+ # Run
	242	+ stage_module = pipe.get_stage_module(self.rank)
	243	+ for _ in range(2):
	244	+ # Zero gradients
	245	+ stage_module.zero_grad()
	246	+ if self.rank == 0:
	247	+ schedule.step(x)
	248	+ elif self.rank == self.world_size - 1:
	249	+ losses = []
	250	+ out = schedule.step(target=target, losses=losses)
	251	+ else:
	252	+ schedule.step()
	253	+
	254	+ dist.barrier()
	255	+
	256	+ # Last rank checks result
	257	+ if self.rank == self.world_size - 1:
	258	+ # Check output
	259	+ 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)
	264	+ torch.testing.assert_close(pipe_loss, ref_loss)
	265	+
	266	+ # Every rank checks gradients
	267	+ for name, p in stage_module.named_parameters():
	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	+
	275	+ @parametrize("ScheduleClass", [ScheduleGPipe, Schedule1F1B])
	276	+ @parametrize("shape_inference", [True, False])
	277	+ def test_grad_with_manual(self, ScheduleClass, shape_inference):
	278	+ full_mod = MultiMLP(d_hid, n_layers=self.world_size)
	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	+
	290	+ # Run reference
	291	+ for _ in range(2):
	292	+ ref_mod.zero_grad()
	293	+ ref_out = ref_mod(x)
	294	+ ref_loss = loss_fn(ref_out, target)
	295	+ ref_loss.backward()
	296	+
	297	+ # Get a submodule, e.g. `layers.0` or `layers.1`
	298	+ submod_name = f"layers.{self.rank}"
	299	+ stage_module = full_mod.get_submodule(submod_name)
	300	+ chunks = 4
	301	+
	302	+ if shape_inference:
	303	+ input_args = None
	304	+ output_args = None
	305	+ else:
	306	+ input_args = (x.chunk(chunks)[0],)
	307	+ with torch.no_grad():
	308	+ output_args = stage_module(*input_args)
	309	+
	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)
	322	+
	323	+ # Run
	324	+ for _ in range(2):
	325	+ # Zero gradients
	326	+ stage_module.zero_grad()
	327	+ if self.rank == 0:
	328	+ schedule.step(x)
	329	+ elif self.rank == self.world_size - 1:
	330	+ losses = []
	331	+ out = schedule.step(target=target, losses=losses)
	332	+ else:
	333	+ schedule.step()
	334	+
	335	+ dist.barrier()
	336	+
	337	+ # Last rank checks result
	338	+ if self.rank == self.world_size - 1:
	339	+ # Check output
	340	+ 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)
	345	+ torch.testing.assert_close(pipe_loss, ref_loss)
	346	+
	347	+ # Every rank checks gradients
	348	+ ref_submod = ref_mod.get_submodule(submod_name)
	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	+
	357	+ @parametrize(
	358	+ "ScheduleClass",
	359	+ [
	360	+ ScheduleInterleaved1F1B,
	361	+ ScheduleLoopedBFS,
	362	+ ScheduleInterleavedZeroBubble,
	363	+ ],
	364	+ )
	365	+ @parametrize("use_new_runtime", [False, True])
	366	+ def test_grad_with_manual_interleaved(self, ScheduleClass, use_new_runtime):
	367	+ stages_per_rank = 2
	368	+ n_stages = stages_per_rank * self.world_size
	369	+ full_mod = MultiMLP(d_hid, n_layers=n_stages)
	370	+ full_mod.to(self.device)
	371	+
	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	+
	381	+ # Run reference
	382	+ for _ in range(2):
	383	+ ref_mod.zero_grad()
	384	+ ref_out = ref_mod(x)
	385	+ ref_loss = loss_fn(ref_out, target)
	386	+ ref_loss.backward()
	387	+
	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 = (
	401	+ ScheduleClass.num_microbatches
	402	+ if hasattr(ScheduleClass, "num_microbatches")
	403	+ else 8
	404	+ )
	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	+
	415	+ # Attach to a schedule
	416	+ schedule = ScheduleClass(
	417	+ stages, num_microbatches, loss_fn=loss_fn, scale_grads=False
	418	+ )
	419	+ if use_new_runtime:
	420	+ old_schedule = schedule
	421	+ tmp_schedule = _PipelineScheduleRuntime(
	422	+ stages,
	423	+ num_microbatches,
	424	+ loss_fn=loss_fn,
	425	+ scale_grads=False,
	426	+ )
	427	+ tmp_schedule._load_actions(old_schedule.pipeline_order)
	428	+ # test that csv round-trip works for compute_comms schedule
	429	+ schedule = _PipelineScheduleRuntime(
	430	+ stages,
	431	+ num_microbatches,
	432	+ loss_fn=loss_fn,
	433	+ scale_grads=False,
	434	+ )
	435	+ with tempfile.NamedTemporaryFile() as f:
	436	+ tmp_schedule._dump_csv(f.name)
	437	+ f.seek(0)
	438	+ schedule._load_csv(f.name, format="compute_comms")
	439	+ one_more_schedule = _PipelineScheduleRuntime(
	440	+ stages,
	441	+ num_microbatches,
	442	+ loss_fn=loss_fn,
	443	+ scale_grads=False,
	444	+ )
	445	+ one_more_schedule._load_actions(
	446	+ schedule.pipeline_order_with_comms, format="compute_comms"
	447	+ )
	448	+ self.assertEqual(
	449	+ len(schedule.pipeline_order_with_comms),
	450	+ len(
	451	+ one_more_schedule.pipeline_order_with_comms,
	452	+ ),
	453	+ )
	454	+ for rank in schedule.pipeline_order_with_comms:
	455	+ self.assertEqual(
	456	+ len(schedule.pipeline_order_with_comms[rank]),
	457	+ len(
	458	+ one_more_schedule.pipeline_order_with_comms[rank],
	459	+ ),
	460	+ )
	461	+ for a, b in zip(
	462	+ schedule.pipeline_order_with_comms[rank],
	463	+ one_more_schedule.pipeline_order_with_comms[rank],
	464	+ ):
	465	+ self.assertEqual(a, b)
	466	+
	467	+ # Run
	468	+ with check_leaked_tensors() as garbage_tensors:
	469	+ for _ in range(2):
	470	+ # Zero gradients
	471	+ for stage_module in stage_modules:
	472	+ stage_module.zero_grad()
	473	+ if self.rank == 0:
	474	+ schedule.step(x)
	475	+ elif self.rank == self.world_size - 1:
	476	+ losses = []
	477	+ out = schedule.step(target=target, losses=losses)
	478	+ else:
	479	+ schedule.step()
	480	+ self.assertEqual(
	481	+ len(garbage_tensors),
	482	+ 0,
	483	+ "Found leaked tensors, check logs above for debug info",
	484	+ )
	485	+ dist.barrier()
	486	+
	487	+ # Last rank checks result
	488	+ if self.rank == self.world_size - 1:
	489	+ # Check output
	490	+ 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)
	495	+ torch.testing.assert_close(pipe_loss, ref_loss)
	496	+
	497	+ # Every rank checks gradients
	498	+ for stage_module, submod_name in zip(stage_modules, submod_names):
	499	+ # Get corresponding submodule from reference model
	500	+ ref_submod = ref_mod.get_submodule(submod_name)
	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	+ )
	562	+
	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])
	811	+ def test_schedule_with_weight_update_mlp_e2e(self, ScheduleClass):
	812	+ stages_per_rank = 2
	813	+ n_stages = stages_per_rank * self.world_size
	814	+ full_mod = MultiMLPWithDw(d_hid, n_layers=n_stages)
	815	+ full_mod.to(self.device)
	816	+
	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()
	828	+
	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	+
	840	+ # Run reference
	841	+ for _ in range(2):
	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	+
	849	+ ref_mod.zero_grad()
	850	+ ref_out = ref_mod(x)
	851	+ ref_loss = ref_loss_fn(ref_out, target)
	852	+ ref_loss.backward()
	853	+
	854	+ class CustomState:
	855	+ def __init__(self, stage_module, stage_idx, rank):
	856	+ self.i = 0
	857	+ self.stage_module = stage_module
	858	+ self.stage_idx = stage_idx
	859	+ self.rank = rank
	860	+
	861	+ def dw_builder(self):
	862	+ def dw_runner():
	863	+ # This inner function would be called by PipelineStage during `backward_weight_one_chunk`
	864	+ self.i += 1
	865	+ print(
	866	+ f"[Rank {self.rank}] dw_count={self.i} stage={self.stage_idx}"
	867	+ )
	868	+ self.stage_module.compute_dW()
	869	+
	870	+ return dw_runner
	871	+
	872	+ cs = {}
	873	+ for stage_module, stage_idx in zip(stage_modules, stage_indices):
	874	+ cs[stage_idx] = CustomState(stage_module, stage_idx, self.rank)
	875	+
	876	+ # Create a pipeline stage to wrap that submodule
	877	+ chunks = 2
	878	+ stages = [
	879	+ PipelineStage(
	880	+ stage_module,
	881	+ stage_idx,
	882	+ n_stages,
	883	+ self.device,
	884	+ dw_builder=cs[stage_idx].dw_builder,
	885	+ )
	886	+ for stage_module, stage_idx in zip(stage_modules, stage_indices)
	887	+ ]
	888	+
	889	+ # Attach to a schedule
	890	+ schedule = ScheduleClass(
	891	+ stages, chunks, loss_fn=full_loss_fn, scale_grads=False
	892	+ )
	893	+
	894	+ for _ in range(2):
	895	+ # Zero gradients
	896	+ for stage_module in stage_modules:
	897	+ stage_module.zero_grad()
	898	+ if self.rank == 0:
	899	+ schedule.step(x)
	900	+ elif self.rank == self.world_size - 1:
	901	+ losses = []
	902	+ out = schedule.step(target=target, losses=losses)
	903	+ else:
	904	+ schedule.step()
	905	+
	906	+ 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	+
	912	+ # Check loss
	913	+ # Since the reduction used in the loss function above is "sum", we use
	914	+ # "sum" here to reduce microbatch losses into a single value too.
	915	+ pipe_loss = sum(losses)
	916	+ torch.testing.assert_close(pipe_loss, ref_loss)
	917	+
	918	+ # Every rank checks gradients
	919	+ for stage_module, submod_name in zip(stage_modules, submod_names):
	920	+ # Get corresponding submodule from reference model
	921	+ ref_submod = ref_mod.get_submodule(submod_name)
	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	+
	927	+
	928	+instantiate_parametrized_tests(ScheduleTest)
	929	+
	930	+
	931	+if __name__ == "__main__":
	932	+ # Check if NPU and HCCL are available
	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	+ )

Atest/distributed/pipelining/test_stage.py+342-0

Atest/distributed/pipelining/test_transformer.py+75-0

Atest/distributed/pipelining/test_unflatten.py+75-0