Coverage for / home / jenkins / .local / lib / python3.10 / site-packages / hyper_parallel / compile / examples / automodel_tp / train.py: 0%
135 statements
« prev ^ index » next coverage.py v7.13.1, created at 2026-10-25 05:45 +0800
1#!/usr/bin/env python3
2# Copyright 2026 Huawei Technologies Co., Ltd
3#
4# Licensed under the Apache License, Version 2.0 (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# http://www.apache.org/licenses/LICENSE-2.0
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.
15# ============================================================================
16"""TP + FSDP graph-mode demo.
18Combines automodel's TP sharding with the graph-mode FSDP pass:
20 1. automodel ``MeshContext.build_meshs`` builds a ``(dp, cp, tp)`` mesh.
21 The TP group is created here and the boundary forwards will hold the
22 group object directly.
23 2. ``ShardingPlanner.plan`` derives a TP ``ShardingPlan``.
24 3. ``apply_sharding_plan(production)`` turns parameters into plain TP
25 shards (``to_local``) and wraps module forwards with boundary
26 redistribution — TP collectives are inserted on **activations**
27 (colwise all-reduce on exit, rowwise input gather, etc.).
28 4. FSDP2Manager is NOT used (no ``fully_shard`` hooks). Instead the
29 TP-sharded model is handed to ``GraphTrainer``:
30 - the automodel mesh is reused — only the dp sub-mesh is registered
31 under ``"fsdp"`` (the TP group is already wired into the boundary
32 forwards);
33 - ``trace_model_graph`` traces a joint fwd+bwd FX graph; parameters
34 are static inputs whose shape is the TP shard;
35 - ``FSDPPass`` shards each already-TP-sharded parameter along dim 0
36 (FSDP axis, orthogonal to the TP axis), inserts all_gather on the
37 parameter placeholders (recovers the TP shard) and reduce_scatter
38 on the gradient outputs;
39 - ``AutoOverlapPass`` moves wait_tensor for comm/compute overlap.
41The two communication domains are orthogonal: FSDP talks on parameters/grads
42(the "fsdp" group), TP talks on activations (the "tp" group baked into the
43boundary forwards). FSDPPass only sees the FX graph — it is TP-agnostic.
44"""
46import argparse
47import logging
48import sys
49from pathlib import Path
51import torch
52import torch.distributed as dist
53import torch.nn.functional as F
54import yaml
56_REPO_ROOT = str(Path(__file__).resolve().parent.parent.parent.parent.parent)
57if _REPO_ROOT not in sys.path:
58 sys.path.append(_REPO_ROOT)
60from model import ( # pylint: disable=C0413
61 build_model as build_model_from_config,
62 DataSampler,
63)
64from hyper_parallel.compile import ( # pylint: disable=C0413
65 GraphParallelPlan,
66 GraphTrainer,
67 PassConfig,
68)
69from hyper_parallel.distributed.mesh import MeshContext # pylint: disable=C0413
70from hyper_parallel.distributed import ( # pylint: disable=C0413
71 ShardingPlanner,
72 apply_sharding_plan,
73)
75_LOG = logging.getLogger(__name__)
78def parse_args():
79 """Parse example CLI arguments."""
80 parser = argparse.ArgumentParser(
81 description="TP + FSDP graph-mode demo (automodel TP-shard + FSDPPass)"
82 )
83 parser.add_argument("--config", type=str, default="config.yaml")
84 return parser.parse_args()
87def load_config(path: str) -> dict:
88 """Load the example YAML config."""
89 with open(path, encoding="utf-8") as f:
90 return yaml.safe_load(f)
93def setup_distributed() -> str:
94 """Init the process group; return the device type to build meshes on."""
95 if not dist.is_initialized():
96 if hasattr(torch, "npu") and torch.npu.is_available():
97 dist.init_process_group(backend="hccl")
98 torch.npu.set_device(dist.get_rank() % torch.npu.device_count())
99 if dist.get_rank() == 0:
100 _LOG.info(
101 "[setup] backend=hccl, world_size=%s, npu_device=%s",
102 dist.get_world_size(),
103 torch.npu.current_device(),
104 )
105 return "npu"
106 dist.init_process_group(backend="gloo")
107 if dist.get_rank() == 0:
108 _LOG.info("[setup] backend=gloo, world_size=%s", dist.get_world_size())
109 return "cpu"
110 return "cpu"
113def build_model(cfg: dict, device: torch.device) -> torch.nn.Module:
114 """Build the model on ``device``.
116 Delegates to ``build_model`` from ``model.py``, which creates an
117 ``AutoModelAdapterForCausalLM`` — a structurally correct model where
118 ``rotary_emb`` lives inside ``Attention.forward`` (not in the inner
119 model's forward). This means cos/sin are computed *after* the SP
120 boundary all-gather, so they naturally have the full sequence length
121 — no runtime patch needed.
122 """
123 return build_model_from_config(cfg["model"], device)
126def build_mesh_context(parallel_cfg: dict, device_type: str) -> MeshContext:
127 """Build the automodel (dp, cp, tp) mesh + the dense FSDP sub-mesh."""
128 tp = parallel_cfg["tp_size"]
129 dp = parallel_cfg["dp_size"]
130 sequence_parallel = parallel_cfg.get("sequence_parallel", False)
131 loss_parallel = parallel_cfg.get("loss_parallel", False)
132 world_size = dist.get_world_size()
133 if tp * dp != world_size:
134 raise ValueError(
135 f"tp_size({tp}) * dp_size({dp}) != world_size({world_size}) — this "
136 "demo assumes tp*dp == world_size (cp=pp=ep=1)."
137 )
138 ctx = MeshContext(
139 tp_size=tp,
140 dp_size=dp,
141 cp_size=1,
142 pp_size=1,
143 ep_size=1,
144 dp_replicate_size=1,
145 dp_shard_size=dp,
146 edp_shard_size=1,
147 sequence_parallel=sequence_parallel,
148 loss_parallel=loss_parallel,
149 )
150 ctx.build_meshs(device_type, world_size)
151 ctx.tp_rank = dist.get_rank() % tp
152 ctx.dp_rank = dist.get_rank() // tp
153 return ctx
156def train_fn(model, *, input_ids, labels):
157 """Standard CE loss on the boundary-wrapped model forward.
159 In SP mode the lm_head boundary all-gathers hidden_states to full
160 sequence, so logits and labels both have the full seq_len. The
161 min_seq alignment is a safety no-op.
162 """
163 outputs = model(input_ids)
164 if hasattr(outputs, "logits"):
165 logits = outputs.logits
166 elif isinstance(outputs, dict) and "logits" in outputs:
167 logits = outputs["logits"]
168 else:
169 logits = outputs[0]
170 # Align sequence dimension: logits may be shorter than labels in SP mode
171 min_seq = min(logits.shape[1], labels.shape[1])
172 logits = logits[:, :min_seq, :]
173 labels = labels[:, :min_seq]
174 shift_logits = logits[..., :-1, :].contiguous()
175 shift_labels = labels[..., 1:].contiguous()
176 loss = F.cross_entropy(
177 shift_logits.view(-1, shift_logits.size(-1)),
178 shift_labels.view(-1),
179 ignore_index=-100,
180 )
181 return loss
184def inspect_graph(trainer: GraphTrainer) -> None:
185 """Print FSDP and TP comm nodes in the compiled joint graph.
187 FSDPPass tags its nodes with ``meta['comm_type']`` starting with ``fsdp_``.
188 TP collectives live inside the boundary forwards and are plain
189 ``call_function`` nodes whose op name mentions all_gather/all_reduce/
190 reduce_scatter.
191 """
192 gm = trainer._compiler._joint_graph.graph_module # pylint: disable=W0212
193 fsdp, tp_nodes = [], []
194 for node in gm.graph.nodes:
195 if node.op != "call_function":
196 continue
197 comm = node.meta.get("comm_type")
198 target = getattr(node.target, "_qualified_op_name", None) or str(node.target)
199 if comm and "fsdp" in str(comm):
200 fsdp.append((node.name, comm))
201 elif any(k in target for k in ("all_gather", "reduce_scatter", "all_reduce")):
202 tp_nodes.append((node.name, target))
203 rank = dist.get_rank() if dist.is_initialized() else 0
204 if rank == 0:
205 _LOG.info("=" * 70)
206 _LOG.info("[graph] FSDP comm nodes: %s", len(fsdp))
207 for n, c in fsdp[:6]:
208 _LOG.info(" %s: %s", n, c)
209 _LOG.info("[graph] TP(boundary) comm nodes: %s", len(tp_nodes))
210 for n, t in tp_nodes[:6]:
211 _LOG.info(" %s: %s", n, t)
212 _LOG.info("=" * 70)
215def main(): # pylint: disable=too-many-locals
216 """Run the TP + FSDP graph-mode demo end to end."""
217 logging.basicConfig(level=logging.INFO, format="%(message)s")
218 args = parse_args()
219 cfg = load_config(args.config)
221 device_type = setup_distributed()
222 device = (
223 torch.device(device_type)
224 if device_type == "cpu"
225 else torch.device(device_type, dist.get_rank() % torch.npu.device_count())
226 )
227 rank = dist.get_rank()
228 world_size = dist.get_world_size()
230 if rank == 0:
231 _LOG.info("=" * 70)
232 _LOG.info("TP + FSDP graph-mode demo")
233 _LOG.info(" world_size=%s, device=%s", world_size, device_type)
234 _LOG.info("=" * 70)
236 # 1. Model on a real device (distribute_tensor needs real params).
237 tp_size = cfg["parallel"]["tp_size"]
238 sequence_parallel = cfg["parallel"].get("sequence_parallel", False)
239 loss_parallel = cfg["parallel"].get("loss_parallel", False)
240 model = build_model(cfg, device)
242 # 2. automodel mesh: (dp, cp, tp) + dense FSDP sub-mesh. TP group ready.
243 mesh_ctx = build_mesh_context(cfg["parallel"], device_type)
245 # 3. TP sharding plan + production apply. Parameters become plain TP
246 # shards; module forwards are wrapped with boundary TP collectives on
247 # activations. When sequence_parallel=True the boundary I/O contract
248 # changes: attention/mlp boundaries use Shard(1)->Replicate on entry
249 # and Partial->Shard(1) (reduce_scatter) on exit.
250 planner = ShardingPlanner()
251 plan = planner.plan(
252 model,
253 mesh_ctx.device_mesh,
254 tp_size=cfg["parallel"]["tp_size"],
255 sequence_parallel=sequence_parallel,
256 loss_parallel=loss_parallel,
257 )
258 model, _ = apply_sharding_plan(model, plan, mesh_ctx, validate_mode=False)
259 if rank == 0:
260 _LOG.info(
261 "[automodel] TP sharding applied — params are now plain TP shards, "
262 "boundary forwards wrap TP collectives on activations"
263 )
264 if sequence_parallel:
265 _LOG.info("[SP] sequence_parallel=True, tp_size=%s", tp_size)
267 # 3a. No rotary_emb patch needed — the model places
268 # rotary_emb inside Attention.forward (after boundary all-gather),
269 # so cos/sin naturally have the full sequence length.
270 if sequence_parallel and rank == 0:
271 _LOG.info(
272 "[SP] Using AutoModelAdapterForCausalLM — rotary_emb inside "
273 "attention, no patch needed"
274 )
276 # 4. FSDP plan for the graph-mode FSDPPass (mark everything).
277 fsdp_plan = GraphParallelPlan()
278 fsdp_plan.fsdp_mark_pattern("*")
280 # 5. GraphTrainer reuses the automodel mesh (tp group already created);
281 # it registers the dp sub-mesh as "fsdp" and back-fills fsdp_degree.
282 tcfg = cfg["train"]
283 trainer = GraphTrainer(
284 model=model,
285 train_fn=train_fn,
286 pass_config=PassConfig(
287 enable_overlap=cfg["parallel"].get("enable_overlap", True),
288 fsdp_degree=cfg["parallel"]["dp_size"],
289 tp_size=tp_size,
290 sequence_parallel=sequence_parallel,
291 loss_parallel=loss_parallel,
292 ),
293 parallel_plan=fsdp_plan,
294 optimizer_config={
295 "lr": tcfg.get("lr", 1e-4),
296 "grad_clip": tcfg.get("grad_clip", 1.0),
297 },
298 mesh_context=mesh_ctx,
299 device=device,
300 )
302 # 6. Dummy data + explicit compile so we can inspect the graph.
303 # DataSampler yields full-sequence data; SP sharding is handled by
304 # the embedding boundary's reduce-scatter on hidden_states.
305 mcfg = cfg["model"]
306 sampler = DataSampler(
307 vocab_size=mcfg["vocab_size"],
308 batch_size=tcfg["batch_size"],
309 seq_len=tcfg["seq_len"],
310 max_steps=tcfg["max_steps"],
311 tp_size=tp_size,
312 tp_rank=mesh_ctx.tp_rank,
313 sequence_parallel=sequence_parallel,
314 device=device,
315 )
317 # Sample one batch for compilation (shape must match runtime batches)
318 input_batch, label_batch = sampler.sample()
319 if sequence_parallel and rank == 0:
320 _LOG.info(
321 "[SP] sample input full seq_len=%s (SP sharding via embedding "
322 "reduce-scatter)",
323 input_batch.shape[1],
324 )
325 trainer.compile(input_ids=input_batch, labels=label_batch)
326 inspect_graph(trainer)
328 # 7. Training data iterator (DataSampler yields full-sequence batches)
329 _LOG.info("\nStarting training...")
330 trainer.train(
331 iter(sampler),
332 max_steps=tcfg["max_steps"],
333 log_interval=cfg["logging"]["log_interval"],
334 )
335 if rank == 0:
336 _LOG.info("=" * 70)
337 _LOG.info("Training completed!")
338 _LOG.info("=" * 70)
340 if dist.is_initialized():
341 dist.destroy_process_group()
344if __name__ == "__main__":
345 main()