From f560256d599f1923f091f93bba2c128677e91f18 Mon Sep 17 00:00:00 2001 From: gkarfakis19 Date: Sat, 1 Nov 2025 16:52:20 -0700 Subject: [PATCH 01/77] init llm astra --- .gitmodules | 1 + astra-sim | 2 +- 2 files changed, 2 insertions(+), 1 deletion(-) diff --git a/.gitmodules b/.gitmodules index 41a0a515..00b3fb76 100644 --- a/.gitmodules +++ b/.gitmodules @@ -1,3 +1,4 @@ [submodule "astra-sim"] path = astra-sim url = https://github.com/gkarfakis19/astra-sim + branch = faraz diff --git a/astra-sim b/astra-sim index c13cfa86..ce3b64b7 160000 --- a/astra-sim +++ b/astra-sim @@ -1 +1 @@ -Subproject commit c13cfa86c7ba71813faf25fca3145ccc89ef307f +Subproject commit ce3b64b7c02962d77c738e23cb9971818fd4d0a4 From d6d78930f4a1f474228f5fb5d3d5200eeb7878bd Mon Sep 17 00:00:00 2001 From: gkarfakis19 Date: Sat, 1 Nov 2025 18:14:24 -0700 Subject: [PATCH 02/77] fix submodule --- .gitmodules | 3 +-- 1 file changed, 1 insertion(+), 2 deletions(-) diff --git a/.gitmodules b/.gitmodules index 00b3fb76..34f1cb50 100644 --- a/.gitmodules +++ b/.gitmodules @@ -1,4 +1,3 @@ [submodule "astra-sim"] path = astra-sim - url = https://github.com/gkarfakis19/astra-sim - branch = faraz + url = https://github.com/gkarfakis19/astra-sim \ No newline at end of file From 629c02b42175365f03182ff35ae1466af7263dcd Mon Sep 17 00:00:00 2001 From: gkarfakis19 Date: Sat, 1 Nov 2025 19:30:32 -0700 Subject: [PATCH 03/77] updated astra ver --- astra-sim | 2 +- 1 file changed, 1 insertion(+), 1 deletion(-) diff --git a/astra-sim b/astra-sim index ce3b64b7..6a22754b 160000 --- a/astra-sim +++ b/astra-sim @@ -1 +1 @@ -Subproject commit ce3b64b7c02962d77c738e23cb9971818fd4d0a4 +Subproject commit 6a22754be1578ccc7116848f84df5d62898f48ce From 7b019de129efbf3899f0fb3227e98f6a4a8331f2 Mon Sep 17 00:00:00 2001 From: gkarfakis19 Date: Sat, 1 Nov 2025 22:33:23 -0700 Subject: [PATCH 04/77] updated config to handle new astra ver --- astrasim_lib/config_generation.py | 22 +++++++++++++++++++++- 1 file changed, 21 insertions(+), 1 deletion(-) diff --git a/astrasim_lib/config_generation.py b/astrasim_lib/config_generation.py index c68e15c5..69a5e20f 100644 --- a/astrasim_lib/config_generation.py +++ b/astrasim_lib/config_generation.py @@ -55,6 +55,16 @@ def choose_collective(alg: str, topo: str, op: str) -> str: return "direct" if topo == "Switch": return "halvingDoubling" + if topo == "Torus2D": + return "torus2d" + if topo == "Mesh": + return "mesh" + if topo == "Mesh2D": + return "torus2d" + if topo == "HyperCube": + return "hypercube" + if topo == "KingMesh2D": + return "torus2d" return "ring" @@ -70,13 +80,23 @@ def derive_topology_from_hw(hw_obj) -> str: """Map DeepFlow network topology enums to AstraSim names.""" try: topo = getattr(hw_obj.network_topology.intra, "topology", None) - topo_str = (topo or "fc").lower() + topo_str = topo.lower() if topo_str in ("fc", "fullyconnected", "fully_connected", "fully-connected"): return "FullyConnected" if topo_str in ("ring",): return "Ring" if topo_str in ("switch",): return "Switch" + if topo_str in ("torus2d",): + return "Torus2D" + if topo_str in ("mesh",): + return "Mesh" + if topo_str in ("hypercube",): + return "HyperCube" + if topo_str in ("mesh2d",): + return "Mesh2D" + if topo_str in ("kingmesh2d",): + return "KingMesh2D" except Exception: # pragma: no cover - defensive pass return "FullyConnected" From 9b324a5cf6bd8156999e76675db6d7e74171cb42 Mon Sep 17 00:00:00 2001 From: gkarfakis19 Date: Sat, 1 Nov 2025 23:21:32 -0700 Subject: [PATCH 05/77] Some cleanup and fixed major network double counting flattened bug --- LLM_execution.py | 1 - time_calculation_LLM.py | 71 ++++++++++++++++++++++++++++++++--------- time_calculation_inf.py | 15 +++++++-- 3 files changed, 69 insertions(+), 18 deletions(-) diff --git a/LLM_execution.py b/LLM_execution.py index a4e59768..fa3e8b73 100644 --- a/LLM_execution.py +++ b/LLM_execution.py @@ -533,7 +533,6 @@ def run(self, mode: ExecutionMode) -> ExecutionResult: return self._run_full_astrasim_hierarchical() if mode == ExecutionMode.FULL_ASTRASIM_FLATTENED: return self._run_full_astrasim_flattened() - raise ValueError(f"Unsupported execution mode: {mode}") def _run_pipeline_with_analytical_comm(self, declared_mode: ExecutionMode) -> ExecutionResult: if declared_mode == ExecutionMode.HYBRID: diff --git a/time_calculation_LLM.py b/time_calculation_LLM.py index 7ba1c8dd..dca49468 100644 --- a/time_calculation_LLM.py +++ b/time_calculation_LLM.py @@ -678,7 +678,7 @@ def _tensor_parallelism_gemm_forward(self, gemm: Tuple[int, ...], name: str, gem raise ValueError(f"Unsupported gemm type: {gemm_type}") if size_bytes > 0: - total_bytes = size_bytes #we already has the totol bytes for all reduce not bytes per rank + total_bytes = size_bytes # we already has the total bytes for all reduce not bytes per rank reduction_time = self.get_tensor_reduction_time(total_bytes, kind=comm_kind_fwd, participants=participants, name=name) @@ -878,13 +878,16 @@ def get_linear_softmax_f(self, gemm): """ _, effective_m, k, n = self._effective_dims(gemm) - gemm_time, reduction_time, size_bytes,_,_ = self._tensor_parallelism_gemm_forward(gemm, "linear_softmax_f", gemm_type=GemmType.LINEAR_SOFTMAX) + # Previous TP-aware path kept comm time here: + # gemm_time, reduction_time, size_bytes, _, _ = self._tensor_parallelism_gemm_forward( + # gemm, "linear_softmax_f", gemm_type=GemmType.LINEAR_SOFTMAX + # ) + # Linear softmax is modeled as running on a single device within each TP group. + gemm_time, _, _, _, _ = self.single_gpu_gemm_forward(gemm, "linear_softmax_f", gemm_type=GemmType.LINEAR_SOFTMAX) elements = effective_m * n / (self.tp * self.cp) # each tp-cp group holds a shard of the vocab dimension point_flop = elements * SOFTMAX_FORWARD_FLOPS_PER_ELEMENT - # TODO: - # unsure if should be precision.activations or precision.stats point_mem = self.precision.activations * elements * SOFTMAX_FORWARD_MEM_ACCESSES point_time = self.roofline( point_flop, @@ -901,14 +904,19 @@ def get_linear_softmax_f(self, gemm): ) print("Linear Softmax (f) point_time: {:,}\n".format(point_time)) - return gemm_time + reduction_time + point_time, point_mem + return gemm_time + point_time, point_mem def get_linear_softmax_b(self, gemm): _, effective_m, k, n = self._effective_dims(gemm) - gemm_time, reduction_time, size_bytes = self._tensor_parallelism_gemm_backward(gemm, "linear_softmax_b", gemm_type=GemmType.LINEAR_SOFTMAX) + # Previous TP-aware path kept comm time here: + # gemm_time, reduction_time, size_bytes = self._tensor_parallelism_gemm_backward( + # gemm, "linear_softmax_b", gemm_type=GemmType.LINEAR_SOFTMAX + # ) + # Linear softmax is modeled as running on a single device within each TP group. + gemm_time, _, _ = self.single_gpu_gemm_backward(gemm, "linear_softmax_b", gemm_type=GemmType.LINEAR_SOFTMAX) elements = effective_m * n / (self.tp * self.cp) # each tp-cp group holds a shard of the vocab dimension point_flop = elements * SOFTMAX_BACKWARD_FLOPS_PER_ELEMENT # TODO: @@ -930,7 +938,7 @@ def get_linear_softmax_b(self, gemm): ) print("Linear Softmax (b) point_time: {:,}\n".format(point_time)) - return gemm_time + reduction_time + point_time + return gemm_time + point_time def get_scale_softmax_f(self, gemm): """ Estimate time for scale + softmax forward. @@ -1406,12 +1414,14 @@ def compute_all_gemm_and_node_times(self, batch_size, vocab_size, hidden_dim, se # Helpers to handle variable-length returns (supports TP/Single returns with flops/mem and CP/Hybrid without) def _extract_forward(ret): # ret may be (time, red, size, flops, mem) or (time, red, size) - if len(ret) >= 5: - return ret[0], ret[1], ret[2], ret[3], ret[4] + if len(ret) == 5: + time, reduction, size, flops, mem = ret + return time, reduction, size, flops, mem elif len(ret) == 3: - return ret[0], ret[1], ret[2], 0, [] + time, reduction, size = ret + return time, reduction, size, 0, [] else: - return ret[0], 0, 0, 0, [] + raise ValueError(f"Unsupported return length: {len(ret)}") # QKV Projection GEMM qkv_proj_gemm_f, qkv_proj_reduction_f, qkv_proj_size_f, qkv_proj_flops_f, qkv_proj_mem_f = _extract_forward( @@ -1616,6 +1626,16 @@ def _extract_forward(ret): transformer_results['MHA'] = { 'forward': mha_time_f, 'backward': mha_time_b, + 'forward_compute': ( + transformer_results['qkv_proj']['forward_gemm'] + + transformer_results['attention']['forward_gemm'] + + transformer_results['output_proj']['forward_gemm'] + ), + 'backward_compute': ( + transformer_results['qkv_proj']['backward_gemm'] + + transformer_results['attention']['backward_gemm'] + + transformer_results['output_proj']['backward_gemm'] + ), "forward_reduction": qkv_proj_reduction_f + attention_reduction_f + out_proj_reduction_f, "backward_reduction": qkv_proj_reduction_b + attention_reduction_b + out_proj_reduction_b, "comm_size_forward": qkv_proj_size_f + attention_size_f + out_proj_size_f, @@ -1636,6 +1656,16 @@ def _extract_forward(ret): transformer_results['MLP'] = { 'forward': ffn_time_f, 'backward': ffn_time_b, + 'forward_compute': ( + transformer_results['ffn1']['forward_gemm'] + + transformer_results['ffn2']['forward_gemm'] + + transformer_results['gelu']['forward'] + ), + 'backward_compute': ( + transformer_results['ffn1']['backward_gemm'] + + transformer_results['ffn2']['backward_gemm'] + + transformer_results['gelu']['backward'] + ), "forward_reduction": ffn2_reduction_f + ffn1_reduction_f, "backward_reduction": ffn1_reduction_b + ffn2_reduction_b, "comm_size_forward": ffn2_size_f + ffn1_size_f, @@ -1928,6 +1958,17 @@ def _prepare_execution_graphs( transformer_operation_entries: List[Dict[str, Any]] = [] transformer_comm_metadata: Dict[str, Dict[str, Any]] = {} parallelism_mode = self.get_parallelism_mode() + + def _require_compute(spec: Dict[str, Any], direction: str, entry_name: str) -> float: + key_compute = f"{direction}_compute" + key_gemm = f"{direction}_gemm" + if key_compute in spec: + return spec[key_compute] + if key_gemm in spec: + return spec[key_gemm] + print(f"Keys/Values: {spec}") + raise KeyError(f"Missing compute-only duration for transformer entry '{entry_name}' direction '{direction}'") + if parallelism_mode in (ParallelismMode.CONTEXT, ParallelismMode.TENSOR_CONTEXT_HYBRID): for key in ("layernorm1", "qkv_proj", "attention", "output_proj", "layernorm2", "MLP"): try: @@ -1936,8 +1977,8 @@ def _prepare_execution_graphs( print(f"Key {key} not found in transformer_results") print(transformer_results) exit() - fwd_time = spec['forward_gemm'] if 'forward_gemm' in spec else spec['forward_compute'] if 'forward_compute' in spec else spec['forward'] - bwd_time = spec['backward_gemm'] if 'backward_gemm' in spec else spec['backward_compute'] if 'backward_compute' in spec else spec['backward'] + fwd_time = _require_compute(spec, "forward", key) + bwd_time = _require_compute(spec, "backward", key) fwd_red = spec.get('forward_reduction', 0.0) bwd_red = spec.get('backward_reduction', 0.0) comm_bytes_fwd = spec.get('comm_size_forward', 0) @@ -1969,8 +2010,8 @@ def _prepare_execution_graphs( elif parallelism_mode in (ParallelismMode.TENSOR, ParallelismMode.TENSOR_SEQUENCE, ParallelismMode.SINGLE): for key in ("layernorm1", "MHA", "layernorm2", "MLP"): spec = transformer_results[key] - fwd_time = spec['forward'] - bwd_time = spec['backward'] + fwd_time = _require_compute(spec, "forward", key) + bwd_time = _require_compute(spec, "backward", key) fwd_red = spec.get('forward_reduction', 0.0) bwd_red = spec.get('backward_reduction', 0.0) comm_bytes_fwd = spec.get('comm_size_forward', 0) diff --git a/time_calculation_inf.py b/time_calculation_inf.py index 3f08e605..a0b9c259 100644 --- a/time_calculation_inf.py +++ b/time_calculation_inf.py @@ -149,12 +149,15 @@ def _build_decode_transformer_results( attention_forward = ( attention_score_time + attention_scale_softmax_f + attention_output_time + attention_reduction ) + attention_forward_compute = attention_score_time + attention_scale_softmax_f + attention_output_time out_proj_forward = out_proj_time + out_proj_reduction mha_forward = qkv_proj_forward + attention_forward + out_proj_forward + mha_forward_compute = qkv_proj_time + attention_forward_compute + out_proj_time ffn1_forward = ffn1_time + ffn1_reduction ffn2_forward = ffn2_time + ffn2_reduction mlp_forward = ffn1_forward + act_f + ffn2_forward + mlp_forward_compute = ffn1_time + act_f + ffn2_time layernorm1_forward = residual1_f + layernorm1_f layernorm2_forward = residual2_f + layernorm2_f @@ -215,6 +218,8 @@ def _build_decode_transformer_results( "MHA": { "forward": mha_forward, "backward": 0.0, + "forward_compute": mha_forward_compute, + "backward_compute": 0.0, "forward_reduction": qkv_proj_reduction + attention_reduction + out_proj_reduction, "backward_reduction": 0.0, "comm_size_forward": qkv_proj_size + attention_comm_bytes + out_proj_size, @@ -253,6 +258,8 @@ def _build_decode_transformer_results( "MLP": { "forward": mlp_forward, "backward": 0.0, + "forward_compute": mlp_forward_compute, + "backward_compute": 0.0, "forward_reduction": ffn1_reduction + ffn2_reduction, "backward_reduction": 0.0, "comm_size_forward": ffn1_size + ffn2_size, @@ -266,16 +273,18 @@ def _build_decode_transformer_results( "gelu": { "forward": act_f, "backward": 0.0, - "forward_gemm": act_f, + "forward_compute": act_f, + "backward_compute": 0.0, "forward_reduction": 0.0, - "backward_gemm": 0.0, "backward_reduction": 0.0, "comm_size_forward": 0, "comm_size_backward": 0, }, "layernorm1": { "forward": layernorm1_forward, + "forward_compute": layernorm1_f, "backward": 0.0, + "backward_compute": 0.0, "forward_reduction": layernorm1_reduction, "backward_reduction": 0.0, "comm_size_forward": layernorm1_bytes, @@ -283,7 +292,9 @@ def _build_decode_transformer_results( }, "layernorm2": { "forward": layernorm2_forward, + "forward_compute": layernorm2_f, "backward": 0.0, + "backward_compute": 0.0, "forward_reduction": layernorm2_reduction, "backward_reduction": 0.0, "comm_size_forward": layernorm2_bytes, From 75c2ee4a5411dbec80853443b688916838110337 Mon Sep 17 00:00:00 2001 From: gkarfakis19 Date: Sun, 2 Nov 2025 02:32:04 -0800 Subject: [PATCH 06/77] time calc refactor stage 1 --- time_calculation_LLM.py | 919 +++++++++++++++++++++++----------------- time_calculation_inf.py | 402 ++++++++++-------- timing_model.py | 233 ++++++++++ 3 files changed, 988 insertions(+), 566 deletions(-) create mode 100644 timing_model.py diff --git a/time_calculation_LLM.py b/time_calculation_LLM.py index dca49468..a142cd96 100644 --- a/time_calculation_LLM.py +++ b/time_calculation_LLM.py @@ -2,13 +2,15 @@ import os import json from enum import Enum -from typing import Any, Dict, Tuple, Optional, List +from typing import Any, Dict, Tuple, Optional, List, Mapping, Sequence, Set +from collections.abc import Mapping as MappingABC, Sequence as SequenceABC import simulate_LLM from LLM_execution import ExecutionMode, LLMExecutionDispatcher from simulate_LLM import Graph import LLM_util from time_calculation import TimeCalculation from itertools import zip_longest # for element-wise aggregation of memory access lists +from timing_model import CommSpec, DirectionTiming, OperationTiming, OperationGroup def _env_flag(name: str) -> bool: value = os.environ.get(name) @@ -1365,23 +1367,36 @@ def _combine_mem(self, *args): for d in args: mem_levels = self._mem_levels(d) for k, v in mem_levels.items(): - combined[k] = combined.get(k, 0) + v + combined[k] = combined.get(k, 0.0) + (v or 0.0) return combined def _mem_levels(self, arr): + # Accept both dict-like and sequence-like inputs since GEMM helpers return lists while + # aggregated OperationTiming instances already carry dicts. MappingABC/SequenceABC cover + # any object implementing the mapping/sequence protocol (e.g., dict, defaultdict, numpy arrays). + if isinstance(arr, MappingABC): + return {str(k): float(v) for k, v in arr.items()} if isinstance(arr, (int, float)): - return {f"L{i}": (arr if i == self.num_levels - 1 else 0) for i in range(self.num_levels)} - return {f"L{i}": v for i, v in enumerate(arr or [])} + return {f"L{i}": (float(arr) if i == self.num_levels - 1 else 0.0) for i in range(self.num_levels)} + if isinstance(arr, SequenceABC) and not isinstance(arr, (str, bytes)): + return {f"L{i}": float(v) for i, v in enumerate(arr)} + return {} # TODO TODO: # we need a significant refactor here. The comm sizes are ingested in a weird way and never used. Instead we use old precomputed sizes. # FIX at some point! def compute_all_gemm_and_node_times(self, batch_size, vocab_size, hidden_dim, seq_len, num_heads, kv_heads, intermediate_size, num_SMs): - """Compute latency for all GEMM operations and node breakdown times in one function.""" + """Compute latency for all GEMM operations and node breakdown times.""" - # Process GEMM shapes gemm_shapes = LLM_util.process_gemm_shapes( - self, batch_size, seq_len, hidden_dim, num_heads, kv_heads, intermediate_size, vocab_size + self, + batch_size, + seq_len, + hidden_dim, + num_heads, + kv_heads, + intermediate_size, + vocab_size, ) if self.debug: print( @@ -1398,6 +1413,7 @@ def compute_all_gemm_and_node_times(self, batch_size, vocab_size, hidden_dim, se "intermediate_size:", intermediate_size, ) + gemm_qkv_proj = gemm_shapes["qkv_proj"] gemm_attention_score = gemm_shapes["attention_score"] gemm_attention_output = gemm_shapes["attention_output"] @@ -1406,295 +1422,434 @@ def compute_all_gemm_and_node_times(self, batch_size, vocab_size, hidden_dim, se gemm_ffn2 = gemm_shapes["ffn2"] gemm_linear = gemm_shapes["linear"] - # Compute GEMM times and organize in structured dict - transformer_results = {} - - # assuming no context-parallelism for now and only running inference + transformer_timings: Dict[str, OperationTiming] = {} + + def _make_direction( + op_name: str, + direction: str, + compute_time: float, + comm_time: float, + comm_bytes: float, + *, + flops: float = 0.0, + memory: Optional[Mapping[str, float]] = None, + ) -> DirectionTiming: + mem_map = dict(memory) if memory else {} + bytes_int = int(math.ceil(float(comm_bytes or 0.0))) + return DirectionTiming( + compute_time=compute_time, + comm_time=comm_time, + comm_bytes=bytes_int, + flops=flops, + memory_accesses=mem_map, + ) - # Helpers to handle variable-length returns (supports TP/Single returns with flops/mem and CP/Hybrid without) - def _extract_forward(ret): - # ret may be (time, red, size, flops, mem) or (time, red, size) + def _extract_forward(ret: Sequence[Any]) -> Tuple[float, float, float, float, Any]: if len(ret) == 5: time, reduction, size, flops, mem = ret return time, reduction, size, flops, mem - elif len(ret) == 3: + if len(ret) == 3: time, reduction, size = ret - return time, reduction, size, 0, [] - else: - raise ValueError(f"Unsupported return length: {len(ret)}") + return time, reduction, size, 0.0, [] + raise ValueError(f"Unsupported return length: {len(ret)}") - # QKV Projection GEMM - qkv_proj_gemm_f, qkv_proj_reduction_f, qkv_proj_size_f, qkv_proj_flops_f, qkv_proj_mem_f = _extract_forward( + # QKV + qkv_proj_gemm_f, qkv_proj_reduction_f, qkv_proj_size_f, qkv_proj_flops_f, qkv_proj_mem_f = _extract_forward( self.parallelism_gemm_forward(gemm_qkv_proj, "qkv_projection_f", gemm_type=GemmType.QKV) ) - qkv_proj_gemm_b, qkv_proj_reduction_b, qkv_proj_size_b = self.parallelism_gemm_backward(gemm_qkv_proj, "qkv_projection_b", gemm_type=GemmType.QKV) - qkv_proj_f = qkv_proj_gemm_f + qkv_proj_reduction_f - qkv_proj_b = qkv_proj_gemm_b + qkv_proj_reduction_b - transformer_results['qkv_proj'] = { - 'forward': qkv_proj_f, 'backward': qkv_proj_b, - 'forward_gemm': qkv_proj_gemm_f, 'forward_reduction': qkv_proj_reduction_f, - 'backward_gemm': qkv_proj_gemm_b, 'backward_reduction': qkv_proj_reduction_b, - 'comm_size_forward': qkv_proj_size_f, 'comm_size_backward': qkv_proj_size_b, - 'flops': qkv_proj_flops_f, - 'memory_accesses': self._mem_levels(qkv_proj_mem_f), - } + qkv_proj_gemm_b, qkv_proj_reduction_b, qkv_proj_size_b = self.parallelism_gemm_backward( + gemm_qkv_proj, "qkv_projection_b", gemm_type=GemmType.QKV + ) + qkv_forward = _make_direction( + "qkv_proj", + "forward", + compute_time=qkv_proj_gemm_f, + comm_time=qkv_proj_reduction_f, + comm_bytes=qkv_proj_size_f, + flops=qkv_proj_flops_f, + memory=self._mem_levels(qkv_proj_mem_f), + ) + qkv_backward = _make_direction( + "qkv_proj", + "backward", + compute_time=qkv_proj_gemm_b, + comm_time=qkv_proj_reduction_b, + comm_bytes=qkv_proj_size_b, + ) + transformer_timings["qkv_proj"] = OperationTiming("qkv_proj", forward=qkv_forward, backward=qkv_backward) - if self.flash_attention is False: - # Attention Score GEMM - attn_score_gemm_f, attn_score_reduction_f, attn_score_size_f, attn_score_flops_f, attn_score_mem_f = _extract_forward( + if not self.flash_attention: + attn_score_gemm_f, attn_score_reduction_f, attn_score_size_f, attn_score_flops_f, attn_score_mem_f = _extract_forward( self.parallelism_gemm_forward(gemm_attention_score, "attention_score_f", gemm_type=GemmType.ATTENTION_SCORE) ) - attn_score_gemm_b, attn_score_reduction_b, attn_score_size_b = self.parallelism_gemm_backward(gemm_attention_score, "attention_score_b", gemm_type=GemmType.ATTENTION_SCORE) - attention_score_f = attn_score_gemm_f + attn_score_reduction_f - attention_score_b = attn_score_gemm_b + attn_score_reduction_b - transformer_results['attention_score'] = { - 'forward': attention_score_f, 'backward': attention_score_b, - 'forward_gemm': attn_score_gemm_f, 'forward_reduction': attn_score_reduction_f, - 'backward_gemm': attn_score_gemm_b, 'backward_reduction': attn_score_reduction_b, - 'comm_size_forward': attn_score_size_f, 'comm_size_backward': attn_score_size_b, - 'flops': attn_score_flops_f, - 'memory_accesses': self._mem_levels(attn_score_mem_f), - } + attn_score_gemm_b, attn_score_reduction_b, attn_score_size_b = self.parallelism_gemm_backward( + gemm_attention_score, "attention_score_b", gemm_type=GemmType.ATTENTION_SCORE + ) + attn_score_forward = _make_direction( + "attention_score", + "forward", + compute_time=attn_score_gemm_f, + comm_time=attn_score_reduction_f, + comm_bytes=attn_score_size_f, + flops=attn_score_flops_f, + memory=self._mem_levels(attn_score_mem_f), + ) + attn_score_backward = _make_direction( + "attention_score", + "backward", + compute_time=attn_score_gemm_b, + comm_time=attn_score_reduction_b, + comm_bytes=attn_score_size_b, + ) + transformer_timings["attention_score"] = OperationTiming( + "attention_score", + forward=attn_score_forward, + backward=attn_score_backward, + ) + attention_scale_softmax_f = self.get_scale_softmax_f(gemm=gemm_attention_score) attention_scale_softmax_b = self.get_scale_softmax_b(gemm=gemm_attention_score) - transformer_results['attention_scale_softmax'] = {'forward': attention_scale_softmax_f, 'backward': attention_scale_softmax_b} + transformer_timings["attention_scale_softmax"] = OperationTiming( + name="attention_scale_softmax", + forward=_make_direction( + "attention_scale_softmax", + "forward", + compute_time=attention_scale_softmax_f, + comm_time=0.0, + comm_bytes=0.0, + ), + backward=_make_direction( + "attention_scale_softmax", + "backward", + compute_time=attention_scale_softmax_b, + comm_time=0.0, + comm_bytes=0.0, + ), + ) - # Attention Output GEMM - attn_out_gemm_f, attn_out_reduction_f, attn_out_size_f, attn_out_flops_f, attn_out_mem_f = _extract_forward( + attn_out_gemm_f, attn_out_reduction_f, attn_out_size_f, attn_out_flops_f, attn_out_mem_f = _extract_forward( self.parallelism_gemm_forward(gemm_attention_output, "attention_output_f", gemm_type=GemmType.ATTENTION_OUTPUT) ) - attn_out_gemm_b, attn_out_reduction_b, attn_out_size_b = self.parallelism_gemm_backward(gemm_attention_output, "attention_output_b", gemm_type=GemmType.ATTENTION_OUTPUT) - attention_output_f = attn_out_gemm_f + attn_out_reduction_f - attention_output_b = attn_out_gemm_b + attn_out_reduction_b - transformer_results['attention_output'] = { - 'forward': attention_output_f, 'backward': attention_output_b, - 'forward_gemm': attn_out_gemm_f, 'forward_reduction': attn_out_reduction_f, - 'backward_gemm': attn_out_gemm_b, 'backward_reduction': attn_out_reduction_b, - 'comm_size_forward': attn_out_size_f, 'comm_size_backward': attn_out_size_b, - 'flops': attn_out_flops_f, - 'memory_accesses': self._mem_levels(attn_out_mem_f), - } - # Aggregate attention (score + softmax + output) - attention_f = attention_score_f + attention_scale_softmax_f + attention_output_f - attention_b = attention_score_b + attention_scale_softmax_b + attention_output_b - attention_gemm_f = attn_score_gemm_f + attn_out_gemm_f + attention_scale_softmax_f - attention_reduction_f = attn_score_reduction_f + attn_out_reduction_f - attention_gemm_b = attn_score_gemm_b + attn_out_gemm_b + attention_scale_softmax_b - attention_reduction_b = attn_score_reduction_b + attn_out_reduction_b - attention_size_f = attn_score_size_f + attn_out_size_f - attention_size_b = attn_score_size_b + attn_out_size_b - # Aggregate flops and per-level memory accesses (score + output) - attention_flops = (attn_score_flops_f or 0) + (attn_out_flops_f or 0) - attention_mem = [ - (a or 0) + (b or 0) - for a, b in zip_longest(attn_score_mem_f or [], attn_out_mem_f or [], fillvalue=0) - ] - transformer_results["attention"] = { - 'forward': attention_f, - 'backward': attention_b, - "forward_gemm": attention_gemm_f, - "forward_reduction": attention_reduction_f, - "backward_gemm": attention_gemm_b, - "backward_reduction": attention_reduction_b, - "comm_size_forward": attention_size_f, - "comm_size_backward": attention_size_b, - "flops": attention_flops, - "memory_accesses": self._mem_levels(attention_mem), - } - elif self.flash_attention is True: - attention_f, attention_gemm_f, attention_reduction_f, attention_size_f, attention_mem = self.flash_attention_kernel_forward(batch_size,hidden_dim, seq_len, num_heads, kv_heads, num_SMs) - # TODO WIP. - attention_b, attention_gemm_b, attention_reduction_b, attention_size_b = 2 * attention_f, 2 * attention_gemm_f, 2 * attention_reduction_f, 2 * attention_size_f + attn_out_gemm_b, attn_out_reduction_b, attn_out_size_b = self.parallelism_gemm_backward( + gemm_attention_output, "attention_output_b", gemm_type=GemmType.ATTENTION_OUTPUT + ) + attn_out_forward = _make_direction( + "attention_output", + "forward", + compute_time=attn_out_gemm_f, + comm_time=attn_out_reduction_f, + comm_bytes=attn_out_size_f, + flops=attn_out_flops_f, + memory=self._mem_levels(attn_out_mem_f), + ) + attn_out_backward = _make_direction( + "attention_output", + "backward", + compute_time=attn_out_gemm_b, + comm_time=attn_out_reduction_b, + comm_bytes=attn_out_size_b, + ) + transformer_timings["attention_output"] = OperationTiming( + "attention_output", + forward=attn_out_forward, + backward=attn_out_backward, + ) - # flops for flash attention is 2* (QK^T + SV) - attention_flops = 2 * (batch_size * seq_len * (hidden_dim ** 2) + batch_size * (seq_len ** 2) * hidden_dim) - - transformer_results['attention'] = { - 'forward': attention_f, 'backward': attention_b, - 'forward_gemm': attention_gemm_f, 'forward_reduction': attention_reduction_f, - 'backward_gemm': attention_gemm_b, 'backward_reduction': attention_reduction_b, - 'comm_size_forward': attention_size_f, 'comm_size_backward': attention_size_b, - 'flops': attention_flops, - 'memory_accesses': self._mem_levels(attention_mem), # HBM only - } + attention_forward_compute = ( + attn_score_forward.compute_time + + transformer_timings["attention_scale_softmax"].forward.compute_time + + attn_out_forward.compute_time + ) + attention_forward_comm = attn_score_forward.comm_time + attn_out_forward.comm_time + attention_backward_compute = ( + attn_score_backward.compute_time + + transformer_timings["attention_scale_softmax"].backward.compute_time + + attn_out_backward.compute_time + ) + attention_backward_comm = attn_score_backward.comm_time + attn_out_backward.comm_time + attention_comm_bytes_f = attn_score_size_f + attn_out_size_f + attention_comm_bytes_b = attn_score_size_b + attn_out_size_b + attention_mem = self._combine_mem(attn_score_mem_f, attn_out_mem_f) + transformer_timings["attention"] = OperationTiming( + "attention", + forward=_make_direction( + "attention", + "forward", + compute_time=attention_forward_compute, + comm_time=attention_forward_comm, + comm_bytes=attention_comm_bytes_f, + memory=self._mem_levels(attention_mem), + flops=(attn_score_flops_f + attn_out_flops_f), + ), + backward=_make_direction( + "attention", + "backward", + compute_time=attention_backward_compute, + comm_time=attention_backward_comm, + comm_bytes=attention_comm_bytes_b, + ), + ) else: - raise ValueError("flash_attention should be either True or False") + attention_f, attention_gemm_f, attention_reduction_f, attention_size_f, attention_mem = self.flash_attention_kernel_forward( + batch_size, + hidden_dim, + seq_len, + num_heads, + kv_heads, + num_SMs, + ) + attention_b = attention_f * 2 + attention_gemm_b = attention_gemm_f * 2 + attention_reduction_b = attention_reduction_f * 2 + attention_size_b = attention_size_f * 2 + transformer_timings["attention"] = OperationTiming( + "attention", + forward=_make_direction( + "attention", + "forward", + compute_time=attention_gemm_f, + comm_time=attention_reduction_f, + comm_bytes=attention_size_f, + memory=self._mem_levels(attention_mem), + ), + backward=_make_direction( + "attention", + "backward", + compute_time=attention_gemm_b, + comm_time=attention_reduction_b, + comm_bytes=attention_size_b, + ), + ) - # Output Projection GEMM out_proj_gemm_f, out_proj_reduction_f, out_proj_size_f, out_proj_flops_f, out_proj_mem_f = _extract_forward( self.parallelism_gemm_forward(gemm_output_proj, "output_projection_f", gemm_type=GemmType.OUT_PROJ) ) - out_proj_gemm_b, out_proj_reduction_b, out_proj_size_b = self.parallelism_gemm_backward(gemm_output_proj, "output_projection_b", gemm_type=GemmType.OUT_PROJ) - output_proj_f = out_proj_gemm_f + out_proj_reduction_f - output_proj_b = out_proj_gemm_b + out_proj_reduction_b - transformer_results['output_proj'] = { - 'forward': output_proj_f, 'backward': output_proj_b, - 'forward_gemm': out_proj_gemm_f, 'forward_reduction': out_proj_reduction_f, - 'backward_gemm': out_proj_gemm_b, 'backward_reduction': out_proj_reduction_b, - 'comm_size_forward': out_proj_size_f, 'comm_size_backward': out_proj_size_b, - 'flops': out_proj_flops_f, - 'memory_accesses': self._mem_levels(out_proj_mem_f), - } + out_proj_gemm_b, out_proj_reduction_b, out_proj_size_b = self.parallelism_gemm_backward( + gemm_output_proj, "output_projection_b", gemm_type=GemmType.OUT_PROJ + ) + transformer_timings["output_proj"] = OperationTiming( + "output_proj", + forward=_make_direction( + "output_proj", + "forward", + compute_time=out_proj_gemm_f, + comm_time=out_proj_reduction_f, + comm_bytes=out_proj_size_f, + flops=out_proj_flops_f, + memory=self._mem_levels(out_proj_mem_f), + ), + backward=_make_direction( + "output_proj", + "backward", + compute_time=out_proj_gemm_b, + comm_time=out_proj_reduction_b, + comm_bytes=out_proj_size_b, + ), + ) - # FFN1 GEMM - ffn1_gemm_f, ffn1_reduction_f, ffn1_size_f, ffn1_flops_f, ffn1_mem_f = _extract_forward( + ffn1_gemm_f, ffn1_reduction_f, ffn1_size_f, ffn1_flops_f, ffn1_mem_f = _extract_forward( self.parallelism_gemm_forward(gemm_ffn1, "ffn_f", gemm_type=GemmType.FFN1) ) - ffn1_gemm_b, ffn1_reduction_b, ffn1_size_b = self.parallelism_gemm_backward(gemm_ffn1, "ffn_b", gemm_type=GemmType.FFN1) - ffn1_f = ffn1_gemm_f + ffn1_reduction_f - ffn1_b = ffn1_gemm_b + ffn1_reduction_b - transformer_results['ffn1'] = { - 'forward': ffn1_f, 'backward': ffn1_b, - 'forward_gemm': ffn1_gemm_f, 'forward_reduction': ffn1_reduction_f, - 'backward_gemm': ffn1_gemm_b, 'backward_reduction': ffn1_reduction_b, - 'comm_size_forward': ffn1_size_f, 'comm_size_backward': ffn1_size_b, - 'flops': ffn1_flops_f, - 'memory_accesses': self._mem_levels(ffn1_mem_f), - } + ffn1_gemm_b, ffn1_reduction_b, ffn1_size_b = self.parallelism_gemm_backward( + gemm_ffn1, "ffn_b", gemm_type=GemmType.FFN1 + ) + transformer_timings["ffn1"] = OperationTiming( + "ffn1", + forward=_make_direction( + "ffn1", + "forward", + compute_time=ffn1_gemm_f, + comm_time=ffn1_reduction_f, + comm_bytes=ffn1_size_f, + flops=ffn1_flops_f, + memory=self._mem_levels(ffn1_mem_f), + ), + backward=_make_direction( + "ffn1", + "backward", + compute_time=ffn1_gemm_b, + comm_time=ffn1_reduction_b, + comm_bytes=ffn1_size_b, + ), + ) - # FFN2 GEMM ffn2_gemm_f, ffn2_reduction_f, ffn2_size_f, ffn2_flops_f, ffn2_mem_f = _extract_forward( self.parallelism_gemm_forward(gemm_ffn2, "ffn2_f", gemm_type=GemmType.FFN2) ) - ffn2_gemm_b, ffn2_reduction_b, ffn2_size_b = self.parallelism_gemm_backward(gemm_ffn2, "ffn2_b", gemm_type=GemmType.FFN2) - ffn2_f = ffn2_gemm_f + ffn2_reduction_f - ffn2_b = ffn2_gemm_b + ffn2_reduction_b - transformer_results['ffn2'] = { - 'forward': ffn2_f, 'backward': ffn2_b, - 'forward_gemm': ffn2_gemm_f, 'forward_reduction': ffn2_reduction_f, - 'backward_gemm': ffn2_gemm_b, 'backward_reduction': ffn2_reduction_b, - 'comm_size_forward': ffn2_size_f , 'comm_size_backward': ffn2_size_b, - 'flops': ffn2_flops_f, - 'memory_accesses': self._mem_levels(ffn2_mem_f), - } - - + ffn2_gemm_b, ffn2_reduction_b, ffn2_size_b = self.parallelism_gemm_backward( + gemm_ffn2, "ffn2_b", gemm_type=GemmType.FFN2 + ) + transformer_timings["ffn2"] = OperationTiming( + "ffn2", + forward=_make_direction( + "ffn2", + "forward", + compute_time=ffn2_gemm_f, + comm_time=ffn2_reduction_f, + comm_bytes=ffn2_size_f, + flops=ffn2_flops_f, + memory=self._mem_levels(ffn2_mem_f), + ), + backward=_make_direction( + "ffn2", + "backward", + compute_time=ffn2_gemm_b, + comm_time=ffn2_reduction_b, + comm_bytes=ffn2_size_b, + ), + ) - # Calculate non-GEMM operations embedding_f, embedding_mem = self.get_embedding_f(vocab_size=vocab_size, seq_len=seq_len, hidden_dim=hidden_dim) embedding_b = self.get_embedding_b(vocab_size=vocab_size, seq_len=seq_len, hidden_dim=hidden_dim) - transformer_results['embedding'] = {'forward': embedding_f, 'backward': embedding_b, 'memory_accesses': self._mem_levels(embedding_mem)} + transformer_timings["embedding"] = OperationTiming( + "embedding", + forward=_make_direction( + "embedding", + "forward", + compute_time=embedding_f, + comm_time=0.0, + comm_bytes=0.0, + memory=self._mem_levels(embedding_mem), + ), + backward=_make_direction( + "embedding", + "backward", + compute_time=embedding_b, + comm_time=0.0, + comm_bytes=0.0, + ), + ) residual1_f = self.get_residual_f(tensor_shape=gemm_output_proj) residual1_b = self.get_residual_b(tensor_shape=gemm_output_proj) - # transformer_results['residual1'] = {'forward': residual1_f, 'backward': residual1_b} - - layernorm1_f, layernorm_reduction_f, LN1_comm_bytes_f= self.get_layernorm_f(batch=batch_size, seq_len=seq_len, d_model=hidden_dim) - layernorm1_b, layernorm_reduction_b, LN1_comm_bytes_b= self.get_layernorm_b(batch=batch_size, seq_len=seq_len, d_model=hidden_dim) - transformer_results['layernorm1'] = { - 'forward': layernorm1_f + residual1_f +layernorm_reduction_f, - 'forward_compute': layernorm1_f + residual1_f, - 'forward_reduction':layernorm_reduction_f, - 'backward': layernorm1_b + residual1_b + layernorm_reduction_b, - "backward_compute":layernorm1_b + residual1_b , - 'backward_reduction': layernorm_reduction_b, - 'comm_size_forward': LN1_comm_bytes_f, 'comm_size_backward': LN1_comm_bytes_b - } + layernorm1_f, layernorm1_reduction_f, LN1_comm_bytes_f = self.get_layernorm_f( + batch=batch_size, + seq_len=seq_len, + d_model=hidden_dim, + ) + layernorm1_b, layernorm1_reduction_b, LN1_comm_bytes_b = self.get_layernorm_b( + batch=batch_size, + seq_len=seq_len, + d_model=hidden_dim, + ) + transformer_timings["layernorm1"] = OperationTiming( + "layernorm1", + forward=_make_direction( + "layernorm1", + "forward", + compute_time=layernorm1_f + residual1_f, + comm_time=layernorm1_reduction_f, + comm_bytes=LN1_comm_bytes_f, + ), + backward=_make_direction( + "layernorm1", + "backward", + compute_time=layernorm1_b + residual1_b, + comm_time=layernorm1_reduction_b, + comm_bytes=LN1_comm_bytes_b, + ), + ) residual2_f = self.get_residual_f(tensor_shape=gemm_ffn2) residual2_b = self.get_residual_b(tensor_shape=gemm_ffn2) - # transformer_results['residual2'] = {'forward': residual2_f, 'backward': residual2_b} + layernorm2_f, layernorm2_reduction_f, LN2_comm_bytes_f = self.get_layernorm_f( + batch=batch_size, + seq_len=seq_len, + d_model=hidden_dim, + ) + layernorm2_b, layernorm2_reduction_b, LN2_comm_bytes_b = self.get_layernorm_b( + batch=batch_size, + seq_len=seq_len, + d_model=hidden_dim, + ) + transformer_timings["layernorm2"] = OperationTiming( + "layernorm2", + forward=_make_direction( + "layernorm2", + "forward", + compute_time=layernorm2_f + residual2_f, + comm_time=layernorm2_reduction_f, + comm_bytes=LN2_comm_bytes_f, + ), + backward=_make_direction( + "layernorm2", + "backward", + compute_time=layernorm2_b + residual2_b, + comm_time=layernorm2_reduction_b, + comm_bytes=LN2_comm_bytes_b, + ), + ) - layernorm2_f, layernorm_reduction_f, LN2_comm_bytes_f = self.get_layernorm_f(batch=batch_size, seq_len=seq_len, d_model=hidden_dim) - layernorm2_b, layernorm_reduction_b, LN2_comm_bytes_b = self.get_layernorm_b(batch=batch_size, seq_len=seq_len, d_model=hidden_dim) - transformer_results['layernorm2'] = {'forward': layernorm2_f + residual2_f + layernorm_reduction_f, "forward_compute": layernorm2_f + residual2_f, 'forward_reduction':layernorm_reduction_f,'backward': layernorm2_b + residual2_b + layernorm_reduction_b, "backward_compute":layernorm2_b + residual2_b, 'backward_reduction': layernorm_reduction_b, 'comm_size_forward': LN2_comm_bytes_f, 'comm_size_backward': LN2_comm_bytes_b} - if self.model_type == "llama": act_f = self.get_swiglu_f(tensor_shape=gemm_ffn1) act_b = self.get_swiglu_b(tensor_shape=gemm_ffn1) else: act_f = self.get_gelu_f(tensor_shape=gemm_ffn1) act_b = self.get_gelu_b(tensor_shape=gemm_ffn1) - transformer_results['gelu'] = {'forward': act_f, 'backward': act_b} + transformer_timings["gelu"] = OperationTiming( + "gelu", + forward=_make_direction("gelu", "forward", compute_time=act_f, comm_time=0.0, comm_bytes=0.0), + backward=_make_direction("gelu", "backward", compute_time=act_b, comm_time=0.0, comm_bytes=0.0), + ) linear_softmax_f, linear_softmax_mem = self.get_linear_softmax_f(gemm=gemm_linear) linear_softmax_b = self.get_linear_softmax_b(gemm=gemm_linear) - transformer_results['linear_softmax'] = { - 'forward': linear_softmax_f, - 'backward': linear_softmax_b, - 'memory_accesses': self._mem_levels(linear_softmax_mem)} - - # Calculate MHA and FFN times directly from results dict - mha_time_f = ( - transformer_results['qkv_proj']['forward'] + transformer_results['attention']['forward'] + transformer_results['output_proj']['forward'] - ) - - - mha_time_b = ( - transformer_results['qkv_proj']['backward'] + transformer_results['attention']['backward'] + transformer_results['output_proj']['backward'] - ) - transformer_results['MHA'] = { - 'forward': mha_time_f, - 'backward': mha_time_b, - 'forward_compute': ( - transformer_results['qkv_proj']['forward_gemm'] - + transformer_results['attention']['forward_gemm'] - + transformer_results['output_proj']['forward_gemm'] + transformer_timings["linear_softmax"] = OperationTiming( + "linear_softmax", + forward=_make_direction( + "linear_softmax", + "forward", + compute_time=linear_softmax_f, + comm_time=0.0, + comm_bytes=0.0, + memory=self._mem_levels(linear_softmax_mem), ), - 'backward_compute': ( - transformer_results['qkv_proj']['backward_gemm'] - + transformer_results['attention']['backward_gemm'] - + transformer_results['output_proj']['backward_gemm'] + backward=_make_direction( + "linear_softmax", + "backward", + compute_time=linear_softmax_b, + comm_time=0.0, + comm_bytes=0.0, ), - "forward_reduction": qkv_proj_reduction_f + attention_reduction_f + out_proj_reduction_f, - "backward_reduction": qkv_proj_reduction_b + attention_reduction_b + out_proj_reduction_b, - "comm_size_forward": qkv_proj_size_f + attention_size_f + out_proj_size_f, - "comm_size_backward": qkv_proj_size_b + attention_size_b + out_proj_size_b, - "flops": qkv_proj_flops_f + transformer_results['attention']['flops'] + out_proj_flops_f, - "memory_accesses": self._combine_mem( - qkv_proj_mem_f, - attention_mem, - out_proj_mem_f, - ), - - } - - ffn_time_f = transformer_results['ffn1']['forward'] + transformer_results['ffn2']['forward'] + transformer_results['gelu']['forward'] - ffn_time_b = ( - transformer_results['ffn1']['backward'] + transformer_results['ffn2']['backward'] + transformer_results['gelu']['backward'] ) - transformer_results['MLP'] = { - 'forward': ffn_time_f, - 'backward': ffn_time_b, - 'forward_compute': ( - transformer_results['ffn1']['forward_gemm'] - + transformer_results['ffn2']['forward_gemm'] - + transformer_results['gelu']['forward'] - ), - 'backward_compute': ( - transformer_results['ffn1']['backward_gemm'] - + transformer_results['ffn2']['backward_gemm'] - + transformer_results['gelu']['backward'] + + mha_group = OperationGroup( + "MHA", + operations=( + transformer_timings["qkv_proj"], + transformer_timings["attention"], + transformer_timings["output_proj"], ), - "forward_reduction": ffn2_reduction_f + ffn1_reduction_f, - "backward_reduction": ffn1_reduction_b + ffn2_reduction_b, - "comm_size_forward": ffn2_size_f + ffn1_size_f, - "comm_size_backward": ffn1_size_b + ffn2_size_b, - "flops": ffn1_flops_f + ffn2_flops_f, - "memory_accesses": self._combine_mem( - ffn1_mem_f, - ffn2_mem_f, + ) + mlp_group = OperationGroup( + "MLP", + operations=( + transformer_timings["ffn1"], + transformer_timings["gelu"], + transformer_timings["ffn2"], ), - } - # Calculate transformer times directly - + ) + transformer_time_f = ( - transformer_results['MHA']['forward'] + transformer_results['MLP']['forward'] + - transformer_results['layernorm1']['forward'] + transformer_results['layernorm2']['forward'] + mha_group.forward_total_time() + + mlp_group.forward_total_time() + + transformer_timings["layernorm1"].total_forward_time() + + transformer_timings["layernorm2"].total_forward_time() ) transformer_time_b = ( - transformer_results['MHA']['backward'] + transformer_results['MLP']['backward'] + - transformer_results['layernorm1']['backward'] + transformer_results['layernorm2']['backward'] + mha_group.backward_total_time() + + mlp_group.backward_total_time() + + transformer_timings["layernorm1"].total_backward_time() + + transformer_timings["layernorm2"].total_backward_time() ) - node_breakdown = { - 'transformer_time_f': transformer_time_f, - 'transformer_time_b': transformer_time_b, - 'embedding_f': transformer_results['embedding']['forward'], - 'embedding_b': transformer_results['embedding']['backward'], - 'linear_softmax_f': transformer_results['linear_softmax']['forward'], - 'linear_softmax_b': transformer_results['linear_softmax']['backward'] + "transformer_time_f": transformer_time_f, + "transformer_time_b": transformer_time_b, + "embedding_f": transformer_timings["embedding"].total_forward_time(), + "embedding_b": transformer_timings["embedding"].total_backward_time(), + "linear_softmax_f": transformer_timings["linear_softmax"].total_forward_time(), + "linear_softmax_b": transformer_timings["linear_softmax"].total_backward_time(), } if self._generate_graphs: @@ -1702,7 +1857,9 @@ def _extract_forward(ret): with open(results_path, "w", encoding="utf-8") as results_file: json.dump( { - "transformer_results": transformer_results, + "transformer_results": { + name: timing.to_dict() for name, timing in transformer_timings.items() + }, "node_breakdown": node_breakdown, }, results_file, @@ -1710,7 +1867,9 @@ def _extract_forward(ret): sort_keys=True, ) - return transformer_results, node_breakdown + return transformer_timings, node_breakdown + + def _effective_transformer_batch(self) -> int: @@ -1816,77 +1975,6 @@ def _build_comm_metadata( } return metadata - def _populate_transformer_comm_metadata( - self, - entry: Dict[str, Any], - metadata: Dict[str, Dict[str, Any]], - # gemm_spec: Tuple[int, ...], - comm_bytes_fwd: int, - comm_bytes_bwd: int, - ) -> None: - """Attach tensor-parallel collectives for a GEMM to metadata and entry.""" - - - if not self.tp or self.tp <= 1 and self.cp <= 1: - return - - - - def add_comm(direction: str, suffix: str, kind: str, size_bytes: float, participants: int, interconnect: str) -> None: - if participants <= 1: - return - bytes_int = int(math.ceil(size_bytes)) - if bytes_int <= 0: - return - key = f"{entry['name']}_{direction}_{suffix}" - # ensure uniqueness when multiple collectives share the same suffix - unique_key = key - counter = 1 - while unique_key in metadata: - counter += 1 - unique_key = f"{key}_{counter}" - - metadata[unique_key] = { - 'size': bytes_int, - 'type': kind, - 'participants': int(participants), - 'interconnect_type': interconnect, - } - entry[direction]['comm_keys'].append(unique_key) - - parallelism_mode = self.get_parallelism_mode() - rules_by_mode = COMMUNICATION_RULES.get(parallelism_mode) - if not rules_by_mode: - return - - spec = rules_by_mode.get(entry['name']) - if spec is None: - spec = rules_by_mode.get(COMM_RULE_DEFAULT_KEY) - if not spec: - return - - - participants_lookup = { - 'tp': getattr(self, 'tp', 0), - 'cp': getattr(self, 'cp', 0), - 'dp': getattr(self, 'dp', 0), - 'lp': getattr(self, 'lp', 0), - } - - for direction in ('forward', 'backward'): - rule = spec.get(direction) - if not rule: - continue - size_bytes = comm_bytes_fwd if direction == 'forward' else comm_bytes_bwd - scope = rule.get('participants') - participants = participants_lookup.get(scope or '', 0) - if participants <= 1: - continue - kind = rule['kind'] - suffix = rule.get('suffix', kind) - interconnect = rule.get('interconnect', scope or 'tp') - add_comm(direction, suffix, kind, size_bytes, participants, interconnect) - def _build_interconnect_params(self) -> Dict[str, Tuple[float, float]]: return { 'dp': (self.IBD, self.LLD), @@ -1899,7 +1987,7 @@ def _prepare_execution_graphs( self, *, node_breakdown: Dict[str, float], - transformer_results: Dict[str, Dict[str, Any]], + transformer_timings: Dict[str, OperationTiming], batch_size: int, seq_len: int, hidden_dim: int, @@ -1936,7 +2024,6 @@ def _prepare_execution_graphs( local_comp = self.get_data_parallel_local_computation(hidden_dim, intermediate_size) # these are used for dp all-reduce/reduce-scatter. - # TODO: check the embedding_size below. I think only the first term is needed. embedding_size = math.ceil(self.precision.grad_communication * vocab_size * hidden_dim) + math.ceil(self.precision.grad_communication * seq_len * hidden_dim * batch_size) softmax_size = math.ceil(self.precision.grad_communication * hidden_dim * vocab_size) cross_layer_bytes = self.get_inter_layer_comm_latency_llm(batch_size, hidden_dim, seq_len)[1] @@ -1958,90 +2045,146 @@ def _prepare_execution_graphs( transformer_operation_entries: List[Dict[str, Any]] = [] transformer_comm_metadata: Dict[str, Dict[str, Any]] = {} parallelism_mode = self.get_parallelism_mode() + rules_by_mode = COMMUNICATION_RULES.get(parallelism_mode) or {} + participants_lookup = { + 'tp': int(getattr(self, 'tp', 0) or 0), + 'cp': int(getattr(self, 'cp', 0) or 0), + 'dp': int(getattr(self, 'dp', 0) or 0), + 'lp': int(getattr(self, 'lp', 0) or 0), + } + group_members = { + "MHA": ("qkv_proj", "attention", "output_proj"), + "MLP": ("ffn1", "gelu", "ffn2"), + } - def _require_compute(spec: Dict[str, Any], direction: str, entry_name: str) -> float: - key_compute = f"{direction}_compute" - key_gemm = f"{direction}_gemm" - if key_compute in spec: - return spec[key_compute] - if key_gemm in spec: - return spec[key_gemm] - print(f"Keys/Values: {spec}") - raise KeyError(f"Missing compute-only duration for transformer entry '{entry_name}' direction '{direction}'") + def _register_specs(specs: Sequence[CommSpec]) -> List[str]: + names: List[str] = [] + for spec in specs: + existing = transformer_comm_metadata.get(spec.name) + if existing: + if ( + existing["size"] != spec.size_bytes + or existing["type"] != spec.kind + or existing["participants"] != spec.participants + or existing["interconnect_type"] != spec.interconnect + ): + raise ValueError( + f"Conflicting CommSpec registration for '{spec.name}' " + f"(existing={existing}, new={spec})" + ) + else: + transformer_comm_metadata[spec.name] = { + "size": spec.size_bytes, + "type": spec.kind, + "participants": spec.participants, + "interconnect_type": spec.interconnect, + } + names.append(spec.name) + return names + + def _build_direction_entry(op_name: str, direction: str, timing: Optional[DirectionTiming]) -> Dict[str, Any]: + if timing is None: + return {"duration": 0.0, "reduction": 0.0, "comm_keys": []} + specs = _make_comm_specs(op_name, direction, timing.comm_bytes) + return { + "duration": timing.compute_time, + "reduction": timing.comm_time, + "comm_keys": _register_specs(specs), + } - if parallelism_mode in (ParallelismMode.CONTEXT, ParallelismMode.TENSOR_CONTEXT_HYBRID): - for key in ("layernorm1", "qkv_proj", "attention", "output_proj", "layernorm2", "MLP"): - try: - spec = transformer_results[key] - except KeyError: - print(f"Key {key} not found in transformer_results") - print(transformer_results) - exit() - fwd_time = _require_compute(spec, "forward", key) - bwd_time = _require_compute(spec, "backward", key) - fwd_red = spec.get('forward_reduction', 0.0) - bwd_red = spec.get('backward_reduction', 0.0) - comm_bytes_fwd = spec.get('comm_size_forward', 0) - comm_bytes_bwd = spec.get('comm_size_backward', 0) - - entry = { - "name": key, - "forward": { - "duration": fwd_time, - "reduction": fwd_red, - "comm_keys": [], - }, - "backward": { - "duration": bwd_time, - "reduction": bwd_red, - "comm_keys": [], + def _make_comm_specs(op_name: str, direction: str, total_bytes: float) -> Tuple[CommSpec, ...]: + bytes_int = int(math.ceil(float(total_bytes or 0.0))) + if bytes_int <= 0: + return () + rule_spec = rules_by_mode.get(op_name) or rules_by_mode.get(COMM_RULE_DEFAULT_KEY) + if not rule_spec: + return () + rule = rule_spec.get(direction) + if not rule: + return () + kind = rule.get("kind") + scope = rule.get("participants") + if not kind or not scope: + return () + participants = participants_lookup.get(scope, 0) + if participants <= 1: + return () + interconnect = rule.get("interconnect", scope) + suffix = rule.get("suffix", kind) + name = f"{op_name}_{direction}_{suffix}" + return ( + CommSpec( + name=name, + kind=kind, + size_bytes=bytes_int, + participants=participants, + interconnect=interconnect, + extra={ + "scope": scope, + "direction": direction, + "op_name": op_name, + "parallelism_mode": parallelism_mode.value, }, - } + ), + ) - self._populate_transformer_comm_metadata( - entry=entry, - metadata=transformer_comm_metadata, - comm_bytes_fwd=comm_bytes_fwd, - comm_bytes_bwd=comm_bytes_bwd, + def _build_group_operation(name: str, members: Tuple[str, ...]) -> OperationTiming: + ops = tuple(transformer_timings[m] for m in members) + group = OperationGroup(name, ops) + forward_memory = self._combine_mem(*(op.forward.memory_accesses for op in ops)) + forward_flops = sum(op.forward.flops for op in ops) + forward_dir = DirectionTiming( + compute_time=group.forward_compute_time(), + comm_time=group.forward_comm_time(), + comm_bytes=group.forward_comm_bytes(), + flops=forward_flops, + memory_accesses=forward_memory, + ) + if include_transformer_backward: + backward_memory = self._combine_mem( + *(op.backward.memory_accesses for op in ops if op.backward is not None) ) - - transformer_operation_entries.append(entry) - - elif parallelism_mode in (ParallelismMode.TENSOR, ParallelismMode.TENSOR_SEQUENCE, ParallelismMode.SINGLE): - for key in ("layernorm1", "MHA", "layernorm2", "MLP"): - spec = transformer_results[key] - fwd_time = _require_compute(spec, "forward", key) - bwd_time = _require_compute(spec, "backward", key) - fwd_red = spec.get('forward_reduction', 0.0) - bwd_red = spec.get('backward_reduction', 0.0) - comm_bytes_fwd = spec.get('comm_size_forward', 0) - comm_bytes_bwd = spec.get('comm_size_backward', 0) - - entry = { - "name": key, - "forward": { - "duration": fwd_time, - "reduction": fwd_red, - "comm_keys": [], - }, - "backward": { - "duration": bwd_time, - "reduction": bwd_red, - "comm_keys": [], - }, - } - - self._populate_transformer_comm_metadata( - entry=entry, - metadata=transformer_comm_metadata, - comm_bytes_fwd=comm_bytes_fwd, - comm_bytes_bwd=comm_bytes_bwd, + backward_flops = sum((op.backward.flops if op.backward else 0.0) for op in ops) + backward_dir = DirectionTiming( + compute_time=group.backward_compute_time(), + comm_time=group.backward_comm_time(), + comm_bytes=group.backward_comm_bytes(), + flops=backward_flops, + memory_accesses=backward_memory, ) + else: + backward_dir = None + return OperationTiming(name, forward_dir, backward_dir) - transformer_operation_entries.append(entry) + def _resolve_operation(name: str) -> OperationTiming: + if name in group_members: + return _build_group_operation(name, group_members[name]) + return transformer_timings[name] + if parallelism_mode in (ParallelismMode.CONTEXT, ParallelismMode.TENSOR_CONTEXT_HYBRID): + op_names = ("layernorm1", "qkv_proj", "attention", "output_proj", "layernorm2", "MLP") + elif parallelism_mode in (ParallelismMode.TENSOR, ParallelismMode.TENSOR_SEQUENCE, ParallelismMode.SINGLE): + op_names = ("layernorm1", "MHA", "layernorm2", "MLP") else: raise ValueError(f"Unsupported parallelism mode: {parallelism_mode}") + + for key in op_names: + try: + op_timing = _resolve_operation(key) + except KeyError as exc: + raise KeyError(f"Missing transformer timing for operation '{key}'") from exc + + entry = { + "name": key, + "forward": _build_direction_entry(key, "forward", op_timing.forward), + "backward": _build_direction_entry(key, "backward", op_timing.backward if include_transformer_backward else None), + } + + if not include_transformer_backward: + # Ensure backward section is zeroed when excluded. + entry["backward"] = {"duration": 0.0, "reduction": 0.0, "comm_keys": []} + + transformer_operation_entries.append(entry) transformer_graph: Optional[Graph] = None transformer_forward_root: Optional[Any] = None transformer_backward_root: Optional[Any] = None @@ -2161,7 +2304,7 @@ def calc_time_llm(self): ) num_SMs = self.hw_config.tech_config.core.num_bundles - transformer_results, node_breakdown = self.compute_all_gemm_and_node_times( + transformer_timings, node_breakdown = self.compute_all_gemm_and_node_times( batch_size, vocab_size, hidden_dim, seq_len, num_heads, kv_heads, intermediate_size, num_SMs ) @@ -2237,7 +2380,7 @@ def calc_time_llm(self): interconnect_params, ) = self._prepare_execution_graphs( node_breakdown=node_breakdown, - transformer_results=transformer_results, + transformer_timings=transformer_timings, batch_size=batch_size, seq_len=seq_len, hidden_dim=hidden_dim, diff --git a/time_calculation_inf.py b/time_calculation_inf.py index a0b9c259..4998aab4 100644 --- a/time_calculation_inf.py +++ b/time_calculation_inf.py @@ -2,12 +2,18 @@ import math import os -from typing import Dict, List, Optional, Tuple -from time_calculation_LLM import LLMExecutionDispatcher, TimeCalculationLLM, GemmType +from typing import Dict, List, Optional, Tuple, Mapping, Set +from time_calculation_LLM import ( + LLMExecutionDispatcher, + TimeCalculationLLM, + GemmType, + COMMUNICATION_RULES, + COMM_RULE_DEFAULT_KEY, +) from simulate_inf import DecodeSample, InferenceConfig, InferenceEngine import LLM_util import json -from itertools import zip_longest # added for element-wise aggregation of memory access lists +from timing_model import CommSpec, DirectionTiming, OperationTiming, OperationGroup def convert_prefix(value: float) -> float: """Assign SI unit prefixes to numerical values.""" @@ -35,8 +41,9 @@ def _build_decode_transformer_results( batch_size: int, total_seq_len: int, gemm_shapes: Optional[Dict[str, Tuple[int, ...]]] = None, - ) -> Tuple[Dict[str, Dict[str, float]], Dict[str, float]]: - """Construct transformer spec + node breakdown for a single decode step.""" + ) -> Tuple[Dict[str, OperationTiming], Dict[str, float]]: + """Construct transformer timings + node breakdown for a single decode step.""" + head_dim = self.hidden_dim // self.num_heads token_bytes = LLM_util.kv_cache_token_bytes( @@ -46,7 +53,6 @@ def _build_decode_transformer_results( precision_bytes=self.precision.kv_cache, ) seq_degree = self._sequence_parallel_degree() - comm_kind = "all_reduce" if seq_degree == 1 else "reduce_scatter" intermediate_size = self.intermediate_size @@ -70,15 +76,88 @@ def _build_decode_transformer_results( gemm_ffn2 = gemm_shapes["ffn2"] # FlashAttention is not used during single-token (incremental) decoding. + + output_seq_len = 1 + + transformer_timings: Dict[str, OperationTiming] = {} + + def _make_forward( + op_name: str, + compute_time: float, + comm_time: float, + comm_bytes: float, + *, + flops: float = 0.0, + memory: Optional[Mapping[str, float]] = None, + ) -> DirectionTiming: + bytes_int = int(math.ceil(float(comm_bytes or 0.0))) + return DirectionTiming( + compute_time=compute_time, + comm_time=comm_time, + comm_bytes=bytes_int, + flops=flops, + memory_accesses=dict(memory) if memory else {}, + ) + + # QKV projection qkv_proj_time, qkv_proj_reduction, qkv_proj_size, qkv_proj_flops, qkv_proj_mem = self.parallelism_gemm_forward( gemm_qkv_proj, "decode_qkv_proj_f", gemm_type=GemmType.QKV ) + transformer_timings["qkv_proj"] = OperationTiming( + "qkv_proj", + forward=_make_forward( + "qkv_proj", + compute_time=qkv_proj_time, + comm_time=qkv_proj_reduction, + comm_bytes=qkv_proj_size, + flops=qkv_proj_flops, + memory=self._mem_levels(qkv_proj_mem), + ), + backward=None, + ) + + # Attention components attention_score_time, attention_score_reduction, attention_score_size, attention_score_flops, attention_score_mem = self.parallelism_gemm_forward( gemm_attention_score, "decode_attention_score_f", gemm_type=GemmType.ATTENTION_SCORE ) attention_output_time, attention_output_reduction, attention_output_size, attention_output_flops, attention_output_mem = self.parallelism_gemm_forward( gemm_attention_output, "decode_attention_output_f", gemm_type=GemmType.ATTENTION_OUTPUT ) + attention_scale_softmax_f = self.get_scale_softmax_f(gemm_attention_score) + + attention_reduction = attention_score_reduction + attention_output_reduction + attention_comm_bytes = attention_score_size + attention_output_size + attention_forward_compute = attention_score_time + attention_scale_softmax_f + attention_output_time + attention_forward_time = attention_forward_compute + attention_reduction + attention_flops = (attention_score_flops or 0.0) + (attention_output_flops or 0.0) + attention_mem = self._combine_mem(attention_score_mem, attention_output_mem) + + transformer_timings["attention"] = OperationTiming( + "attention", + forward=_make_forward( + "attention", + compute_time=attention_forward_compute, + comm_time=attention_reduction, + comm_bytes=attention_comm_bytes, + flops=attention_flops, + memory=attention_mem, + ), + backward=None, + ) + + attention_scale_softmax_op = OperationTiming( + "attention_scale_softmax", + forward=_make_forward( + "attention_scale_softmax", + compute_time=attention_scale_softmax_f, + comm_time=0.0, + comm_bytes=0.0, + ), + backward=None, + ) + transformer_timings["attention_scale_softmax"] = attention_scale_softmax_op + + # Output projection out_proj_time, _, out_proj_size, out_proj_flops, out_proj_mem = self.parallelism_gemm_forward( gemm_output_proj, "decode_output_projection_f", gemm_type=GemmType.OUT_PROJ ) @@ -87,7 +166,20 @@ def _build_decode_transformer_results( if out_proj_size else 0.0 ) + transformer_timings["output_proj"] = OperationTiming( + "output_proj", + forward=_make_forward( + "output_proj", + compute_time=out_proj_time, + comm_time=out_proj_reduction, + comm_bytes=out_proj_size, + flops=out_proj_flops, + memory=self._mem_levels(out_proj_mem), + ), + backward=None, + ) + # FFN layers ffn1_time, ffn1_reduction, ffn1_size, ffn1_flops, ffn1_mem = self.parallelism_gemm_forward( gemm_ffn1, "decode_ffn1_f", gemm_type=GemmType.FFN1 ) @@ -95,13 +187,45 @@ def _build_decode_transformer_results( gemm_ffn2, "decode_ffn2_f", gemm_type=GemmType.FFN2 ) ffn2_reduction = ( - self.get_tensor_reduction_time(ffn2_size, comm_kind, "decode_ffn2") - if ffn2_size - else 0.0 + self.get_tensor_reduction_time(ffn2_size, comm_kind, "decode_ffn2") if ffn2_size else 0.0 + ) + transformer_timings["ffn1"] = OperationTiming( + "ffn1", + forward=_make_forward( + "ffn1", + compute_time=ffn1_time, + comm_time=ffn1_reduction, + comm_bytes=ffn1_size, + flops=ffn1_flops, + memory=self._mem_levels(ffn1_mem), + ), + backward=None, + ) + transformer_timings["ffn2"] = OperationTiming( + "ffn2", + forward=_make_forward( + "ffn2", + compute_time=ffn2_time, + comm_time=ffn2_reduction, + comm_bytes=ffn2_size, + flops=ffn2_flops, + memory=self._mem_levels(ffn2_mem), + ), + backward=None, ) - output_seq_len = 1 + # GELU/SwiGLU activation + if self.model_type == "llama": + act_f = self.get_swiglu_f(gemm_ffn1) + else: + act_f = self.get_gelu_f(gemm_ffn1) + transformer_timings["gelu"] = OperationTiming( + "gelu", + forward=_make_forward("gelu", compute_time=act_f, comm_time=0.0, comm_bytes=0.0), + backward=None, + ) + # Layer norms output_proj_shape = ( batch_size, output_seq_len, @@ -112,8 +236,17 @@ def _build_decode_transformer_results( layernorm1_f, layernorm1_reduction, layernorm1_bytes = self.get_layernorm_f( batch=batch_size, seq_len=output_seq_len, d_model=self.hidden_dim ) + transformer_timings["layernorm1"] = OperationTiming( + "layernorm1", + forward=_make_forward( + "layernorm1", + compute_time=layernorm1_f + residual1_f, + comm_time=layernorm1_reduction, + comm_bytes=layernorm1_bytes, + ), + backward=None, + ) - intermediate_size = self.intermediate_size ffn2_shape = ( batch_size, output_seq_len, @@ -124,6 +257,16 @@ def _build_decode_transformer_results( layernorm2_f, layernorm2_reduction, layernorm2_bytes = self.get_layernorm_f( batch=batch_size, seq_len=output_seq_len, d_model=self.hidden_dim ) + transformer_timings["layernorm2"] = OperationTiming( + "layernorm2", + forward=_make_forward( + "layernorm2", + compute_time=layernorm2_f + residual2_f, + comm_time=layernorm2_reduction, + comm_bytes=layernorm2_bytes, + ), + backward=None, + ) linear_shape = ( batch_size, @@ -131,33 +274,46 @@ def _build_decode_transformer_results( self.hidden_dim, self.vocab_size, ) - linear_softmax_f, _ = self.get_linear_softmax_f(linear_shape) + linear_softmax_f, linear_softmax_mem = self.get_linear_softmax_f(linear_shape) + transformer_timings["linear_softmax"] = OperationTiming( + "linear_softmax", + forward=_make_forward( + "linear_softmax", + compute_time=linear_softmax_f, + comm_time=0.0, + comm_bytes=0.0, + memory=self._mem_levels(linear_softmax_mem), + ), + backward=None, + ) - if self.model_type == "llama": - act_f = self.get_swiglu_f(gemm_ffn1) - else: - act_f = self.get_gelu_f(gemm_ffn1) + # Aggregate operations + mha_group = OperationGroup( + "MHA", + operations=( + transformer_timings["qkv_proj"], + transformer_timings["attention"], + transformer_timings["output_proj"], + ), + ) + mlp_group = OperationGroup( + "MLP", + operations=( + transformer_timings["ffn1"], + transformer_timings["gelu"], + transformer_timings["ffn2"], + ), + ) - attention_scale_softmax_f = self.get_scale_softmax_f(gemm_attention_score) - # Preserve the same per-op structure and collective accounting used in the training path. - # Most decode ops (qkv projection, attention score/output, ffn1) still have zero-sized - # reductions under tensor parallelism - we pass through the values returned by the shared - # helpers so the context/tensor-context graph builder sees the exact same semantics. + # Match exact floating-point operation order from original code qkv_proj_forward = qkv_proj_time + qkv_proj_reduction - attention_reduction = attention_score_reduction + attention_output_reduction - attention_comm_bytes = attention_score_size + attention_output_size - attention_forward = ( - attention_score_time + attention_scale_softmax_f + attention_output_time + attention_reduction - ) - attention_forward_compute = attention_score_time + attention_scale_softmax_f + attention_output_time + attention_forward = attention_score_time + attention_scale_softmax_f + attention_output_time + attention_reduction out_proj_forward = out_proj_time + out_proj_reduction mha_forward = qkv_proj_forward + attention_forward + out_proj_forward - mha_forward_compute = qkv_proj_time + attention_forward_compute + out_proj_time ffn1_forward = ffn1_time + ffn1_reduction ffn2_forward = ffn2_time + ffn2_reduction mlp_forward = ffn1_forward + act_f + ffn2_forward - mlp_forward_compute = ffn1_time + act_f + ffn2_time layernorm1_forward = residual1_f + layernorm1_f layernorm2_forward = residual2_f + layernorm2_f @@ -171,147 +327,17 @@ def _build_decode_transformer_results( + layernorm2_reduction ) - # Aggregate attention flops and memory accesses from score + output - attention_flops = (attention_score_flops or 0) + (attention_output_flops or 0) - attention_mem = [ - (a or 0) + (b or 0) - for a, b in zip_longest(attention_score_mem, attention_output_mem, fillvalue=0) - ] - - transformer_results = { - "qkv_proj": { - "forward": qkv_proj_time + qkv_proj_reduction, - "backward": 0.0, - "forward_gemm": qkv_proj_time, - "forward_reduction": qkv_proj_reduction, - "backward_gemm": 0.0, - "backward_reduction": 0.0, - "comm_size_forward": qkv_proj_size, - "comm_size_backward": 0, - "flops": qkv_proj_flops, - "memory_accesses": self._mem_levels(qkv_proj_mem), # changed to per-level dict - }, - "attention": { - "forward": attention_forward, - "backward": 0.0, - "forward_gemm": attention_score_time + attention_output_time + attention_scale_softmax_f, - "forward_reduction": attention_reduction, - "backward_gemm": 0.0, - "backward_reduction": 0.0, - "comm_size_forward": attention_comm_bytes, - "comm_size_backward": 0, - "flops": attention_flops, - "memory_accesses": self._mem_levels(attention_mem), # changed to per-level dict - }, - "output_proj": { - "forward": out_proj_time + out_proj_reduction, - "backward": 0.0, - "forward_gemm": out_proj_time, - "forward_reduction": out_proj_reduction, - "backward_gemm": 0.0, - "backward_reduction": 0.0, - "comm_size_forward": out_proj_size, - "comm_size_backward": 0, - "flops": out_proj_flops, - "memory_accesses": self._mem_levels(out_proj_mem), # changed to per-level dict - }, - "MHA": { - "forward": mha_forward, - "backward": 0.0, - "forward_compute": mha_forward_compute, - "backward_compute": 0.0, - "forward_reduction": qkv_proj_reduction + attention_reduction + out_proj_reduction, - "backward_reduction": 0.0, - "comm_size_forward": qkv_proj_size + attention_comm_bytes + out_proj_size, - "comm_size_backward": 0, - "flops": qkv_proj_flops + attention_flops + out_proj_flops, - "memory_accesses": self._combine_mem( - qkv_proj_mem, - attention_mem, - out_proj_mem, - ), - }, - "ffn1": { - "forward": ffn1_forward, - "backward": 0.0, - "forward_gemm": ffn1_time, - "forward_reduction": ffn1_reduction, - "backward_gemm": 0.0, - "backward_reduction": 0.0, - "comm_size_forward": ffn1_size, - "comm_size_backward": 0, - "flops": ffn1_flops, - "memory_accesses": self._mem_levels(ffn1_mem), # changed to per-level dict - }, - "ffn2": { - "forward": ffn2_forward, - "backward": 0.0, - "forward_gemm": ffn2_time, - "forward_reduction": ffn2_reduction, - "backward_gemm": 0.0, - "backward_reduction": 0.0, - "comm_size_forward": ffn2_size, - "comm_size_backward": 0, - "flops": ffn2_flops, - "memory_accesses": self._mem_levels(ffn2_mem), # changed to per-level dict - }, - "MLP": { - "forward": mlp_forward, - "backward": 0.0, - "forward_compute": mlp_forward_compute, - "backward_compute": 0.0, - "forward_reduction": ffn1_reduction + ffn2_reduction, - "backward_reduction": 0.0, - "comm_size_forward": ffn1_size + ffn2_size, - "comm_size_backward": 0, - "flops": ffn1_flops + ffn2_flops, - "memory_accesses": self._combine_mem( - ffn1_mem, - ffn2_mem, - ), - }, - "gelu": { - "forward": act_f, - "backward": 0.0, - "forward_compute": act_f, - "backward_compute": 0.0, - "forward_reduction": 0.0, - "backward_reduction": 0.0, - "comm_size_forward": 0, - "comm_size_backward": 0, - }, - "layernorm1": { - "forward": layernorm1_forward, - "forward_compute": layernorm1_f, - "backward": 0.0, - "backward_compute": 0.0, - "forward_reduction": layernorm1_reduction, - "backward_reduction": 0.0, - "comm_size_forward": layernorm1_bytes, - "comm_size_backward": 0, - }, - "layernorm2": { - "forward": layernorm2_forward, - "forward_compute": layernorm2_f, - "backward": 0.0, - "backward_compute": 0.0, - "forward_reduction": layernorm2_reduction, - "backward_reduction": 0.0, - "comm_size_forward": layernorm2_bytes, - "comm_size_backward": 0, - }, - } - node_breakdown = { "transformer_time_f": transformer_forward, "transformer_time_b": 0.0, - "linear_softmax_f": linear_softmax_f, + "linear_softmax_f": transformer_timings["linear_softmax"].total_forward_time(), "linear_softmax_b": 0.0, "embedding_f": 0.0, "embedding_b": 0.0, } - return transformer_results, node_breakdown + return transformer_timings, node_breakdown + def prepare_decode_graphs( self, @@ -321,21 +347,23 @@ def prepare_decode_graphs( gemm_shapes: Optional[Dict[str, Tuple[int, ...]]] = None, ): intermediate_size = self.intermediate_size - transformer_results, node_breakdown = self._build_decode_transformer_results( + transformer_timings, node_breakdown = self._build_decode_transformer_results( batch_size=batch_size, total_seq_len=total_seq_len, gemm_shapes=gemm_shapes, ) output_act_bytes = gemm_shapes["qkv_proj"][0] * gemm_shapes["qkv_proj"][1] * self.precision_bytes - energy = self.calc_energy(transformer_results, output_act_bytes) + energy = self.calc_energy(transformer_timings, output_act_bytes) if self._generate_graphs: results_path = os.path.join(self.output_dir, "decode_transformer_results.txt") with open(results_path, "w", encoding="utf-8") as results_file: json.dump( { - "transformer_results": transformer_results, + "transformer_results": { + name: timing.to_dict() for name, timing in transformer_timings.items() + }, "node_breakdown": node_breakdown, }, results_file, @@ -345,7 +373,7 @@ def prepare_decode_graphs( return self._prepare_execution_graphs( node_breakdown=node_breakdown, - transformer_results=transformer_results, + transformer_timings=transformer_timings, batch_size=batch_size, seq_len=1, hidden_dim=self.hidden_dim, @@ -356,18 +384,36 @@ def prepare_decode_graphs( gemm_shapes=gemm_shapes, ), energy - def calc_energy(self, transformer_results: Dict[str, Dict[str, float]], cross_layer_comm) -> float: + def calc_energy(self, transformer_timings: Dict[str, OperationTiming], cross_layer_comm) -> float: """ Calculate energy consumption based on transformer results. """ - total_flops = 0 - total_hbm_bytes = 0 - inter_comm_bytes = 0 # data parallelism? - for op_name, op_results in transformer_results.items(): - if op_name in ['MHA', 'MLP', 'layernorm1', 'layernorm2', 'embedding', 'linear_softmax']: - total_flops += op_results.get("flops", 0) - total_hbm_bytes += op_results.get("memory_accesses", {}).get("L3", 0) - inter_comm_bytes += op_results.get("comm_size_forward", 0) + total_flops = 0.0 + total_hbm_bytes = 0.0 + inter_comm_bytes = 0.0 # data parallelism? + + aggregate_groups = { + "MHA": ("qkv_proj", "attention", "output_proj"), + "MLP": ("ffn1", "gelu", "ffn2"), + } + solo_ops = ("layernorm1", "layernorm2", "embedding", "linear_softmax") + + for members in aggregate_groups.values(): + for name in members: + op = transformer_timings.get(name) + if op is None: + continue + total_flops += op.forward.flops + total_hbm_bytes += op.forward.memory_accesses.get("L3", 0.0) + inter_comm_bytes += op.forward.comm_bytes + + for name in solo_ops: + op = transformer_timings.get(name) + if op is None: + continue + total_flops += op.forward.flops + total_hbm_bytes += op.forward.memory_accesses.get("L3", 0.0) + inter_comm_bytes += op.forward.comm_bytes total_comm_bytes = self.num_layers * inter_comm_bytes + (self.lp - 1) * cross_layer_comm @@ -400,7 +446,7 @@ def calc_time(self) -> float: self.readjust_type() num_SMs = self.hw_config.tech_config.core.num_bundles - transformer_results, node_breakdown = self.compute_all_gemm_and_node_times( + transformer_timings, node_breakdown = self.compute_all_gemm_and_node_times( batch_size, vocab_size, hidden_dim, @@ -413,7 +459,7 @@ def calc_time(self) -> float: output_act_bytes = batch_size * self.seq_len * hidden_dim * self.precision_bytes - total_energy = self.calc_energy(transformer_results, output_act_bytes) + total_energy = self.calc_energy(transformer_timings, output_act_bytes) head_dim = hidden_dim // num_heads token_bytes = LLM_util.kv_cache_token_bytes( @@ -432,7 +478,7 @@ def calc_time(self) -> float: interconnect_params, ) = self._prepare_execution_graphs( node_breakdown=node_breakdown, - transformer_results=transformer_results, + transformer_timings=transformer_timings, batch_size=batch_size, seq_len=prefill_len, hidden_dim=hidden_dim, diff --git a/timing_model.py b/timing_model.py new file mode 100644 index 00000000..d9bb35f5 --- /dev/null +++ b/timing_model.py @@ -0,0 +1,233 @@ +from dataclasses import dataclass, field +from typing import Any, Dict, Mapping, Optional, Sequence, Tuple + + +@dataclass(frozen=True) +class CommSpec: + """ + Metadata describing a single communication collective. + + Attributes: + name: Unique identifier used when wiring graph edges. + kind: Collective type (e.g. all_reduce, reduce_scatter). + size_bytes: Total bytes transferred (int >= 0). + participants: Number of ranks involved (int >= 1). + interconnect: Logical interconnect label (tp/dp/cp/lp/etc.). + extra: Free-form metadata (e.g. debug labels). + """ + + name: str + kind: str + size_bytes: int + participants: int + interconnect: str + extra: Mapping[str, Any] = field(default_factory=dict) + + def __post_init__(self) -> None: + if not self.name: + raise ValueError("CommSpec.name must be non-empty") + if not self.kind: + raise ValueError(f"CommSpec '{self.name}' missing kind") + if self.size_bytes < 0: + raise ValueError(f"CommSpec '{self.name}' has negative size_bytes ({self.size_bytes})") + if self.participants <= 0: + raise ValueError(f"CommSpec '{self.name}' must have participants >= 1 (got {self.participants})") + if not self.interconnect: + raise ValueError(f"CommSpec '{self.name}' missing interconnect label") + + +@dataclass(frozen=True) +class DirectionTiming: + """ + Timing information for a single direction (forward/backward) of an operation. + + Attributes: + compute_time: Pure computation time (seconds). + comm_time: Time attributed to communication collectives (seconds). + comm_bytes: Raw communication bytes (int >= 0). CommSpec instances are created + later during graph building based on COMMUNICATION_RULES. + flops: Floating-point operation count (optional). + memory_accesses: Tuple describing memory accesses per level (optional). + notes: Misc metadata for debugging/reporting. + """ + + compute_time: float + comm_time: float = 0.0 + comm_bytes: int = 0 + flops: float = 0.0 + memory_accesses: Mapping[str, float] = field(default_factory=dict) + notes: Mapping[str, Any] = field(default_factory=dict) + + def __post_init__(self) -> None: + if self.compute_time < 0: + raise ValueError(f"DirectionTiming compute_time must be >= 0 (got {self.compute_time})") + if self.comm_time < 0: + raise ValueError(f"DirectionTiming comm_time must be >= 0 (got {self.comm_time})") + if self.comm_bytes < 0: + raise ValueError(f"DirectionTiming comm_bytes must be >= 0 (got {self.comm_bytes})") + if self.flops < 0: + raise ValueError(f"DirectionTiming flops must be >= 0 (got {self.flops})") + for level_bytes in self.memory_accesses.values(): + if level_bytes < 0: + raise ValueError("DirectionTiming memory_accesses entries must be >= 0") + + def total_time(self) -> float: + """Return compute + comm time.""" + return self.compute_time + self.comm_time + + def total_comm_bytes(self) -> int: + """Return raw communication bytes.""" + return self.comm_bytes + + def has_comm(self) -> bool: + return self.comm_bytes > 0 + + def to_dict(self) -> Dict[str, Any]: + """Serialize to a dict for backwards compatibility (e.g. logging).""" + return { + "compute_time": self.compute_time, + "comm_time": self.comm_time, + "comm_bytes": self.comm_bytes, + "flops": self.flops, + "memory_accesses": dict(self.memory_accesses), + "notes": dict(self.notes), + } + + +@dataclass(frozen=True) +class OperationTiming: + """ + Timing details for an operation (e.g., QKV GEMM, MLP). + + Attributes: + name: Logical operation name. + forward: Forward-direction timing (required). + backward: Backward-direction timing (optional). + metadata: Additional operation-level metadata (e.g. layer index). + """ + + name: str + forward: DirectionTiming + backward: Optional[DirectionTiming] = None + metadata: Mapping[str, Any] = field(default_factory=dict) + + def __post_init__(self) -> None: + if not self.name: + raise ValueError("OperationTiming.name must be non-empty") + if not isinstance(self.forward, DirectionTiming): + raise TypeError("OperationTiming.forward must be a DirectionTiming instance") + if self.backward is not None and not isinstance(self.backward, DirectionTiming): + raise TypeError("OperationTiming.backward must be a DirectionTiming instance or None") + + def has_backward(self) -> bool: + return self.backward is not None + + def total_forward_time(self) -> float: + return self.forward.total_time() + + def total_backward_time(self) -> float: + if self.backward is None: + raise RuntimeError(f"Operation '{self.name}' does not have backward timing") + return self.backward.total_time() + + def total_time(self) -> float: + total = self.forward.total_time() + if self.backward is not None: + total += self.backward.total_time() + return total + + def to_dict(self) -> Dict[str, Any]: + """Serialize to dict for reporting or artifact dumps.""" + data: Dict[str, Any] = { + "name": self.name, + "forward": self.forward.to_dict(), + } + if self.backward is not None: + data["backward"] = self.backward.to_dict() + if self.metadata: + data["metadata"] = dict(self.metadata) + return data + + def validate(self, require_backward: bool = False) -> None: + if require_backward and self.backward is None: + raise ValueError(f"Operation '{self.name}' is missing backward timing") + + +@dataclass(frozen=True) +class OperationGroup: + """ + Aggregates a set of OperationTiming instances (e.g. the GEMMs that form MHA). + Provides convenience accessors for summary statistics. + """ + + name: str + operations: Tuple[OperationTiming, ...] + metadata: Mapping[str, Any] = field(default_factory=dict) + + def __post_init__(self) -> None: + if not self.name: + raise ValueError("OperationGroup.name must be non-empty") + if not self.operations: + raise ValueError(f"OperationGroup '{self.name}' must contain at least one operation") + for op in self.operations: + if not isinstance(op, OperationTiming): + raise TypeError("OperationGroup.operations must contain OperationTiming instances") + + def forward_compute_time(self) -> float: + return sum(op.forward.compute_time for op in self.operations) + + def forward_comm_time(self) -> float: + return sum(op.forward.comm_time for op in self.operations) + + def forward_comm_bytes(self) -> int: + return sum(op.forward.comm_bytes for op in self.operations) + + def forward_total_time(self) -> float: + return sum(op.total_forward_time() for op in self.operations) + + def backward_compute_time(self) -> float: + total = 0.0 + for op in self.operations: + if op.backward is None: + raise RuntimeError( + f"OperationGroup '{self.name}' member '{op.name}' is missing backward timing" + ) + total += op.backward.compute_time + return total + + def backward_comm_time(self) -> float: + total = 0.0 + for op in self.operations: + if op.backward is None: + raise RuntimeError( + f"OperationGroup '{self.name}' member '{op.name}' is missing backward timing" + ) + total += op.backward.comm_time + return total + + def backward_comm_bytes(self) -> int: + total = 0 + for op in self.operations: + if op.backward is None: + raise RuntimeError( + f"OperationGroup '{self.name}' member '{op.name}' is missing backward timing" + ) + total += op.backward.comm_bytes + return total + + def backward_total_time(self) -> float: + total = 0.0 + for op in self.operations: + if op.backward is None: + raise RuntimeError( + f"OperationGroup '{self.name}' member '{op.name}' is missing backward timing" + ) + total += op.total_backward_time() + return total + + def to_dict(self) -> Dict[str, Any]: + return { + "name": self.name, + "operations": [op.to_dict() for op in self.operations], + "metadata": dict(self.metadata), + } From 4f0a3c4c1f59b7399be087e300486e087e2dd108 Mon Sep 17 00:00:00 2001 From: gkarfakis19 Date: Sun, 2 Nov 2025 05:14:36 -0800 Subject: [PATCH 07/77] first stage config update. Astrasim integration missing, everything else supported. --- README.md | 8 +- astrasim_lib/config_generation.py | 41 +- config.py | 505 ++++++++++++------ configs/hardware-config/H100_SXM5_80GB.yaml | 56 +- configs/hardware-config/a100_80GB.yaml | 102 ++-- .../a100_80GB_no_parallelism.yaml | 49 +- configs/hardware-config/v100.yaml | 50 +- .../waferscale_20v100_80hbm.yaml | 46 +- configs/model-config/LLM.yaml | 2 +- configs/model-config/LLM_inf.yaml | 8 +- .../H100_SXM5_80GB.yaml | 6 +- .../old-detailed-hw-configs/a100_80GB.yaml | 6 +- configs/old-detailed-hw-configs/v100.yaml | 6 +- .../waferscale_20v100_80hbm.yaml | 6 +- .../old_lstm_hw_configs/H100_SXM5_80GB.yaml | 150 ++++++ configs/old_lstm_hw_configs/a100_80GB.yaml | 143 +++++ configs/old_lstm_hw_configs/v100.yaml | 145 +++++ GD_search.py => deprecated/GD_search.py | 9 +- .../deviceMapping.py | 0 energy.py => deprecated/energy.py | 0 main.py => deprecated/main.py | 0 topology.py => deprecated/topology.py | 0 .../topology_hack.py | 0 examples/llm_astra_inference.sh | 2 +- .../llm_astra_inference_debug_graphviz.sh | 2 +- hw_component.py | 214 ++------ pyproject.toml | 2 - requirements.txt | 4 +- run_perf.py | 21 +- time_calculation.py | 134 ++--- time_calculation_LLM.py | 21 +- time_calculation_inf.py | 3 +- uv.lock | 39 -- 33 files changed, 1032 insertions(+), 748 deletions(-) create mode 100644 configs/old_lstm_hw_configs/H100_SXM5_80GB.yaml create mode 100644 configs/old_lstm_hw_configs/a100_80GB.yaml create mode 100644 configs/old_lstm_hw_configs/v100.yaml rename GD_search.py => deprecated/GD_search.py (99%) rename deviceMapping.py => deprecated/deviceMapping.py (100%) rename energy.py => deprecated/energy.py (100%) rename main.py => deprecated/main.py (100%) rename topology.py => deprecated/topology.py (100%) rename topology_hack.py => deprecated/topology_hack.py (100%) diff --git a/README.md b/README.md index 80addecd..978018e6 100644 --- a/README.md +++ b/README.md @@ -173,7 +173,7 @@ python configs/model-config/hf_to_config.py Qwen/Qwen2.5-3B --run-type inference 2. Open the provided A100 hardware file and set tensor parallel degree to 1: -Edit `configs/hardware-config/a100_80GB.yaml` so that `system_hierarchy.num_devices_per_node: 1` and `scheduling_param.tp: 1`. This models one GPU without tensor parallelism. +Edit `configs/hardware-config/a100_80GB.yaml` so that `parallelism.tp: 1`. This models one GPU without tensor parallelism. Run the inference estimation using the example configs: @@ -183,7 +183,7 @@ python run_perf.py --hardware_config configs/hardware-config/a100_80GB.yaml --mo 3. To see the effects of tensor-parallelism, modify the hardware file, switching to two devices per node, and tp = 2. -Edit `configs/hardware-config/a100_80GB.yaml` so that `system_hierarchy.num_devices_per_node: 2` and `scheduling_param.tp: 2`. This models two GPU-inference with tensor parallelism degree of 2. +Edit `configs/hardware-config/a100_80GB.yaml` so that `parallelism.tp: 2`. This models two GPU-inference with tensor parallelism degree of 2. Re-run the same inference command with the updated hardware config. @@ -249,6 +249,10 @@ Comparing the two runs will show how increasing tensor parallelism changes the p - **Supported:** Very rudimentary, non validated model. Does not consider static power. Inference only for now (Training is WIP). Energy due to prefill and decode is reported separately. - **Work in progress:** Training support, and heavy refinement of current implementation. +### Mixed Precision +- **Supported:** Mixed precision is supported for all supported parallelism and model types. +- **Work in progress:** KV Cache precision cannot yet be set (config value is ignored). + ### Validation - **Status:** Validation scripts are under development. Limited inference and training validation has been performed. - **Supported:** For inference, we validate against Koyeb's single NVIDIA A100 data on LLama3.1-8B. The relevant script can be found in `validation_scripts/koyeb.py`. Below is the resulting plot, comparing DeepFlow's predicted decode throughput with Koyeb's reported decode throughput, available at https://www.koyeb.com/docs/hardware/gpu-benchmarks. Average absolute error is ~6%. diff --git a/astrasim_lib/config_generation.py b/astrasim_lib/config_generation.py index 69a5e20f..7c65321c 100644 --- a/astrasim_lib/config_generation.py +++ b/astrasim_lib/config_generation.py @@ -78,27 +78,26 @@ def compute_intra_inter_ib_ll_from_hw(hw_obj) -> Tuple[Tuple[float, float], Tupl def derive_topology_from_hw(hw_obj) -> str: """Map DeepFlow network topology enums to AstraSim names.""" - try: - topo = getattr(hw_obj.network_topology.intra, "topology", None) - topo_str = topo.lower() - if topo_str in ("fc", "fullyconnected", "fully_connected", "fully-connected"): - return "FullyConnected" - if topo_str in ("ring",): - return "Ring" - if topo_str in ("switch",): - return "Switch" - if topo_str in ("torus2d",): - return "Torus2D" - if topo_str in ("mesh",): - return "Mesh" - if topo_str in ("hypercube",): - return "HyperCube" - if topo_str in ("mesh2d",): - return "Mesh2D" - if topo_str in ("kingmesh2d",): - return "KingMesh2D" - except Exception: # pragma: no cover - defensive - pass + layout = getattr(hw_obj, "network_layout", None) + primary = layout.primary_dimension() if layout else None + topo = (primary.topology_type if primary else "ring") or "ring" + topo_str = str(topo).lower() + if topo_str in ("fc", "fullyconnected", "fully_connected", "fully-connected"): + return "FullyConnected" + if topo_str in ("ring",): + return "Ring" + if topo_str in ("switch",): + return "Switch" + if topo_str in ("torus2d", "torus"): + return "Torus2D" + if topo_str in ("mesh",): + return "Mesh" + if topo_str in ("hypercube",): + return "HyperCube" + if topo_str in ("mesh2d", "mesh-2d"): + return "Mesh2D" + if topo_str in ("kingmesh2d", "king-mesh2d"): + return "KingMesh2D" return "FullyConnected" diff --git a/config.py b/config.py index 716d8064..71a684b2 100644 --- a/config.py +++ b/config.py @@ -1,10 +1,10 @@ from dataclasses import dataclass, field -import ruamel as _ruamel -import ruamel.yaml as _yaml -from ruamel.yaml import YAMLError as _YAMLError import math from collections import namedtuple as _namedtuple -from typing import Optional +from typing import Dict, List, Optional, Sequence, Tuple + +import yaml as _yaml +from yaml import YAMLError as _YAMLError _PRECISION_DTYPE_BYTES = { @@ -234,50 +234,6 @@ def from_dict(cls, sram_config_dict): ) -@dataclass -class SubNetworkConfig: - latency: float - nominal_freq: float - nominal_voltage: float - nominal_energy_per_link: float - nominal_area_per_link: float - threshold_voltage: float - margin_voltage: float - num_links_per_mm: int - util: float - operating_frequency: float = None # Optional: overrides voltage-scaled frequency - bandwidth: float = None # Optional: master parameter that directly sets throughput - - @classmethod - def from_dict(cls, config_dict): - return cls( - latency=config_dict["latency"], - nominal_freq=config_dict.get("nominal_frequency", 0.1), - nominal_voltage=config_dict.get("nominal_voltage", 0.1), - nominal_energy_per_link=config_dict.get("nominal_energy_per_link", 0.1), - nominal_area_per_link=config_dict.get("nominal_area_per_link", 0.1), - threshold_voltage=config_dict.get("threshold_voltage", 0.1), - margin_voltage=config_dict.get("margin_voltage", 0.1), - num_links_per_mm=config_dict.get("num_links_per_mm", 1), - util=config_dict["util"], - operating_frequency=config_dict.get("operating_frequency", None), - bandwidth=parse_bandwidth_string(config_dict.get("bandwidth", None)), - ) - - -@dataclass -class NetworkConfig: - intra_node: SubNetworkConfig - inter_node: SubNetworkConfig - - @classmethod - def from_dict(cls, d): - return cls( - intra_node=SubNetworkConfig.from_dict(d["intra_node"]), - inter_node=SubNetworkConfig.from_dict(d["inter_node"]), - ) - - @dataclass class TechConfig: core: CoreConfig @@ -285,7 +241,6 @@ class TechConfig: SRAML2: SRAMConfig SRAML1: SRAMConfig SRAMR: SRAMConfig - network: NetworkConfig @classmethod def from_dict(cls, tech_config_dict): @@ -295,7 +250,6 @@ def from_dict(cls, tech_config_dict): SRAML2=SRAMConfig.from_dict(tech_config_dict["SRAM-L2"]), SRAML1=SRAMConfig.from_dict(tech_config_dict["SRAM-L1"]), SRAMR=SRAMConfig.from_dict(tech_config_dict["SRAM-R"]), - network=NetworkConfig.from_dict(tech_config_dict["network"]), ) @@ -388,69 +342,267 @@ def from_dict(cls, network_power_config_dict): ) -@dataclass -class SystemHierarchyConfig: - num_nodes_per_wafer: int - num_wafers: int - num_workers: int - inter_derate: float - intra_derate: float - kp1_inter: float - kp2_inter: float - dp_inter: float - lp_inter: float - tp_inter: float - par2cross: dict +@dataclass(frozen=True) +class NetworkDimensionLayout: + id: str + label: str + size: int + topology_type: str + bandwidth: float + util: float + latency: float + collective_override: Dict[str, str] = field(default_factory=dict) + parallelisms: Tuple[str, ...] = field(default_factory=tuple) + energy_per_bit: float = 0.0 @classmethod - def from_dict(cls, system_config_dict): + def from_raw( + cls, + raw: dict, + *, + parallelism_params: Dict[str, object], + index: int, + ) -> "NetworkDimensionLayout": + if not isinstance(raw, dict): + raise TypeError("each network dimension must be a mapping") + + raw_id = raw.get("id") + dim_id = str(raw_id) if raw_id is not None else f"dim{index}" + label = str(raw.get("label", dim_id)) + + if "size" not in raw: + raise ValueError(f"network dimension '{label}' is missing required field 'size'") + try: + size = int(raw["size"]) + except (TypeError, ValueError) as exc: + raise ValueError(f"network dimension '{label}' size must be an integer") from exc + if size < 1: + raise ValueError(f"network dimension '{label}' size must be >= 1") + + topo_dict = raw.get("topology") + if not isinstance(topo_dict, dict): + raise ValueError(f"network dimension '{label}' requires a 'topology' mapping") + if "type" not in topo_dict: + raise ValueError(f"network dimension '{label}' topology missing required 'type'") + topo_type = str(topo_dict["type"]) + + if "bandwidth" not in topo_dict: + raise ValueError(f"network dimension '{label}' topology missing required 'bandwidth'") + bandwidth = parse_bandwidth_string(topo_dict["bandwidth"]) + + try: + util = float(topo_dict.get("util", 1.0)) + except (TypeError, ValueError) as exc: + raise ValueError( + f"network dimension '{label}' topology util must be numeric" + ) from exc + if util <= 0: + raise ValueError(f"network dimension '{label}' topology util must be > 0") + + if "energy_per_bit" not in topo_dict: + raise ValueError( + f"network dimension '{label}' topology missing required 'energy_per_bit'" + ) + try: + energy_per_bit = float(topo_dict["energy_per_bit"]) + except (TypeError, ValueError) as exc: + raise ValueError( + f"network dimension '{label}' energy_per_bit must be numeric" + ) from exc + if energy_per_bit < 0: + raise ValueError( + f"network dimension '{label}' energy_per_bit must be >= 0" + ) + + if "latency" not in topo_dict: + raise ValueError(f"network dimension '{label}' topology missing required 'latency'") + try: + latency = float(topo_dict["latency"]) + except (TypeError, ValueError) as exc: + raise ValueError(f"network dimension '{label}' latency must be numeric") from exc + + collectives_raw = raw.get("collective_override") + if collectives_raw is None and "collectives" in raw: + collectives_raw = raw.get("collectives") + if not collectives_raw: + collectives_raw = {} + if not isinstance(collectives_raw, dict): + raise ValueError(f"network dimension '{label}' collective_override must be a mapping if provided") + collective_override = {str(k): str(v) for k, v in collectives_raw.items()} + + parallelisms_raw = raw.get("parallelisms", []) + if parallelisms_raw is None: + parallelisms_raw = [] + if not isinstance(parallelisms_raw, Sequence) or isinstance(parallelisms_raw, (str, bytes)): + raise ValueError( + f"network dimension '{label}' parallelisms must be a sequence of names" + ) + + normalized_parallelisms: List[str] = [] + alias_map: Dict[str, str] = {} + for entry in parallelisms_raw: + name = str(entry).strip() + if not name: + raise ValueError(f"network dimension '{label}' has an empty parallelism name") + normalized = name.lower() + normalized_parallelisms.append(normalized) + alias_map[normalized] = name + + _validate_dimension_parallelisms( + dimension_label=label, + dimension_size=size, + normalized_names=tuple(normalized_parallelisms), + alias_map=alias_map, + parallelism_params=parallelism_params, + ) + return cls( - num_nodes_per_wafer=system_config_dict["num_devices_per_node"], - num_wafers=system_config_dict["num_nodes"], - num_workers=int( - system_config_dict["num_nodes"] - * system_config_dict["num_devices_per_node"] - ), - inter_derate=system_config_dict.get("inter_derate", 1.0), - intra_derate=system_config_dict.get("intra_derate", 1.0), - kp1_inter=system_config_dict["kp1_inter"], - kp2_inter=system_config_dict["kp2_inter"], - dp_inter=system_config_dict["dp_inter"], - lp_inter=system_config_dict["lp_inter"], - tp_inter=system_config_dict["tp_inter"], - par2cross={ - "kp1": system_config_dict["kp1_inter"], - "kp2": system_config_dict["kp2_inter"], - "dp": system_config_dict["dp_inter"], - "lp": system_config_dict["lp_inter"], - "tp": system_config_dict["tp_inter"], - }, + id=dim_id, + label=label, + size=size, + topology_type=topo_type, + bandwidth=bandwidth, + util=util, + latency=latency, + collective_override=collective_override, + parallelisms=tuple(normalized_parallelisms), + energy_per_bit=energy_per_bit, ) + @property + def effective_bandwidth(self) -> float: + return float(self.bandwidth) * float(self.util) + + +def _validate_dimension_parallelisms( + *, + dimension_label: str, + dimension_size: int, + normalized_names: Tuple[str, ...], + alias_map: Dict[str, str], + parallelism_params: Dict[str, object], +) -> None: + if not normalized_names: + return + + product = 1 + for name in normalized_names: + if name not in parallelism_params: + alias = alias_map.get(name, name) + raise ValueError( + f"network dimension '{dimension_label}' references parallelism '{alias}' " + "which is not defined in parallelism" + ) + value = parallelism_params[name] + if value in (None, False): + alias = alias_map.get(name, name) + raise ValueError( + f"network dimension '{dimension_label}' parallelism '{alias}' must have a " + "positive parallelism factor" + ) + try: + factor = int(value) + except (TypeError, ValueError) as exc: + alias = alias_map.get(name, name) + raise ValueError( + f"parallelism.{alias} must be an integer to compute network dimension sizes" + ) from exc + if factor < 1: + alias = alias_map.get(name, name) + raise ValueError( + f"parallelism.{alias} must be >= 1 to compute network dimension sizes" + ) + product *= factor -@dataclass -class TopologyConfig: - topology: str = None + if product != dimension_size: + readable = [alias_map.get(name, name) for name in normalized_names] + raise ValueError( + f"network dimension '{dimension_label}' size mismatch: declared size {dimension_size} " + f"but product of parallelisms {readable} is {product}" + ) - @classmethod - def from_dict(cls, d): - if d == "hybrid": - NotImplemented() - else: - return cls(topology=d) +def _parse_network_layout( + network_spec, + parallelism_params: Dict[str, object], +) -> Tuple[NetworkDimensionLayout, ...]: + if network_spec is None: + return tuple() -@dataclass -class NetworkTopologyConfig: - inter: TopologyConfig - intra: TopologyConfig + if not isinstance(network_spec, Sequence) or isinstance(network_spec, (str, bytes)): + raise ValueError("network section must be a sequence of dimension mappings") - @classmethod - def from_dict(cls, d): - return cls( - inter=TopologyConfig.from_dict(d["inter_node"]), - intra=TopologyConfig.from_dict(d["intra_node"]), + dimensions: List[NetworkDimensionLayout] = [] + for index, entry in enumerate(network_spec): + dimensions.append( + NetworkDimensionLayout.from_raw( + entry, + parallelism_params=parallelism_params, + index=index, + ) ) + return tuple(dimensions) + + +@dataclass(frozen=True) +class NetworkLayoutConfig: + dimensions: Tuple[NetworkDimensionLayout, ...] + parallelism_map: Dict[str, NetworkDimensionLayout] = field(default_factory=dict) + + def primary_dimension(self) -> Optional[NetworkDimensionLayout]: + return self.dimensions[0] if self.dimensions else None + + def dimension_for_parallelism(self, name: str) -> Optional[NetworkDimensionLayout]: + normalized = str(name).strip().lower() + if normalized in self.parallelism_map: + return self.parallelism_map[normalized] + return self.primary_dimension() + + def link_for_parallelism(self, name: str) -> Tuple[float, float]: + dim = self.dimension_for_parallelism(name) + if dim is None: + return 0.0, 0.0 + return dim.effective_bandwidth, dim.latency + + +def _build_network_layout_config( + dimensions: Sequence[NetworkDimensionLayout], +) -> NetworkLayoutConfig: + parallelism_map: Dict[str, NetworkDimensionLayout] = {} + for dim in dimensions: + for pname in dim.parallelisms: + if pname in parallelism_map: + raise ValueError( + f"parallelism '{pname}' assigned to multiple network dimensions" + ) + parallelism_map[pname] = dim + return NetworkLayoutConfig(dimensions=tuple(dimensions), parallelism_map=parallelism_map) + + +_PARALLELISM_DEFAULTS: Dict[str, object] = { + "auto": False, + "dp": 1, + "lp": 1, + "mb": 1, + "kp_hidden_dim1": 1, + "kp_softmax_dim1": 1, + "kp_embedding_dim1": 1, + "kp_projection_dim1": 1, + "kp_hidden_dim2": 1, + "kp_softmax_dim2": 1, + "kp_embedding_dim2": 1, + "kp_projection_dim2": 1, + "kp_hidden_type": -1, + "kp_softmax_type": -1, + "kp_embedding_type": -1, + "kp_projection_type": -1, + "t": "CR", + "tp": 1, + "cp": 1, + "tp_sp": False, + "kp1": 1, + "kp2": 1, +} @dataclass @@ -561,7 +713,7 @@ def use_moe(self) -> bool: ) SchedulingConfig = _namedtuple( - "scheduling_param", + "parallelism", [ "auto", "dp", @@ -598,9 +750,8 @@ def use_moe(self) -> bool: "sch_config", "area_breakdown", "perimeter_breakdown", - "system_config", "memory_hierarchy", - "network_topology", + "network_layout", ], ) @@ -658,16 +809,15 @@ def use_moe(self) -> bool: "sch_config", "area_breakdown", "perimeter_breakdown", - "system_config", "memory_hierarchy", - "network_topology", + "network_layout", "execution_backend", "inference_config", ], ) -MODELConfig = _namedtuple( - "MODELConfig", +ModelConfig = _namedtuple( + "ModelConfig", [ "model_config", "inference_config", @@ -675,50 +825,60 @@ def use_moe(self) -> bool: ) +def _convert_scalar_string(value: str): + try: + return float(value) + except ValueError: + pass + + digit = [int(s) for s in value.split() if s.isdigit()] + order = [str(s) for s in value.split() if not s.isdigit()] + if not (order and digit): + return value + + prefix = order[0][0] + bit = order[0][1] if len(order[0]) > 1 else "B" + mult = 1 + + if prefix == "K": + mult = 1024 + elif prefix == "M": + mult = 1024 * 1024 + elif prefix == "G": + mult = 1024 * 1024 * 1024 + elif prefix == "T": + mult = 1024 * 1024 * 1024 * 1024 + else: + raise ValueError(f"Unknown prefix '{prefix}' while parsing value '{value}'") + + if bit == "b": + mult = mult / 8 # Capacity is expected in Bytes + elif bit != "B": + raise ValueError(f"Unknown type '{bit}' while parsing value '{value}'") + + return digit[0] * mult + + +def _convert_value(value): + if isinstance(value, dict): + convert(value) + return value + if isinstance(value, list): + return [_convert_value(item) for item in value] + if isinstance(value, str): + try: + return _convert_scalar_string(value) + except ValueError: + return value + return value + + def convert(d): - for key1, val1 in d.items(): - for key2, val2 in val1.items(): - if isinstance(val2, dict): - for key3, val3 in val2.items(): - if isinstance(val3, str): - digit = [int(s) for s in val3.split() if s.isdigit()] - order = [str(s) for s in val3.split() if not s.isdigit()] - if order and digit: - assert len(order) >= 1 - assert len(digit) >= 1 - - prefix = order[0][0] - bit = order[0][1] - mult = 1 - - if prefix == "K": - mult = 1024 - elif prefix == "M": - mult = 1024 * 1024 - elif prefix == "G": - mult = 1024 * 1024 * 1024 - elif prefix == "T": - mult = 1024 * 1024 * 1024 * 1024 - else: - print( - "Unknown prefix: {} at {}: {}".format( - prefix, key3, val3 - ) - ) - exit(0) - - if bit == "b": - mult = mult / 8 # Capacity is expected in Bytes - elif bit == "B": - mult = mult - else: - print( - "Unknown type: {} at {}: {}".format(bit, key3, val3) - ) - exit(0) - - new_val = digit[0] * mult - d[key1][key2][key3] = new_val + if not isinstance(d, dict): + return d + for key, val in list(d.items()): + d[key] = _convert_value(val) + return d def parse_bandwidth_string(value): @@ -781,14 +941,14 @@ def parse_config(filename, config_type): """ with open(filename, "r") as f: try: - config_dict = _yaml.load(f, Loader=_ruamel.yaml.Loader) + config_dict = _yaml.safe_load(f) except _YAMLError as exc: hint = ( f"Failed to parse YAML config '{filename}'. " "Please check indentation and required sections like 'attention' and parameters such as 'num_experts'." ) raise ValueError(hint) from exc - # print(config_dict) + # print(config_dict) convert(config_dict) if config_type == "hardware": sw_block = dict(config_dict["sw_param"]) @@ -803,12 +963,22 @@ def parse_config(filename, config_type): h2d_bandwidth=h2d_bandwidth, dp_zero_stage=dp_zero_stage, ) - sch_params = dict(config_dict["scheduling_param"]) - if "tp" not in sch_params: - sch_params["tp"] = None - if "cp" not in sch_params: - sch_params["cp"] = None - sch_config = SchedulingConfig(**sch_params) + raw_parallelism_block = config_dict.get("parallelism", {}) + if not isinstance(raw_parallelism_block, dict): + raise ValueError("parallelism must be a mapping") + parallelism_params = dict(_PARALLELISM_DEFAULTS) + parallelism_params.update(raw_parallelism_block) + parallelism_params["auto"] = bool(parallelism_params["auto"]) + parallelism_params["tp_sp"] = bool(parallelism_params["tp_sp"]) + sch_config = SchedulingConfig(**parallelism_params) + scheduling_for_network = {str(k).lower(): v for k, v in parallelism_params.items()} + + network_spec = config_dict.get("network") + network_dimensions = _parse_network_layout(network_spec, scheduling_for_network) + if not network_dimensions: + raise ValueError("network section must define at least one dimension") + network_layout_config = _build_network_layout_config(network_dimensions) + tech_config = TechConfig.from_dict(config_dict["tech_param"]) # Optional breakdown configs (not needed for simplified configs) @@ -836,13 +1006,9 @@ def parse_config(filename, config_type): # Create dummy perimeter config with zeros perimeter_config = PerimeterBreakdownConfig(DRAM=0.1, inter_node=0.1, intra_node=0.1) - system_config = SystemHierarchyConfig.from_dict(config_dict["system_hierarchy"]) memory_hierarchy_config = MemoryHierarchyConfig.from_dict( config_dict["memory_hierarchy"] ) - network_topology_config = NetworkTopologyConfig.from_dict( - config_dict["network_topology"] - ) # execution backend (optional) eb_dict = config_dict.get("execution_backend", {}) eb_model = eb_dict.get("model", "analytical") @@ -891,21 +1057,20 @@ def parse_config(filename, config_type): sch_config=sch_config, area_breakdown=area_config, perimeter_breakdown=perimeter_config, - system_config=system_config, memory_hierarchy=memory_hierarchy_config, - network_topology=network_topology_config, + network_layout=network_layout_config, execution_backend=exec_backend, inference_config=inference_cfg, ) elif config_type == "LSTM": model_config = ModelLSTMConfig(**config_dict["model_param"]) - config = MODELConfig(model_config=model_config, inference_config=None) + config = ModelConfig(model_config=model_config, inference_config=None) elif config_type == "GEMM": mp = dict(config_dict["model_param"]) # copy if "backward" not in mp: mp["backward"] = False model_config = GEMMConfig(**mp) - config = MODELConfig(model_config=model_config, inference_config=None) + config = ModelConfig(model_config=model_config, inference_config=None) elif config_type == "LLM": mp = dict(config_dict["model_param"]) if "attention" not in mp: @@ -1114,13 +1279,13 @@ def _pop_required_int(field: str) -> int: num_experts=num_experts, top_k=top_k, ) - config = MODELConfig(model_config=model_config, inference_config=inference_config) + config = ModelConfig(model_config=model_config, inference_config=inference_config) else: raise ValueError("Invalid config type: {}".format(config_type)) # model_config = ModelConfig(**config_dict["model_param"]) # sw_config = SWConfig(**config_dict["sw_param"]) - # sch_config = SchedulingConfig(**config_dict["scheduling_param"]) + # sch_config = SchedulingConfig(**config_dict["parallelism"]) # tech_config = TechConfig.from_dict(config_dict["tech_param"]) # power_config = PowerBreakdownConfig.from_dict(config_dict["power_breakdown"]) # area_config = AreaBreakdownConfig.from_dict(config_dict["area_breakdown"]) diff --git a/configs/hardware-config/H100_SXM5_80GB.yaml b/configs/hardware-config/H100_SXM5_80GB.yaml index bfadaf2c..95d7ac2b 100644 --- a/configs/hardware-config/H100_SXM5_80GB.yaml +++ b/configs/hardware-config/H100_SXM5_80GB.yaml @@ -65,29 +65,6 @@ tech_param: dynamic_energy_per_bit: 110e-15 latency: 0 util: 1 - network: - intra_node: - latency: 5e-6 - util: 1 - bandwidth: 400 Gb - inter_node: - latency: 5e-6 - util: 0.96 - bandwidth: 400 Gb - -system_hierarchy: - num_devices_per_node: 1 - num_nodes: 1 - kp1_inter: false - kp2_inter: false - dp_inter: false - lp_inter: false - tp_inter: false - -network_topology: - # Allowed values: ring | fc | switch - inter_node: ring - intra_node: ring memory_hierarchy: l0: #Register Memory @@ -103,7 +80,7 @@ memory_hierarchy: type: DRAM scope: global -scheduling_param: +parallelism: auto: false ### LLM PARAMS ### dp: 1 # data parallelism. Acts as replica count during inference @@ -112,29 +89,26 @@ scheduling_param: tp: 1 # tensor parallelism tp_sp: true # tensor-sequence parallelism cp : 1 # context parallelism - ### LSTM/GEMM LEGACY PARAMS ### - t: CR - kp1: 1 - kp2: 1 - kp_hidden_dim1: 1 - kp_hidden_dim2: 1 - kp_softmax_dim1: 1 - kp_softmax_dim2: 1 - kp_embedding_dim1: 1 - kp_embedding_dim2: 1 - kp_projection_dim1: 1 - kp_projection_dim2: 1 - kp_hidden_type: -1 - kp_softmax_type: -1 - kp_embedding_type: -1 - kp_projection_type: -1 + +network: + - id: dim0 + label: device_fabric + size: 1 + topology: + type: Mesh + bandwidth: 400 Gb + latency: 5e-6 + energy_per_bit: 8e-12 # (J/bit) + util: 1.0 + collective_override: {} + parallelisms: [tp, cp, lp, dp, kp1, kp2] execution_backend: model: astra # analytical | astra astra: backend: astra mode: hybrid # hybrid | full_astrasim_hierarchical | full_astrasim_flattened - collectives: + collective_override: # Allowed values per op: auto | ring | direct | halvingDoubling | doubleBinaryTree # Note: leave to auto unless you have a specific reason to change it. all_gather: auto diff --git a/configs/hardware-config/a100_80GB.yaml b/configs/hardware-config/a100_80GB.yaml index 36373ca5..0ac5223c 100644 --- a/configs/hardware-config/a100_80GB.yaml +++ b/configs/hardware-config/a100_80GB.yaml @@ -33,10 +33,10 @@ tech_param: FMA_d1: 8 # tensor core can compute the results for an d1 x d2 x d1 mixed-precision matrix multiplication per clock FMA_d2: 4 dataflow: best # {wst | ast | ost | best}: wst: weight stationary, ast: activation stationary, ost: output stationary, best: best of all 3 - util: 1.0 + util: 100.0 DRAM: size: 80 GB - bandwidth: 1986 GB # (/s) + bandwidth: 198600 GB # (/s) dynamic_energy_per_bit: 4.5e-12 # (J) latency: 100e-9 # (s) util: 1 @@ -58,27 +58,7 @@ tech_param: dynamic_energy_per_bit: 110e-15 # (J) latency: 0 util: 1 - network: - intra_node: - latency: 5e-6 - util: 1 - bandwidth: 400 Gb - inter_node: - latency: 5e-6 - util: 0.96 - bandwidth: 400 Gb -system_hierarchy: - num_devices_per_node: 8 # currently only support astrasim integration with all devices in one node. - num_nodes: 1 - kp1_inter: false # these currently are not support for LLM mode. - kp2_inter: false - dp_inter: false - lp_inter: false - tp_inter: false -network_topology: - # Allowed values: ring | fc | switch - inter_node: ring - intra_node: ring + memory_hierarchy: l0: #Register Memory type: SRAM-R @@ -92,49 +72,69 @@ memory_hierarchy: l3: #Global Memory type: DRAM scope: global -scheduling_param: - auto: false # Not supported for LLM mode. - ### LLM PARAMS ### + +parallelism: dp: 2 # data parallelism. Acts as replica count during inference - lp: 2 # pipeline parallelism. Number of pipeline stages - mb: 1 # micro-batches (for best results set to >= lp) - tp: 2 # tensor parallelism + tp: 4 # tensor parallelism tp_sp: true # tensor-sequence parallelism - cp : 1 # context parallelism - ### LSTM/GEMM LEGACY PARAMS ### - t: "CR" - kp1: 1 - kp2: 1 - kp_hidden_dim1: 1 - kp_hidden_dim2: 1 - kp_softmax_dim1: 1 - kp_softmax_dim2: 1 - kp_embedding_dim1: 1 - kp_embedding_dim2: 1 - kp_projection_dim1: 1 - kp_projection_dim2: 1 - kp_hidden_type: -1 - kp_softmax_type: -1 - kp_embedding_type: -1 - kp_projection_type: -1 + cp: 1 # context parallelism. Not supported for inference. + lp: 2 # pipeline parallelism. Number of pipeline stages + # mb is not a true parallelism parameter, does not affect num of GPUs. + mb: 2 # micro-batches (for best results set to >= lp). + +network: + # currently only support astrasim integration with all devices in one node. + - id: dim0 + label: intra_wafer_mesh_silicon + size: 4 + topology: + type: Ring + bandwidth: 1000 Gb + latency: 5e-6 + energy_per_bit: 8e-12 # (J/bit) + util: 1.0 + collective_override: {} # leave empty to use 'auto' for all collectives unless an override is needed + parallelisms: [tp, cp] + - id: dim1 + label: inter_wafer_ring_infiniband + size: 2 + topology: + type: Ring + bandwidth: 200 Gb + latency: 5e-6 + energy_per_bit: 8e-12 # (J/bit) + util: 1.0 + collective_override: {} # defaults to 'auto' per collective + parallelisms: [lp] + - id: dim2 + label: wafer_dp_chains_ring_infiniband + size: 2 + topology: + type: ring + bandwidth: 200 Gb + latency: 5e-6 + energy_per_bit: 8e-12 # (J/bit) + util: 0.96 + collective_override: {} # defaults to 'auto' per collective + parallelisms: [dp] execution_backend: - model: astra # analytical | astra + model: analytical # analytical | astra astra: backend: astra mode: full_astrasim_flattened # hybrid | full_astrasim_hierarchical | full_astrasim_flattened - collectives: - # Allowed values per op: auto | ring | direct | halvingDoubling | doubleBinaryTree + collective_override: + # Allowed values per op: auto | ring | direct | halvingDoubling | doubleBinaryTree | mesh | mesh2d | hypercube | torus2d # Note: leave to auto unless you have a specific reason to change it. all_gather: auto all_reduce: auto reduce_scatter: auto all_to_all: auto sys_options: - endpoint_delay: 10 + endpoint_delay: 5 active_chunks_per_dimension: 32 # Higher value means significantly slower astrasim runtime, but more accurate especially for tensor parallelism - preferred_dataset_splits: 32 # Should be same as active_chunks_per_dimension + preferred_dataset_splits: 4 # Should be same as active_chunks_per_dimension inference: kvcache_type: hbm_only # Only supported mode for now. diff --git a/configs/hardware-config/a100_80GB_no_parallelism.yaml b/configs/hardware-config/a100_80GB_no_parallelism.yaml index aecc7e1f..a1cee857 100644 --- a/configs/hardware-config/a100_80GB_no_parallelism.yaml +++ b/configs/hardware-config/a100_80GB_no_parallelism.yaml @@ -70,18 +70,6 @@ tech_param: latency: 5e-6 util: 0.96 bandwidth: 400 Gb -system_hierarchy: - num_devices_per_node: 1 - num_nodes: 1 - kp1_inter: false - kp2_inter: false - dp_inter: false - lp_inter: false - tp_inter: false -network_topology: - # Allowed values: ring | fc | switch - inter_node: ring - intra_node: ring memory_hierarchy: l0: type: SRAM-R @@ -95,8 +83,9 @@ memory_hierarchy: l3: type: DRAM scope: global -scheduling_param: - auto: false # Not supported for LLM mode. + + +parallelism: ### LLM PARAMS ### dp: 1 # acts as replica count during inference lp: 1 # pipeline stages @@ -104,22 +93,20 @@ scheduling_param: tp: 1 # tensor parallelism tp_sp: false # tensor-sequence parallelism cp : 1 # context parallelism - ### LSTM/GEMM LEGACY PARAMS ### - t: "CR" - kp1: 1 - kp2: 1 - kp_hidden_dim1: 1 - kp_hidden_dim2: 1 - kp_softmax_dim1: 1 - kp_softmax_dim2: 1 - kp_embedding_dim1: 1 - kp_embedding_dim2: 1 - kp_projection_dim1: 1 - kp_projection_dim2: 1 - kp_hidden_type: -1 - kp_softmax_type: -1 - kp_embedding_type: -1 - kp_projection_type: -1 + +network: + - id: dim0 + label: device_fabric + size: 1 + topology: + type: Ring + bandwidth: 400 Gb + latency: 5e-6 + energy_per_bit: 8e-12 # (J/bit) + util: 1.0 + collective_override: {} + parallelisms: [tp, cp, lp, dp, kp1, kp2] + execution_backend: @@ -127,7 +114,7 @@ execution_backend: astra: backend: astra mode: full_astrasim_flattened # hybrid | full_astrasim_hierarchical | full_astrasim_flattened - collectives: + collective_override: # Allowed values per op: auto | ring | direct | halvingDoubling | doubleBinaryTree # Note: leave to auto unless you have a specific reason to change it. all_gather: auto diff --git a/configs/hardware-config/v100.yaml b/configs/hardware-config/v100.yaml index a68ff884..9cc82ade 100644 --- a/configs/hardware-config/v100.yaml +++ b/configs/hardware-config/v100.yaml @@ -70,20 +70,6 @@ tech_param: util: 0.96 bandwidth: 300 Gb -system_hierarchy: - num_devices_per_node: 1 - num_nodes: 1 - kp1_inter: false - kp2_inter: false - dp_inter: false - lp_inter: false - tp_inter: false - -network_topology: - # Allowed values: ring | fc | switch - inter_node: ring - intra_node: ring - memory_hierarchy: l0: #Register Memory type: SRAM-R @@ -98,8 +84,8 @@ memory_hierarchy: type: DRAM scope: global -scheduling_param: - auto: false + +parallelism: ### LLM PARAMS ### dp: 1 # acts as replica count during inference lp: 1 # pipeline stages @@ -107,29 +93,27 @@ scheduling_param: tp: 1 # tensor parallelism tp_sp: true # tensor-sequence parallelism cp : 1 # context parallelism - ### LSTM/GEMM LEGACY PARAMS ### - t: CR - kp1: 1 - kp2: 1 - kp_hidden_dim1: 1 - kp_hidden_dim2: 1 - kp_softmax_dim1: 1 - kp_softmax_dim2: 1 - kp_embedding_dim1: 1 - kp_embedding_dim2: 1 - kp_projection_dim1: 1 - kp_projection_dim2: 1 - kp_hidden_type: -1 - kp_softmax_type: -1 - kp_embedding_type: -1 - kp_projection_type: -1 + +network: + - id: dim0 + label: device_fabric + size: 1 + topology: + type: Mesh + bandwidth: 300 Gb + latency: 5e-6 + energy_per_bit: 8e-12 # (J/bit) + util: 1.0 + collective_override: {} + parallelisms: [tp, cp, lp, dp, kp1, kp2] + execution_backend: model: astra # analytical | astra astra: backend: astra mode: hybrid # hybrid | full_astrasim_hierarchical | full_astrasim_flattened - collectives: + collective_override: # Allowed values per op: auto | ring | direct | halvingDoubling | doubleBinaryTree # Note: leave to auto unless you have a specific reason to change it. all_gather: auto diff --git a/configs/hardware-config/waferscale_20v100_80hbm.yaml b/configs/hardware-config/waferscale_20v100_80hbm.yaml index 0ffdf7b3..c9f6876e 100644 --- a/configs/hardware-config/waferscale_20v100_80hbm.yaml +++ b/configs/hardware-config/waferscale_20v100_80hbm.yaml @@ -88,19 +88,6 @@ tech_param: util: 0.96 bandwidth: 300 Gb -system_hierarchy: - num_devices_per_node: 1 # devices per package essentially - num_nodes: 20 # one package - kp1_inter: true - kp2_inter: true - dp_inter: false - lp_inter: false - tp_inter: true - -network_topology: - inter_node: none #mesh, torus, crossbar, custom, hierarchical - intra_node: mesh - memory_hierarchy: l0: #Register Memory type: SRAM-R @@ -115,7 +102,7 @@ memory_hierarchy: type: DRAM scope: global -scheduling_param: +parallelism: auto: false dp: 1 lp: 1 @@ -123,27 +110,26 @@ scheduling_param: tp: 20 tp_sp: false cp: 1 - t: CR - kp1: 20 - kp2: 1 - kp_hidden_dim1: 1 - kp_hidden_dim2: 1 - kp_softmax_dim1: 1 - kp_softmax_dim2: 1 - kp_embedding_dim1: 1 - kp_embedding_dim2: 1 - kp_projection_dim1: 1 - kp_projection_dim2: 1 - kp_hidden_type: -1 - kp_softmax_type: -1 - kp_embedding_type: -1 - kp_projection_type: -1 + +network: + - id: dim0 + label: wafer_mesh + size: 20 + topology: + type: Mesh + bandwidth: 5000 Gb + latency: 5e-6 + energy_per_bit: 8e-12 # (J/bit) + util: 0.95 + collective_override: {} + parallelisms: [tp, cp, lp, dp] + execution_backend: model: astra # analytical | astra astra: backend: astra mode: hybrid # hybrid | full_astrasim_hierarchical | full_astrasim_flattened - collectives: + collective_override: # Allowed values per op: auto | ring | direct | halvingDoubling | doubleBinaryTree # Note: leave to auto unless you have a specific reason to change it. all_gather: auto diff --git a/configs/model-config/LLM.yaml b/configs/model-config/LLM.yaml index 7ced3ae6..22072e7c 100644 --- a/configs/model-config/LLM.yaml +++ b/configs/model-config/LLM.yaml @@ -16,4 +16,4 @@ model_param: top_k: 1 # Number of active experts per token in MoE intermediate_size: 32768 # Output dimension of the first FFN layer (typically 4×hidden_dim) vocab_size: 800000 # Vocabulary size - num_layers: 2 # Number of transformer layers + num_layers: 1 # Number of transformer layers diff --git a/configs/model-config/LLM_inf.yaml b/configs/model-config/LLM_inf.yaml index f66a6db9..b16307bb 100644 --- a/configs/model-config/LLM_inf.yaml +++ b/configs/model-config/LLM_inf.yaml @@ -3,18 +3,18 @@ model_param: run_type: "inference" # Type of run: "training" or "inference" tied_embeddings: true # Share input and output embedding weights (WIP: affects memory estimation only) model_type: "gpt" # gpt assumes dense-style architecture with GELU; llama assumes llama-style architecture with SwiGLU - batch_size: 64 # Sequences per global batch - seq_len: 4096 # Total context length (prefill + decode_len for inference) + batch_size: 2048 # Sequences per global batch + seq_len: 2048 # Total context length (prefill + decode_len for inference) decode_len: 1024 # Number of autoregressive decode steps. hidden_dim: 8192 # Transformer hidden dimension attention: # Attention-related parameters attention_type: "mha" # Supported attention types: "mha" or "gqa" ("mla" not implemented) num_heads: 64 # Total number of attention heads - kv_heads: 8 # Number of key/value heads (used only when attention_type = "gqa") + kv_heads: 64 # Number of key/value heads (used only when attention_type = "gqa") use_flashattention: false # Enable FlashAttention optimization attention_tile_size: 256 # Tile size used in attention computation num_experts: 1 # Number of experts in MoE layers. Set to 1 to disable MoE. top_k: 1 # Number of active experts per token in MoE intermediate_size: 32768 # Output dimension of the first FFN layer (typically 4×hidden_dim) vocab_size: 129280 # Vocabulary size - num_layers: 96 # Number of transformer layers + num_layers: 32 # Number of transformer layers diff --git a/configs/old-detailed-hw-configs/H100_SXM5_80GB.yaml b/configs/old-detailed-hw-configs/H100_SXM5_80GB.yaml index 44319ea2..3c4f0450 100644 --- a/configs/old-detailed-hw-configs/H100_SXM5_80GB.yaml +++ b/configs/old-detailed-hw-configs/H100_SXM5_80GB.yaml @@ -99,7 +99,7 @@ tech_param: latency: 5e-6 nominal_frequency: 3e9 nominal_voltage: 1 - nominal_energy_per_link: 1e-12 + energy_per_bit: 1e-12 nominal_area_per_link: 100e-6 # (mm^2) num_links_per_mm: 400 threshold_voltage: 0.25 @@ -109,7 +109,7 @@ tech_param: latency: 5e-6 nominal_frequency: 20e9 nominal_voltage: 1.2 - nominal_energy_per_link: 13e-12 + energy_per_bit: 13e-12 nominal_area_per_link: 0.045 num_links_per_mm: 20 #Assuming 4 layers on PCB and 200 um pitch for high speed signalling threshold_voltage: 0.35 @@ -204,7 +204,7 @@ execution_backend: astra: backend: astra mode: hybrid # hybrid | full_astrasim_hierarchical | full_astrasim_flattened - collectives: + collective_override: # Allowed values per op: auto | ring | direct | halvingDoubling | doubleBinaryTree # Note: leave to auto unless you have a specific reason to change it. all_gather: auto diff --git a/configs/old-detailed-hw-configs/a100_80GB.yaml b/configs/old-detailed-hw-configs/a100_80GB.yaml index 75a5a3c6..debc8ce6 100644 --- a/configs/old-detailed-hw-configs/a100_80GB.yaml +++ b/configs/old-detailed-hw-configs/a100_80GB.yaml @@ -89,7 +89,7 @@ tech_param: latency: 5e-6 nominal_frequency: 3e9 nominal_voltage: 1 - nominal_energy_per_link: 8e-12 + energy_per_bit: 8e-12 nominal_area_per_link: 100e-6 num_links_per_mm: 400 threshold_voltage: 0.25 @@ -99,7 +99,7 @@ tech_param: latency: 5e-6 nominal_frequency: 20e9 nominal_voltage: 1.2 - nominal_energy_per_link: 13e-12 + energy_per_bit: 13e-12 nominal_area_per_link: 0.045 num_links_per_mm: 20 threshold_voltage: 0.35 @@ -189,7 +189,7 @@ execution_backend: astra: backend: astra mode: full_astrasim_flattened # hybrid | full_astrasim_hierarchical | full_astrasim_flattened - collectives: + collective_override: # Allowed values per op: auto | ring | direct | halvingDoubling | doubleBinaryTree # Note: leave to auto unless you have a specific reason to change it. all_gather: auto diff --git a/configs/old-detailed-hw-configs/v100.yaml b/configs/old-detailed-hw-configs/v100.yaml index 71721296..f2a59f9c 100644 --- a/configs/old-detailed-hw-configs/v100.yaml +++ b/configs/old-detailed-hw-configs/v100.yaml @@ -91,7 +91,7 @@ tech_param: latency: 5e-6 nominal_frequency: 3e9 nominal_voltage: 1 - nominal_energy_per_link: 1e-12 + energy_per_bit: 1e-12 nominal_area_per_link: 100e-6 #mm^2 num_links_per_mm: 400 threshold_voltage: 0.25 @@ -101,7 +101,7 @@ tech_param: latency: 5e-6 nominal_frequency: 20e9 nominal_voltage: 1.2 - nominal_energy_per_link: 13e-12 + energy_per_bit: 13e-12 nominal_area_per_link: 0.045 num_links_per_mm: 20 #Assuming 4 layers on PCB and 200 um pitch for high speed signalling threshold_voltage: 0.35 @@ -197,7 +197,7 @@ execution_backend: astra: backend: astra mode: hybrid # hybrid | full_astrasim_hierarchical | full_astrasim_flattened - collectives: + collective_override: # Allowed values per op: auto | ring | direct | halvingDoubling | doubleBinaryTree # Note: leave to auto unless you have a specific reason to change it. all_gather: auto diff --git a/configs/old-detailed-hw-configs/waferscale_20v100_80hbm.yaml b/configs/old-detailed-hw-configs/waferscale_20v100_80hbm.yaml index 7e99ef83..2cc8fa6f 100644 --- a/configs/old-detailed-hw-configs/waferscale_20v100_80hbm.yaml +++ b/configs/old-detailed-hw-configs/waferscale_20v100_80hbm.yaml @@ -111,7 +111,7 @@ tech_param: nominal_frequency: 12.48e9 # THIS IS BITS/SECOND PER LINK!!!! Look at lines 502 in hw_component.py, 678 in perf.py nominal_voltage: 1 # energy per link - nominal_energy_per_link: 8e-12 # This is probably pessimistic, as it is likely stated for 300mm. Our avg distance is 150mm. 8pj/Bit based on https://www.techpowerup.com/forums/threads/nvidia-is-preparing-co-packaged-photonics-for-nvlink.276139/#:~:text=While%20the%20current%20NVLink%202.0,(4%20pJ%2Fb), https://www.techpowerup.com/img/fpXXsSk543ZUFWcx_thm.jpg + energy_per_bit: 8e-12 # This is probably pessimistic, as it is likely stated for 300mm. Our avg distance is 150mm. 8pj/Bit based on https://www.techpowerup.com/forums/threads/nvidia-is-preparing-co-packaged-photonics-for-nvlink.276139/#:~:text=While%20the%20current%20NVLink%202.0,(4%20pJ%2Fb), https://www.techpowerup.com/img/fpXXsSk543ZUFWcx_thm.jpg nominal_area_per_link: 1e-12 #mm^2. We are setting this to an extremely low value on purpose to ignore mem ctrl area. This is basically 0 but we cant set to 0 because of div errors. num_links_per_mm: 3.363 #increased a tiny bit to account for integer division and the fact that I dropped digits threshold_voltage: 0.25 @@ -121,7 +121,7 @@ tech_param: latency: 5e-6 nominal_frequency: 20e9 nominal_voltage: 1.2 - nominal_energy_per_link: 13e-12 + energy_per_bit: 13e-12 nominal_area_per_link: 0.045 num_links_per_mm: 20 #Assuming 4 layers on PCB and 200 um pitch for high speed signalling threshold_voltage: 0.35 @@ -212,7 +212,7 @@ execution_backend: astra: backend: astra mode: hybrid # hybrid | full_astrasim_hierarchical | full_astrasim_flattened - collectives: + collective_override: # Allowed values per op: auto | ring | direct | halvingDoubling | doubleBinaryTree # Note: leave to auto unless you have a specific reason to change it. all_gather: auto diff --git a/configs/old_lstm_hw_configs/H100_SXM5_80GB.yaml b/configs/old_lstm_hw_configs/H100_SXM5_80GB.yaml new file mode 100644 index 00000000..183a9f33 --- /dev/null +++ b/configs/old_lstm_hw_configs/H100_SXM5_80GB.yaml @@ -0,0 +1,150 @@ +# DeepFlow configuration for NVIDIA H100 PCIe 80 GB +# specifications from the following sources: +# [NVIDIA] https://resources.nvidia.com/en-us-hopper-architecture/nvidia-h100-tensor-c +# [TePU] https://www.techpowerup.com/gpu-specs/h100-sxm5-80-gb.c3900 +# [Semi] https://semianalysis.com/2025/06/23/nvidia-tensor-core-evolution-from-volta-to-blackwell/#4th-generation-tensor-core-warpgroup-level-asynchronous-mma + +sw_param: + kernel_launch_overhead: 9e-6 # (s) + # dp_zero_stage selects ZeRO-style data-parallel sharding: 0=DDP, 1=ZeRO-1, 2=ZeRO-2, 3=ZeRO-3 + # only used if data parallelism (dp) > 1 + dp_zero_stage: 0 + precision: + # tensor_format accepts numeric byte counts or dtype strings (fp16, bf16, fp32). + # Take care that the number of flops defined below has been carefully tuned for FP16/BF16 formats. + # For FP32 or other formats they need to be manually adjusted. + # The following are equivalent: (tensor_format: 2.0, tensor_format: fp16, tensor_format: bf16). + tensor_format: bf16 # baseline storage width for activations/gradients. + # mixed_precision: true | false + # -> false forces all activations/weights/gradients/optimizer states/stats to use tensor_format bytes. + # -> true enables AMP-style FP32 accumulations/optimizer states depending on param_storage_mode. + # if true, optimizer states and stats are always FP32. + mixed_precision: false + # param_storage_mode: as_tensor_format | tensor_plus_fp32_master | fp32_params. + # -> as_tensor_format makes weights/gradients use tensor_format bytes. + # -> tensor_plus_fp32_master makes weights/gradients use tensor_format bytes and + # keeps a separate FP32 master copy for weights. + # -> fp32_params makes weights *and* gradients use FP32 bytes without a master copy. + param_storage_mode: as_tensor_format + # kv_cache precision may be set as bytes, dtype string, or "as_tensor_format". + kv_cache: as_tensor_format + h2d_bandwidth: 13297674240 # Bytes/s; set -1 to disable H2D timing + +tech_param: + core: + operating_frequency: 1.98e9 # (Hz) GPU boost clock frequency [TePU] + num_bundles: 132 # number of bundles (SMs) [TePU] + nominal_flop_rate_per_mcu: 1024 # F16 FLOP / cycle / SM is 4096 [Semi] + nominal_energy_per_flop: 1.8e-13 + num_mcu_per_bundle: 4 # number of tensorcores per SM [TePU] + FMA_d1: 8 # tensor core can compute the results for an d1 x d2 x d1 mixed-precision matrix multiplication per clock + FMA_d2: 4 + dataflow: best #{wst. ast, ost, best, none}: wst: weight stationary, ast: activation stationary, ost: output stationary + util: 0.75 #util should be 0.75 (12/16) for tensorcore + DRAM: + size: 80 GB # [TePU] + bandwidth: 3440 GB # (/s) [TePU] + dynamic_energy_per_bit: 4.5e-12 # (J) [https://passlab.github.io/mchpc/mchpc2019/presentations/MCHPC_Pawlowski_keynote.pdf, slide 36] + latency: 100e-9 # (s) + util: 1 + SRAM-L2: + size: 50 MB # [TePU] + bandwidth: 8138 GB # NCU L2 peak traffic + dynamic_energy_per_bit: 130e-15 + latency: 0 + util: 1 + SRAM-L1: + size: 30096 KB # max configurable shared memory per SM * number of SMs [NVIDIA] + bandwidth: 22392 GB # NCU L1/TEX peak traffic + dynamic_energy_per_bit: 130e-15 + latency: 0 + util: 1 + SRAM-R: + size: 33 MB # register file memory per SM * number of SMs [NVIDIA] + bandwidth: 244 TB # traffic for 2 reads and 1 write for GEMMs computed per clock + dynamic_energy_per_bit: 110e-15 + latency: 0 + util: 1 + network: + intra_node: + latency: 5e-6 + util: 1 + bandwidth: 400 Gb + inter_node: + latency: 5e-6 + util: 0.96 + bandwidth: 400 Gb + +system_hierarchy: + num_devices_per_node: 1 + num_nodes: 1 + kp1_inter: false + kp2_inter: false + dp_inter: false + lp_inter: false + tp_inter: false + +network_topology: + # Allowed values: ring | fc | switch + inter_node: ring + intra_node: ring + +memory_hierarchy: + l0: #Register Memory + type: SRAM-R + scope: mcu + l1: #Shared Memory + type: SRAM-L1 + scope: mcu-bundle + l2: #L2 + type: SRAM-L2 + scope: global + l3: #Global Memory + type: DRAM + scope: global + +scheduling_param: + auto: false + ### LLM PARAMS ### + dp: 1 # data parallelism. Acts as replica count during inference + lp: 1 # pipeline parallelism. Number of pipeline stages + mb: 1 # micro-batches (for best results set to >= lp) + tp: 1 # tensor parallelism + tp_sp: true # tensor-sequence parallelism + cp : 1 # context parallelism + ### LSTM/GEMM LEGACY PARAMS ### + t: CR + kp1: 1 + kp2: 1 + kp_hidden_dim1: 1 + kp_hidden_dim2: 1 + kp_softmax_dim1: 1 + kp_softmax_dim2: 1 + kp_embedding_dim1: 1 + kp_embedding_dim2: 1 + kp_projection_dim1: 1 + kp_projection_dim2: 1 + kp_hidden_type: -1 + kp_softmax_type: -1 + kp_embedding_type: -1 + kp_projection_type: -1 + +execution_backend: + model: astra # analytical | astra + astra: + backend: astra + mode: hybrid # hybrid | full_astrasim_hierarchical | full_astrasim_flattened + collective_override: + # Allowed values per op: auto | ring | direct | halvingDoubling | doubleBinaryTree + # Note: leave to auto unless you have a specific reason to change it. + all_gather: auto + all_reduce: auto + reduce_scatter: auto + all_to_all: auto + sys_options: + endpoint_delay: 10 + active_chunks_per_dimension: 32 # Higher value means significantly slower astrasim runtime, but more accurate especially for tensor parallelism + preferred_dataset_splits: 32 # Should be same as active_chunks_per_dimension + +inference: + kvcache_type: hbm_only diff --git a/configs/old_lstm_hw_configs/a100_80GB.yaml b/configs/old_lstm_hw_configs/a100_80GB.yaml new file mode 100644 index 00000000..cfbebc6c --- /dev/null +++ b/configs/old_lstm_hw_configs/a100_80GB.yaml @@ -0,0 +1,143 @@ +sw_param: + kernel_launch_overhead: 9e-6 # (s) + # dp_zero_stage selects ZeRO-style data-parallel sharding: 0=DDP, 1=ZeRO-1, 2=ZeRO-2, 3=ZeRO-3 + # only used if data parallelism (dp) > 1 + dp_zero_stage: 2 + precision: + # tensor_format accepts numeric byte counts or dtype strings (fp16, bf16, fp32). + # Take care that the number of flops defined below has been carefully tuned for FP16/BF16 formats. + # For FP32 or other formats they need to be manually adjusted. + # The following are equivalent: (tensor_format: 2.0, tensor_format: fp16, tensor_format: bf16). + tensor_format: bf16 # baseline storage width for activations/gradients. + # mixed_precision: true | false + # -> false forces all activations/weights/gradients/optimizer states/stats to use tensor_format bytes. + # -> true enables AMP-style FP32 accumulations/optimizer states depending on param_storage_mode. + # if true, optimizer states and stats are always FP32. + mixed_precision: false + # param_storage_mode: as_tensor_format | tensor_plus_fp32_master | fp32_params. + # -> as_tensor_format makes weights/gradients use tensor_format bytes. + # -> tensor_plus_fp32_master makes weights/gradients use tensor_format bytes and + # keeps a separate FP32 master copy for weights. + # -> fp32_params makes weights *and* gradients use FP32 bytes without a master copy. + param_storage_mode: as_tensor_format + # kv_cache precision may be set as bytes, dtype string, or "as_tensor_format". + kv_cache: as_tensor_format + h2d_bandwidth: -1 # Bytes/s; set -1 to disable H2D timing +tech_param: + core: + operating_frequency: 1.41e9 # (Hz) + num_bundles: 108 + nominal_flop_rate_per_mcu: 512 + nominal_energy_per_flop: 1.8e-13 # (J) + num_mcu_per_bundle: 4 # number of tensorcores per SM + FMA_d1: 8 # tensor core can compute the results for an d1 x d2 x d1 mixed-precision matrix multiplication per clock + FMA_d2: 4 + dataflow: best # {wst | ast | ost | best}: wst: weight stationary, ast: activation stationary, ost: output stationary, best: best of all 3 + util: 100.0 + DRAM: + size: 80 GB + bandwidth: 198600 GB # (/s) + dynamic_energy_per_bit: 4.5e-12 # (J) + latency: 100e-9 # (s) + util: 1 + SRAM-L2: + size: 40 MB + bandwidth: 7050 GB + dynamic_energy_per_bit: 130e-15 # (J) + latency: 0 + util: 1 + SRAM-L1: + size: 20736 KB + bandwidth: 18 TB + dynamic_energy_per_bit: 130e-15 # (J) + latency: 0 + util: 1 + SRAM-R: + size: 27 MB + bandwidth: 122 TB + dynamic_energy_per_bit: 110e-15 # (J) + latency: 0 + util: 1 + network: + intra_node: + latency: 5e-6 + util: 1 + bandwidth: 200 Gb + inter_node: + latency: 5e-6 + util: 0.96 + bandwidth: 200 Gb +system_hierarchy: + num_devices_per_node: 8 # currently only support astrasim integration with all devices in one node. + num_nodes: 1 + ## LLM PARAMS ## + dp_inter: false + lp_inter: false + tp_inter: false + cp_inter: false + ## LEGACY LSTM PARAMS ## + kp1_inter: false + kp2_inter: false +network_topology: + # Allowed values: ring | fc | switch | HyperCube | Mesh | Mesh2D | Torus2D + inter_node: ring + intra_node: ring +memory_hierarchy: + l0: #Register Memory + type: SRAM-R + scope: mcu + l1: #Shared Memory + type: SRAM-L1 + scope: mcu-bundle + l2: #L2 + type: SRAM-L2 + scope: global + l3: #Global Memory + type: DRAM + scope: global +scheduling_param: + auto: false # Not supported for LLM mode. + ### LLM PARAMS ### + dp: 2 # data parallelism. Acts as replica count during inference + lp: 2 # pipeline parallelism. Number of pipeline stages + mb: 2 # micro-batches (for best results set to >= lp) + tp: 2 # tensor parallelism + tp_sp: true # tensor-sequence parallelism + cp : 1 # context parallelism + ### LSTM/GEMM LEGACY PARAMS ### + t: "CR" + kp1: 1 + kp2: 1 + kp_hidden_dim1: 1 + kp_hidden_dim2: 1 + kp_softmax_dim1: 1 + kp_softmax_dim2: 1 + kp_embedding_dim1: 1 + kp_embedding_dim2: 1 + kp_projection_dim1: 1 + kp_projection_dim2: 1 + kp_hidden_type: -1 + kp_softmax_type: -1 + kp_embedding_type: -1 + kp_projection_type: -1 + + +execution_backend: + model: astra # analytical | astra + astra: + backend: astra + mode: full_astrasim_flattened # hybrid | full_astrasim_hierarchical | full_astrasim_flattened + collective_override: + # Allowed values per op: auto | ring | direct | halvingDoubling | doubleBinaryTree | mesh | mesh2d | hypercube | torus2d + # Note: leave to auto unless you have a specific reason to change it. + all_gather: auto + all_reduce: auto + reduce_scatter: auto + all_to_all: auto + sys_options: + endpoint_delay: 5 + active_chunks_per_dimension: 32 # Higher value means significantly slower astrasim runtime, but more accurate especially for tensor parallelism + preferred_dataset_splits: 4 # Should be same as active_chunks_per_dimension + +inference: + kvcache_type: hbm_only # Only supported mode for now. diff --git a/configs/old_lstm_hw_configs/v100.yaml b/configs/old_lstm_hw_configs/v100.yaml new file mode 100644 index 00000000..b930572b --- /dev/null +++ b/configs/old_lstm_hw_configs/v100.yaml @@ -0,0 +1,145 @@ + +sw_param: + kernel_launch_overhead: 9e-6 #9us + # dp_zero_stage selects ZeRO-style data-parallel sharding: 0=DDP, 1=ZeRO-1, 2=ZeRO-2, 3=ZeRO-3 + # only used if data parallelism (dp) > 1 + dp_zero_stage: 0 + precision: + # tensor_format accepts numeric byte counts or dtype strings (fp16, bf16, fp32). + # Take care that the number of flops defined below has been carefully tuned for FP16/BF16 formats. + # For FP32 or other formats they need to be manually adjusted. + # The following are equivalent: (tensor_format: 2.0, tensor_format: fp16, tensor_format: bf16). + tensor_format: bf16 # baseline storage width for activations/gradients. + # mixed_precision: true | false + # -> false forces all activations/weights/gradients/optimizer states/stats to use tensor_format bytes. + # -> true enables AMP-style FP32 accumulations/optimizer states depending on param_storage_mode. + # if true, optimizer states and stats are always FP32. + mixed_precision: false + # param_storage_mode: as_tensor_format | tensor_plus_fp32_master | fp32_params. + # -> as_tensor_format makes weights/gradients use tensor_format bytes. + # -> tensor_plus_fp32_master makes weights/gradients use tensor_format bytes and + # keeps a separate FP32 master copy for weights. + # -> fp32_params makes weights *and* gradients use FP32 bytes without a master copy. + param_storage_mode: as_tensor_format + # kv_cache precision may be set as bytes, dtype string, or "as_tensor_format". + kv_cache: as_tensor_format + h2d_bandwidth: 13297674240 # Bytes/s; set -1 to disable H2D timing + +tech_param: + core: + operating_frequency: 1.33e9 #Hz + num_bundles: 80 # V100 has 80 SMs + nominal_flop_rate_per_mcu: 128 + nominal_energy_per_flop: 1.8e-13 + num_mcu_per_bundle: 8 #a bundle can be thought of as an SM, number of tensorcores per SM + FMA_d1: 4 # tensor core can compute the results for an d1 x d2 x d1 mixed-precision matrix multiplication per clock + FMA_d2: 4 + dataflow: best #{wst. ast, ost, best, none}: wst: weight stationary, ast: activation stationary, ost: output stationary + util: 0.85 #util should be 0.75 (12/16) for tensorcore + DRAM: + size: 16 GB # V100 has 16GB HBM2 + bandwidth: 900 GB # V100 has ~900 GB/s bandwidth + dynamic_energy_per_bit: 4.2e-12 #Joules + latency: 100e-9 + util: 1 + SRAM-L2: + size: 6 MB # V100 L2 cache + bandwidth: 3072 GB # Estimated L2 bandwidth + dynamic_energy_per_bit: 130e-15 + latency: 0 + util: 0.45 + SRAM-L1: + size: 6400 KB # 80 KB per SM * 80 SMs + bandwidth: 15 TB # Estimated L1 bandwidth + dynamic_energy_per_bit: 130e-15 + latency: 0 + util: 1 + SRAM-R: + size: 20 MB # Register file + bandwidth: 80 TB # Estimated register bandwidth + dynamic_energy_per_bit: 110e-15 + latency: 0 + util: 1 + network: + intra_node: + latency: 5e-6 + util: 1 + bandwidth: 300 Gb + inter_node: + latency: 5e-6 + util: 0.96 + bandwidth: 300 Gb + +system_hierarchy: + num_devices_per_node: 1 + num_nodes: 1 + kp1_inter: false + kp2_inter: false + dp_inter: false + lp_inter: false + tp_inter: false + +network_topology: + # Allowed values: ring | fc | switch + inter_node: ring + intra_node: ring + +memory_hierarchy: + l0: #Register Memory + type: SRAM-R + scope: mcu + l1: #Shared Memory + type: SRAM-L1 + scope: mcu-bundle + l2: #L2 + type: SRAM-L2 + scope: global + l3: #Global Memory + type: DRAM + scope: global + +scheduling_param: + auto: false + ### LLM PARAMS ### + dp: 1 # acts as replica count during inference + lp: 1 # pipeline stages + mb: 1 # number of micro-batches + tp: 1 # tensor parallelism + tp_sp: true # tensor-sequence parallelism + cp : 1 # context parallelism + ### LSTM/GEMM LEGACY PARAMS ### + t: CR + kp1: 1 + kp2: 1 + kp_hidden_dim1: 1 + kp_hidden_dim2: 1 + kp_softmax_dim1: 1 + kp_softmax_dim2: 1 + kp_embedding_dim1: 1 + kp_embedding_dim2: 1 + kp_projection_dim1: 1 + kp_projection_dim2: 1 + kp_hidden_type: -1 + kp_softmax_type: -1 + kp_embedding_type: -1 + kp_projection_type: -1 + +execution_backend: + model: astra # analytical | astra + astra: + backend: astra + mode: hybrid # hybrid | full_astrasim_hierarchical | full_astrasim_flattened + collective_override: + # Allowed values per op: auto | ring | direct | halvingDoubling | doubleBinaryTree + # Note: leave to auto unless you have a specific reason to change it. + all_gather: auto + all_reduce: auto + reduce_scatter: auto + all_to_all: auto + sys_options: + endpoint_delay: 10 + active_chunks_per_dimension: 32 # Higher value means significantly slower astrasim runtime, but more accurate especially for tensor parallelism + preferred_dataset_splits: 32 # Should be same as active_chunks_per_dimension + +inference: + kvcache_type: hbm_only diff --git a/GD_search.py b/deprecated/GD_search.py similarity index 99% rename from GD_search.py rename to deprecated/GD_search.py index d496160b..d8ddde80 100644 --- a/GD_search.py +++ b/deprecated/GD_search.py @@ -12,8 +12,7 @@ import util import perf -import ruamel as _ruamel -import ruamel.yaml as _yaml +import yaml as _yaml import math @@ -184,7 +183,7 @@ def initialize(self): random.seed(random_seed) with open(self.exp_config, "r") as f: - config_dict = _yaml.load(f, Loader=_ruamel.yaml.Loader) + config_dict = _yaml.safe_load(f) for param_class in self.parameters: if param_class == 'perimeter_breakdown': @@ -263,7 +262,7 @@ def populate_config(self, params, exp_dir): config_dict = {} with open(exp_config, "r") as f: - config_dict = _yaml.load(f, Loader=_ruamel.yaml.Loader) + config_dict = _yaml.safe_load(f) config_dict['area_breakdown']['proc_chip_area_budget'] = self.chip_area_budget if self.chip_area_budget != -1 else config_dict['area_breakdown']['proc_chip_area_budget']#in mm^2 for param_class in params: @@ -305,7 +304,7 @@ def populate_config(self, params, exp_dir): config_dict['system_hierarchy'][param] = self.system_hierarchy_params[param] with open(config_file, 'w') as yaml_file: - _yaml.dump(config_dict, yaml_file, default_flow_style=False) + _yaml.safe_dump(config_dict, yaml_file, default_flow_style=False, sort_keys=False) #print(config_file) #if self.debug: diff --git a/deviceMapping.py b/deprecated/deviceMapping.py similarity index 100% rename from deviceMapping.py rename to deprecated/deviceMapping.py diff --git a/energy.py b/deprecated/energy.py similarity index 100% rename from energy.py rename to deprecated/energy.py diff --git a/main.py b/deprecated/main.py similarity index 100% rename from main.py rename to deprecated/main.py diff --git a/topology.py b/deprecated/topology.py similarity index 100% rename from topology.py rename to deprecated/topology.py diff --git a/topology_hack.py b/deprecated/topology_hack.py similarity index 100% rename from topology_hack.py rename to deprecated/topology_hack.py diff --git a/examples/llm_astra_inference.sh b/examples/llm_astra_inference.sh index c8cb3551..553f31f1 100755 --- a/examples/llm_astra_inference.sh +++ b/examples/llm_astra_inference.sh @@ -4,4 +4,4 @@ # Setup model parameter in configs/model-config/Llama2-7B_inf.yaml # Setup hardware parameters (incl. parallelism/network) in configs/hardware-config/a100_80GB.yaml -python run_perf.py --hardware_config configs/hardware-config/a100_80GB.yaml --model_config configs/model-config/Llama2-7B_inf.yaml \ No newline at end of file +python run_perf.py --hardware_config configs/hardware-config/a100_80GB.yaml --model_config configs/model-config/LLM_inf.yaml \ No newline at end of file diff --git a/examples/llm_astra_inference_debug_graphviz.sh b/examples/llm_astra_inference_debug_graphviz.sh index eda86b7a..70426991 100755 --- a/examples/llm_astra_inference_debug_graphviz.sh +++ b/examples/llm_astra_inference_debug_graphviz.sh @@ -4,4 +4,4 @@ # Setup model parameter in configs/model-config/Llama2-7B.yaml # Setup hardware parameters (incl. parallelism/network) in configs/hardware-config/a100_80GB.yaml -DEEPFLOW_PERSIST_ASTRASIM_ARTIFACTS=1 DEEPFLOW_VISUALIZE_GRAPHS=1 DEEPFLOW_PERSIST_ARTIFACT_VIZ=1 python run_perf.py --hardware_config configs/hardware-config/a100_80GB.yaml --model_config configs/model-config/Llama2-7B_inf.yaml \ No newline at end of file +DEEPFLOW_PERSIST_ASTRASIM_ARTIFACTS=1 DEEPFLOW_VISUALIZE_GRAPHS=1 DEEPFLOW_PERSIST_ARTIFACT_VIZ=1 python run_perf.py --hardware_config configs/hardware-config/a100_80GB.yaml --model_config configs/model-config/LLM_inf.yaml \ No newline at end of file diff --git a/hw_component.py b/hw_component.py index 9e2f9133..73abb13b 100644 --- a/hw_component.py +++ b/hw_component.py @@ -1,17 +1,13 @@ import math +from typing import Tuple + import numpy as np import util -from topology_hack import Topology kilo = 1024.0 giga = 1024.0 * 1024.0 * 1024.0 -class System: - def __init__(self, exp_config): - self.num_nodes_per_wafer = exp_config.system_hierarchy.num_nodes_per_wafer - - class Base: def __init__(self, exp_config): self.exp_config = exp_config @@ -708,188 +704,48 @@ def printStats(self, f): class Network(Base): def __init__(self, exp_config): super().__init__(exp_config) - self.intra_network = SubNetwork( - exp_config, - exp_config.tech_config.network.intra_node, - exp_config.power_breakdown.network.intra_node, - exp_config.area_breakdown.network.intra_node, - "intra", - ) - self.inter_network = SubNetwork( - exp_config, - exp_config.tech_config.network.inter_node, - exp_config.power_breakdown.network.inter_node, - exp_config.area_breakdown.network.inter_node, - "inter", - ) + layout = getattr(exp_config, "network_layout", None) + if layout is None: + raise ValueError("hardware config is missing network layout information") + self.layout = layout + self.energies_per_bit = [ + float(getattr(dim, "energy_per_bit", 0.0)) + for dim in layout.dimensions + ] def calcThroughput(self): - inter_throughput = self.inter_network.calcThroughput() - intra_throughput = self.intra_network.calcThroughput() - - return intra_throughput, inter_throughput + primary = self.layout.primary_dimension() + if primary is None: + return 0.0, 0.0 + bw = float(primary.effective_bandwidth) + return bw, bw def calcLatency(self): - inter_latency = self.inter_network.latency - intra_latency = self.intra_network.latency - - return intra_latency, inter_latency - - def printStats(self, f): - self.inter_network.printStats(f, "inter") - self.intra_network.printStats(f, "intra") - - -class SubNetwork(Base): - def __init__( - self, exp_config, net_config, power_breakdown, area_breakdown, netLevel - ): - super().__init__(exp_config) - self.tot_power = power_breakdown # * self.TDP - self.tot_area = area_breakdown # * self.proc_chip_area_budget - # TODO: Rename core_perimeter to proc_chip_perimeter - self.latency = net_config.latency - self.nominal_freq = net_config.nominal_freq - self.nominal_voltage = net_config.nominal_voltage - self.nominal_energy_per_link = net_config.nominal_energy_per_link - self.nominal_area_per_link = net_config.nominal_area_per_link - self.threshold_voltage = net_config.threshold_voltage - self.margin_voltage = net_config.margin_voltage - self.num_links_per_mm = net_config.num_links_per_mm - self.util = net_config.util - self.config_operating_frequency = net_config.operating_frequency # Master parameter - self.config_bandwidth = net_config.bandwidth # Master parameter - overrides all - self.inter = True if netLevel == "inter" else False - self.intra = True if netLevel == "intra" else False - - inter_frac = exp_config.perimeter_breakdown.inter_node - intra_frac = exp_config.perimeter_breakdown.intra_node - perimeter_fraction = inter_frac if self.inter else intra_frac - self.tot_perimeter = perimeter_fraction * self.core_perimeter - self.num_links = int( - min( - self.tot_area / self.nominal_area_per_link, - perimeter_fraction * self.core_perimeter * self.num_links_per_mm, - ) - ) - - self.throughput = 0 - self.operating_freq = 0 - self.operating_voltage = 0 - - # If bandwidth is explicitly set, use it directly (master parameter - overrides everything) - if self.config_bandwidth: - self.throughput = self.config_bandwidth * self.util - # Set nominal values when bandwidth is overridden (actual freq/voltage don't matter) - self.operating_freq = self.nominal_freq - self.operating_voltage = self.nominal_voltage - elif self.num_links > 0: - self.calcOperatingVoltageFrequency() - # Calculate throughput based on topology - # 4 edges for intra (mesh), 1 edge for inter (mesh extension) - degree = 4 if self.intra else 1 - self.throughput = (self.num_links * self.operating_freq * self.util) / (8 * degree) - - # self.energy_per_bit = self.calcEnergyPerBit() - - def calcOperatingVoltageFrequency(self): - # If operating_frequency is explicitly set, use it directly (master parameter) - if self.config_operating_frequency: - self.operating_freq = self.config_operating_frequency - self.operating_voltage = self.nominal_voltage # Assume nominal voltage when freq is overridden - return - - # Otherwise, calculate from voltage scaling based on power constraints - self.nominal_power = ( - self.nominal_energy_per_link * self.num_links * self.nominal_freq - ) - # minimum voltage to meet the power constraint - self.operating_voltage = self.solve_poly( - p0=1, - p1=-2 * self.threshold_voltage, - p2=self.threshold_voltage**2, - p3=-1 - * self.tot_power - / self.nominal_power - * self.nominal_voltage - * (self.nominal_voltage - self.threshold_voltage) ** 2, - ) - # operating frequency at minimum voltage - self.operating_freq = self.nominal_freq * ( - ( - (self.operating_voltage - self.threshold_voltage) ** 2 - / (self.operating_voltage) - ) - / ( - (self.nominal_voltage - self.threshold_voltage) ** 2 - / self.nominal_voltage - ) - ) - self.frequency_scaling_factor = 1 - if self.operating_voltage < (self.threshold_voltage + self.margin_voltage): - self.scaled_voltage = self.threshold_voltage + self.margin_voltage - self.frequency_scaling_factor = ( - self.operating_voltage / self.scaled_voltage - ) ** 2 + primary = self.layout.primary_dimension() + if primary is None: + return 0.0, 0.0 + latency = float(primary.latency) + return latency, latency - self.operating_freq = self.frequency_scaling_factor * self.operating_freq - - def calcEnergyPerBit(self): - self.operating_energy_per_link = ( - self.nominal_energy_per_link - * (self.operating_voltage / self.nominal_voltage) ** 2 - ) - energy_per_bit = self.operating_energy_per_link - return energy_per_bit + def get_link(self, parallelism: str) -> Tuple[float, float]: + bandwidth, latency = self.layout.link_for_parallelism(parallelism) + return float(bandwidth), float(latency) - # Return P2P bw - def calcThroughput(self): - return self.throughput - - def printStats(self, f, name): - self.calcEnergyPerBit() - self.eff_power = ( - self.num_links * self.operating_freq * self.operating_energy_per_link - ) - self.eff_area = self.num_links * self.nominal_area_per_link - self.eff_perimeter = self.num_links / self.num_links_per_mm + def dimension_for_parallelism(self, parallelism: str): + return self.layout.dimension_for_parallelism(parallelism) - if self.eff_power > 0 and self.eff_area > 0: + def printStats(self, f): + for dim in self.layout.dimensions: f.write("\n\n=============\n") - f.write("Network: {}\n".format(name)) + f.write(f"Network Dimension: {dim.label} ({dim.id})\n") f.write("=============\n") + f.write(f"Topology: {dim.topology_type}\n") + f.write(f"Size: {dim.size}\n") f.write( - "operating_voltage: {0:.2f}, operating_freq: {1:.2f} (Ghz)\n".format( - self.operating_voltage, self.operating_freq / 1e9 - ) - ) - f.write( - "voltage_lowerbound: {0:.2f}\n".format( - self.threshold_voltage + self.margin_voltage + "Effective Bandwidth: {0:.4f} (GB/s)\n".format( + dim.effective_bandwidth / 1024 / 1024 / 1024 ) ) - f.write("#links: {0:5d}\n".format(self.num_links)) - if self.tot_area != 0: - f.write( - "eff_area: {0:.2f} (mm2), tot_area: {1:.2f} (mm2), util: {2:.2f}%\n".format( - self.eff_area, - self.tot_area, - self.eff_area / self.tot_area * 100, - ) - ) - if self.tot_power != 0: - f.write( - "eff_power: {0:.2f} (watt), tot_power: {1:.2f} (watt), util: {2:.2f}%\n".format( - self.eff_power, - self.tot_power, - self.eff_power / self.tot_power * 100, - ) - ) - if self.tot_perimeter != 0: - f.write( - "eff_perimeter: {0:.2f} (mm), tot_perimeter: {1:.2f} (mm), util: {2:.2f}%\n".format( - self.eff_perimeter, - self.tot_perimeter, - self.eff_perimeter / self.tot_perimeter * 100, - ) - ) + f.write(f"Latency: {dim.latency * 1e9:.4f} (ns)\n") + f.write(f"Util: {dim.util:.4f}\n") + f.write(f"Parallelisms: {', '.join(dim.parallelisms) or 'None'}\n") diff --git a/pyproject.toml b/pyproject.toml index 0bee730a..4580b16f 100644 --- a/pyproject.toml +++ b/pyproject.toml @@ -15,8 +15,6 @@ dependencies = [ "python-dateutil==2.9.0.post0", "pytz==2025.2", "pyyaml==6.0.2", - "ruamel-yaml==0.17.40", - "ruamel-yaml-clib==0.2.8", "seaborn>=0.13.2", "six==1.17.0", "tzdata==2025.2", diff --git a/requirements.txt b/requirements.txt index 524f5a5e..746d674e 100644 --- a/requirements.txt +++ b/requirements.txt @@ -4,12 +4,10 @@ pandas==2.3.1 python-dateutil==2.9.0.post0 pytz==2025.2 pyyaml==6.0.2 -ruamel-yaml==0.17.40 -ruamel-yaml-clib==0.2.8 six==1.17.0 tzdata==2025.2 protobuf==5.28.2 pydot==3.0.0 graphviz==0.21 joblib>=1.5.2 -seaborn \ No newline at end of file +seaborn diff --git a/run_perf.py b/run_perf.py index 28020b25..7086868b 100755 --- a/run_perf.py +++ b/run_perf.py @@ -85,17 +85,16 @@ def _validate_astrasim_dependencies(hw_config) -> None: def _validate_network_topology(hw_config) -> None: backend = getattr(hw_config, "execution_backend", None) model = getattr(backend, "model", "analytical") if backend else "analytical" - network_topology = getattr(hw_config, "network_topology", None) - - if str(model).lower() == "analytical" and network_topology: - inter_topology = getattr(network_topology.inter, "topology", "ring") - intra_topology = getattr(network_topology.intra, "topology", "ring") - - if str(inter_topology).lower() != "ring" or str(intra_topology).lower() != "ring": - raise RuntimeError( - "Non-ring network topologies are not supported in analytical mode. " - "Only execution_backend.model='astra' (requires a valid AstraSim install) supports non-ring networks." - ) + network_layout = getattr(hw_config, "network_layout", None) + + if str(model).lower() == "analytical" and network_layout: + for dim in getattr(network_layout, "dimensions", ()): + topo = str(getattr(dim, "topology_type", "ring")).lower() + if topo != "ring": + raise RuntimeError( + "Non-ring network topologies are not supported in analytical mode. " + "Only execution_backend.model='astra' (requires a valid AstraSim install) supports non-ring networks." + ) def run_LSTM( exp_hw_config_path, diff --git a/time_calculation.py b/time_calculation.py index 6e05bdef..9fae111b 100644 --- a/time_calculation.py +++ b/time_calculation.py @@ -8,7 +8,6 @@ # import numpy as np from parallelism import Parallelism -from topology import Topology from simulate import Graph import util from hw_component import Core, MemoryHierarchy, Network, DRAM @@ -208,83 +207,20 @@ def __init__(self, hw_config, model_config, mode, *, astra_policy_override: Opti self.num_levels = self.memoryHierarchy.num_levels self.memLayer = self.memoryHierarchy.memLayer self.tileSpace = None - - # TODO: move this to config file self.H2Dbw = self.h2d_bandwidth - - - # System Parameters - self.num_wafer = hw_config.system_config.num_wafers - self.num_workers = hw_config.system_config.num_workers - - + self.num_workers = self._derive_num_workers(hw_config) level = 0 mem_config = hw_config.memory_hierarchy.mem_hr[level] self.DRAM = DRAM(hw_config, mem_config, level) - self.memory_capacity = self.DRAM.size * self.num_workers# in bytes - - # self.memory_capacity = hw_config.perimeter_breakdown.DRAM - self.network = Network(hw_config) - intra_throughput, inter_throughput = self.network.calcThroughput() - intra_latency, inter_latency = self.network.calcLatency() - - inter_derate = hw_config.system_config.inter_derate - intra_derate = hw_config.system_config.intra_derate - par2cross = hw_config.system_config.par2cross - - derated_inter_throughput = -1 - derated_intra_throughput = -1 - - # inter-wafer communications will pass through intra links too - if self.num_wafer > 1 and self.num_workers > 1: - if intra_derate != 0: - derated_inter_throughput = min( - intra_throughput / intra_derate, inter_throughput / inter_derate - ) - # print(f'intra_throughput / intra_derate: {intra_throughput / intra_derate}, inter_throughput / inter_derate: {inter_throughput / inter_derate}') - else: - derated_inter_throughput = inter_throughput / inter_derate - # print(f'inter_throughput / inter_derate: {inter_throughput / inter_derate}') - else: - derated_inter_throughput = 0 - - if self.num_workers > 1 and intra_derate != 0: - derated_intra_throughput = intra_throughput / intra_derate - else: - derated_intra_throughput = 0 - - self.IBK1, self.LLK1 = ( - (derated_inter_throughput, inter_latency) - if par2cross["kp1"] - else (derated_intra_throughput, intra_latency) - ) - self.IBK2, self.LLK2 = ( - (derated_inter_throughput, inter_latency) - if par2cross["kp2"] - else (derated_intra_throughput, intra_latency) - ) - self.IBD, self.LLD = ( #interconnect bandwidth and latency for data parallelism - (derated_inter_throughput, inter_latency) - if par2cross["dp"] - else (derated_intra_throughput, intra_latency) - ) - self.IBL, self.LLL = ( - (derated_inter_throughput, inter_latency) - if par2cross["lp"] - else (derated_intra_throughput, intra_latency) - ) - - self.IBTP, self.LLTP = ( - (derated_inter_throughput, inter_latency) - if par2cross["tp"] - else (derated_intra_throughput, intra_latency) - ) + self.IBK1, self.LLK1 = self.network.get_link("kp1") + self.IBK2, self.LLK2 = self.network.get_link("kp2") + self.IBD, self.LLD = self.network.get_link("dp") + self.IBL, self.LLL = self.network.get_link("lp") + self.IBTP, self.LLTP = self.network.get_link("tp") - - # Scheduling Parameters par = Parallelism(hw_config) par.findParallelStrategy() @@ -312,17 +248,17 @@ def __init__(self, hw_config, model_config, mode, *, astra_policy_override: Opti try: self.tp = int(tp_value) except (TypeError, ValueError) as exc: - raise ValueError("scheduling_param.tp must be an integer") from exc + raise ValueError("parallelism.tp must be an integer") from exc if self.tp < 1: - raise ValueError("scheduling_param.tp must be >= 1") + raise ValueError("parallelism.tp must be >= 1") cp_value = par.cp if par.cp not in (None, 0) else 1 try: self.cp = int(cp_value) except (TypeError, ValueError) as exc: - raise ValueError("scheduling_param.cp must be an integer") from exc + raise ValueError("parallelism.cp must be an integer") from exc if self.cp < 1: - raise ValueError("scheduling_param.cp must be >= 1") + raise ValueError("parallelism.cp must be >= 1") self.tp_sp = par.tp_sp @@ -338,25 +274,11 @@ def __init__(self, hw_config, model_config, mode, *, astra_policy_override: Opti if self.cp > 1: raise ValueError( "Context parallelism (cp) is not supported for LLM inference. " - "Please set scheduling_param.cp to 1 for inference runs." + "Please set parallelism.cp to 1 for inference runs." ) if self.mb > 1: print(f"[WARNING]: LLM inference configured with mb={self.mb} (>1). \n Pipeline micro-batching is ill-defined for autoregressive decode and should be avoided.") - if self.mode == "LLM": - expected_workers = self.tp * self.cp * self.dp * self.lp - label = f"tp({self.tp}) * cp({self.cp}) * dp({self.dp}) * lp({self.lp})" - else: - kp1 = int(self.kp1) if self.kp1 else 1 - kp2 = int(self.kp2) if self.kp2 else 1 - expected_workers = self.dp * self.lp * kp1 * kp2 - label = f"dp({self.dp}) * lp({self.lp}) * kp1({kp1}) * kp2({kp2})" - - if expected_workers != self.num_workers: - raise ValueError( - f"Parallelism mismatch: {label} = {expected_workers}, " - f"but system_hierarchy reports num_workers={self.num_workers}." - ) # Statistics Param self.tot_flop = 0 @@ -377,12 +299,11 @@ def __init__(self, hw_config, model_config, mode, *, astra_policy_override: Opti self.roofline, astra_policy=self._astra_policy, ) - + # Dynamically select and instantiate the model class model_class = self.get_model_class(mode) self.model = model_class(model_config) # Instantiate the model class - # Model Parameters # self.model = self.get_model_class(mode) if mode == "LSTM": @@ -432,6 +353,37 @@ def __init__(self, hw_config, model_config, mode, *, astra_policy_override: Opti if self.moe_top_k > self.moe_num_experts: raise ValueError("model_param.top_k cannot exceed model_param.num_experts") self.use_moe = self.moe_num_experts > 1 + + + def _derive_num_workers(self, hw_config) -> int: + layout = getattr(hw_config, "network_layout", None) + if layout and getattr(layout, "dimensions", None): + total = 1 + for dim in layout.dimensions: + try: + size = int(dim.size) + except (TypeError, ValueError): + size = 1 + if size < 1: + size = 1 + total *= size + if total >= 1: + return int(total) + + factors = [] + for name in ("dp", "lp", "tp", "cp", "kp1", "kp2"): + value = getattr(hw_config.sch_config, name, 1) + try: + factor = int(value) if value else 1 + except (TypeError, ValueError): + factor = 1 + if factor < 1: + factor = 1 + factors.append(factor) + total = 1 + for factor in factors: + total *= factor + return max(1, total) def get_model_class(self, model_type): diff --git a/time_calculation_LLM.py b/time_calculation_LLM.py index a142cd96..89b41017 100644 --- a/time_calculation_LLM.py +++ b/time_calculation_LLM.py @@ -132,8 +132,8 @@ def __init__(self, hw_config, model_config, mode, output_dir: Optional[str] = No mode, astra_policy_override=astra_policy, ) - if self.num_workers > 1 and self.use_moe: - raise ValueError("MoE is only supported for single-worker runs (system_hierarchy.num_workers must be 1).") + if self.num_workers > 1 and model_config.model_config.num_experts > 1: + raise ValueError("MoE is only supported for single-worker runs (total worker count must be 1).") self.output_dir = os.path.abspath(output_dir) if output_dir else os.getcwd() os.makedirs(self.output_dir, exist_ok=True) @@ -2354,23 +2354,6 @@ def calc_time_llm(self): self.zero3_ephemeral_peak_bytes ) - - - if self.debug: - print("self.tp:", self.tp) - print("Calculating LLM time...") - print( - "simulating parallelism with dp = {}, lp = {}, total data batch = {}, " - "for each dp node, data batch = {}, for each pipeline stage, data batch = {}".format( - self.dp, self.lp, self.batch_size, self.miniB, self.microB - ) - ) - print("data-parallel zero stage:", self.zero_stage) - print("total number of workers: {}".format(self.num_workers)) - print("number of workers for each data parallelism batch: {}".format(self.num_workers_dp)) - print("number of workers for each pipeline stage: {}".format(self.num_workers_lp)) - - ( pipeline_graph_obj, graph_root, diff --git a/time_calculation_inf.py b/time_calculation_inf.py index 4998aab4..5f84d759 100644 --- a/time_calculation_inf.py +++ b/time_calculation_inf.py @@ -419,7 +419,8 @@ def calc_energy(self, transformer_timings: Dict[str, OperationTiming], cross_lay energy_per_flop = self.core.nominal_energy_per_flop energy_hbm_byte = self.DRAM.dynamic_energy_per_bit * 8 - energy_comm_byte = self.network.intra_network.nominal_energy_per_link * 8 # 8 pJ/bit for interconnect + # TODO: honor per-dimension interconnect energies; currently assumes all comms use dimension 0. + energy_comm_byte = (self.network.energies_per_bit[0] if self.network.energies_per_bit else 0.0) * 8 total_energy = (total_flops * energy_per_flop) + \ (total_hbm_bytes * energy_hbm_byte) + \ diff --git a/uv.lock b/uv.lock index ae382dd3..ffb4f894 100644 --- a/uv.lock +++ b/uv.lock @@ -133,8 +133,6 @@ dependencies = [ { name = "python-dateutil" }, { name = "pytz" }, { name = "pyyaml" }, - { name = "ruamel-yaml" }, - { name = "ruamel-yaml-clib" }, { name = "seaborn" }, { name = "six" }, { name = "tzdata" }, @@ -152,8 +150,6 @@ requires-dist = [ { name = "python-dateutil", specifier = "==2.9.0.post0" }, { name = "pytz", specifier = "==2025.2" }, { name = "pyyaml", specifier = "==6.0.2" }, - { name = "ruamel-yaml", specifier = "==0.17.40" }, - { name = "ruamel-yaml-clib", specifier = "==0.2.8" }, { name = "seaborn", specifier = ">=0.13.2" }, { name = "six", specifier = "==1.17.0" }, { name = "tzdata", specifier = "==2025.2" }, @@ -689,41 +685,6 @@ wheels = [ { url = "https://files.pythonhosted.org/packages/fa/de/02b54f42487e3d3c6efb3f89428677074ca7bf43aae402517bc7cca949f3/PyYAML-6.0.2-cp313-cp313-win_amd64.whl", hash = "sha256:8388ee1976c416731879ac16da0aff3f63b286ffdd57cdeb95f3f2e085687563", size = 156446 }, ] -[[package]] -name = "ruamel-yaml" -version = "0.17.40" -source = { registry = "https://pypi.org/simple" } -dependencies = [ - { name = "ruamel-yaml-clib", marker = "python_full_version < '3.13' and platform_python_implementation == 'CPython'" }, -] -sdist = { url = "https://files.pythonhosted.org/packages/d1/d6/eb2833ccba5ea36f8f4de4bcfa0d1a91eb618f832d430b70e3086821f251/ruamel.yaml-0.17.40.tar.gz", hash = "sha256:6024b986f06765d482b5b07e086cc4b4cd05dd22ddcbc758fa23d54873cf313d", size = 137672 } -wheels = [ - { url = "https://files.pythonhosted.org/packages/35/79/5e2cffa1c77432f11cd93a5351f30732c997a239d3a3090856a72d6d8ba7/ruamel.yaml-0.17.40-py3-none-any.whl", hash = "sha256:b16b6c3816dff0a93dca12acf5e70afd089fa5acb80604afd1ffa8b465b7722c", size = 113666 }, -] - -[[package]] -name = "ruamel-yaml-clib" -version = "0.2.8" -source = { registry = "https://pypi.org/simple" } -sdist = { url = "https://files.pythonhosted.org/packages/46/ab/bab9eb1566cd16f060b54055dd39cf6a34bfa0240c53a7218c43e974295b/ruamel.yaml.clib-0.2.8.tar.gz", hash = "sha256:beb2e0404003de9a4cab9753a8805a8fe9320ee6673136ed7f04255fe60bb512", size = 213824 } -wheels = [ - { url = "https://files.pythonhosted.org/packages/b1/15/971b385c098e8d0d170893f5ba558452bb7b776a0c90658b8f4dd0e3382b/ruamel.yaml.clib-0.2.8-cp311-cp311-macosx_10_9_universal2.whl", hash = "sha256:bef08cd86169d9eafb3ccb0a39edb11d8e25f3dae2b28f5c52fd997521133069", size = 148870 }, - { url = "https://files.pythonhosted.org/packages/01/b0/4ddef56e9f703d7909febc3a421d709a3482cda25826816ec595b73e3847/ruamel.yaml.clib-0.2.8-cp311-cp311-macosx_13_0_arm64.whl", hash = "sha256:b16420e621d26fdfa949a8b4b47ade8810c56002f5389970db4ddda51dbff248", size = 134475 }, - { url = "https://files.pythonhosted.org/packages/a4/f7/22d6b620ed895a05d40802d8281eff924dc6190f682d933d4efff60db3b5/ruamel.yaml.clib-0.2.8-cp311-cp311-manylinux_2_17_x86_64.manylinux2014_x86_64.manylinux_2_24_x86_64.whl", hash = "sha256:25c515e350e5b739842fc3228d662413ef28f295791af5e5110b543cf0b57d9b", size = 544020 }, - { url = "https://files.pythonhosted.org/packages/7c/e4/0d19d65e340f93df1c47f323d95fa4b256bb28320290f5fddef90837853a/ruamel.yaml.clib-0.2.8-cp311-cp311-manylinux_2_24_aarch64.whl", hash = "sha256:1707814f0d9791df063f8c19bb51b0d1278b8e9a2353abbb676c2f685dee6afe", size = 642643 }, - { url = "https://files.pythonhosted.org/packages/c9/ff/f781eb5e2ae011e586d5426e2086a011cf1e0f59704a6cad1387975c5a62/ruamel.yaml.clib-0.2.8-cp311-cp311-musllinux_1_1_i686.whl", hash = "sha256:46d378daaac94f454b3a0e3d8d78cafd78a026b1d71443f4966c696b48a6d899", size = 695832 }, - { url = "https://files.pythonhosted.org/packages/e3/41/f62e67ac651358b8f0d60cfb12ab2daf99b1b69eeaa188d0cec809d943a6/ruamel.yaml.clib-0.2.8-cp311-cp311-musllinux_1_1_x86_64.whl", hash = "sha256:09b055c05697b38ecacb7ac50bdab2240bfca1a0c4872b0fd309bb07dc9aa3a9", size = 730923 }, - { url = "https://files.pythonhosted.org/packages/9f/f0/19ab8acbf983cd1b37f47d27ceb8b10a738d60d36316a54bad57e0d73fbb/ruamel.yaml.clib-0.2.8-cp311-cp311-win32.whl", hash = "sha256:53a300ed9cea38cf5a2a9b069058137c2ca1ce658a874b79baceb8f892f915a7", size = 99999 }, - { url = "https://files.pythonhosted.org/packages/ec/54/d8a795997921d87224c65d44499ca595a833093fb215b133f920c1062956/ruamel.yaml.clib-0.2.8-cp311-cp311-win_amd64.whl", hash = "sha256:c2a72e9109ea74e511e29032f3b670835f8a59bbdc9ce692c5b4ed91ccf1eedb", size = 118008 }, - { url = "https://files.pythonhosted.org/packages/7a/a2/eb5e9d088cb9d15c24d956944c09dca0a89108ad6e2e913c099ef36e3f0d/ruamel.yaml.clib-0.2.8-cp312-cp312-macosx_10_9_universal2.whl", hash = "sha256:ebc06178e8821efc9692ea7544aa5644217358490145629914d8020042c24aa1", size = 144636 }, - { url = "https://files.pythonhosted.org/packages/66/98/8de4f22bbfd9135deb3422e96d450c4bc0a57d38c25976119307d2efe0aa/ruamel.yaml.clib-0.2.8-cp312-cp312-macosx_13_0_arm64.whl", hash = "sha256:edaef1c1200c4b4cb914583150dcaa3bc30e592e907c01117c08b13a07255ec2", size = 135684 }, - { url = "https://files.pythonhosted.org/packages/30/d3/5fe978cd01a61c12efd24d65fa68c6f28f28c8073a06cf11db3a854390ca/ruamel.yaml.clib-0.2.8-cp312-cp312-manylinux_2_17_x86_64.manylinux2014_x86_64.manylinux_2_28_x86_64.whl", hash = "sha256:d176b57452ab5b7028ac47e7b3cf644bcfdc8cacfecf7e71759f7f51a59e5c92", size = 734571 }, - { url = "https://files.pythonhosted.org/packages/55/b3/e2531a050758b717c969cbf76c103b75d8a01e11af931b94ba656117fbe9/ruamel.yaml.clib-0.2.8-cp312-cp312-manylinux_2_24_aarch64.whl", hash = "sha256:1dc67314e7e1086c9fdf2680b7b6c2be1c0d8e3a8279f2e993ca2a7545fecf62", size = 643946 }, - { url = "https://files.pythonhosted.org/packages/0d/aa/06db7ca0995b513538402e11280282c615b5ae5f09eb820460d35fb69715/ruamel.yaml.clib-0.2.8-cp312-cp312-musllinux_1_1_i686.whl", hash = "sha256:3213ece08ea033eb159ac52ae052a4899b56ecc124bb80020d9bbceeb50258e9", size = 692169 }, - { url = "https://files.pythonhosted.org/packages/27/38/4cf4d482b84ecdf51efae6635cc5483a83cf5ca9d9c13e205a750e251696/ruamel.yaml.clib-0.2.8-cp312-cp312-musllinux_1_1_x86_64.whl", hash = "sha256:aab7fd643f71d7946f2ee58cc88c9b7bfc97debd71dcc93e03e2d174628e7e2d", size = 740325 }, - { url = "https://files.pythonhosted.org/packages/6f/67/c62c6eea53a4feb042727a3d6c18f50dc99683c2b199c06bd2a9e3db8e22/ruamel.yaml.clib-0.2.8-cp312-cp312-win32.whl", hash = "sha256:5c365d91c88390c8d0a8545df0b5857172824b1c604e867161e6b3d59a827eaa", size = 98639 }, - { url = "https://files.pythonhosted.org/packages/10/d2/52a3d810d0b5b3720725c0504a27b3fced7b6f310fe928f7019d79387bc1/ruamel.yaml.clib-0.2.8-cp312-cp312-win_amd64.whl", hash = "sha256:1758ce7d8e1a29d23de54a16ae867abd370f01b5a69e1a3ba75223eaa3ca1a1b", size = 115305 }, -] [[package]] name = "seaborn" From b79cf26b90393f70a373ff72677822820f14a9ee Mon Sep 17 00:00:00 2001 From: gkarfakis19 Date: Sun, 2 Nov 2025 05:15:15 -0800 Subject: [PATCH 08/77] minor fix --- astrasim_lib/config_generation.py | 2 +- 1 file changed, 1 insertion(+), 1 deletion(-) diff --git a/astrasim_lib/config_generation.py b/astrasim_lib/config_generation.py index 7c65321c..92e1d432 100644 --- a/astrasim_lib/config_generation.py +++ b/astrasim_lib/config_generation.py @@ -88,7 +88,7 @@ def derive_topology_from_hw(hw_obj) -> str: return "Ring" if topo_str in ("switch",): return "Switch" - if topo_str in ("torus2d", "torus"): + if topo_str in ("torus2d"): return "Torus2D" if topo_str in ("mesh",): return "Mesh" From f8834f8fe6c9cf2cb329db601ccce8c4add4013e Mon Sep 17 00:00:00 2001 From: gkarfakis19 Date: Sun, 2 Nov 2025 11:21:38 -0800 Subject: [PATCH 09/77] small config cleanup --- config.py | 2 -- configs/hardware-config/H100_SXM5_80GB.yaml | 7 ++++--- configs/hardware-config/a100_80GB.yaml | 9 +++++---- configs/hardware-config/a100_80GB_no_parallelism.yaml | 7 ++++--- configs/hardware-config/v100.yaml | 7 ++++--- configs/hardware-config/waferscale_20v100_80hbm.yaml | 7 ++++--- 6 files changed, 21 insertions(+), 18 deletions(-) diff --git a/config.py b/config.py index 71a684b2..b2b691ff 100644 --- a/config.py +++ b/config.py @@ -796,7 +796,6 @@ def use_moe(self) -> bool: "InferenceHWConfig", [ "kvcache_type", - "kvcache_fetch_overlap", ], ) @@ -1047,7 +1046,6 @@ def parse_config(filename, config_type): inference_dict = config_dict.get("inference", {}) or {} inference_cfg = InferenceHWConfig( kvcache_type=inference_dict.get("kvcache_type", "hbm_only"), - kvcache_fetch_overlap=bool(inference_dict.get("kvcache_fetch_overlap", False)), ) config = HWConfig( diff --git a/configs/hardware-config/H100_SXM5_80GB.yaml b/configs/hardware-config/H100_SXM5_80GB.yaml index 95d7ac2b..b6a642c2 100644 --- a/configs/hardware-config/H100_SXM5_80GB.yaml +++ b/configs/hardware-config/H100_SXM5_80GB.yaml @@ -100,7 +100,7 @@ network: latency: 5e-6 energy_per_bit: 8e-12 # (J/bit) util: 1.0 - collective_override: {} + collective_override: {} # overrides execution_backend.collectives when provided parallelisms: [tp, cp, lp, dp, kp1, kp2] execution_backend: @@ -108,8 +108,9 @@ execution_backend: astra: backend: astra mode: hybrid # hybrid | full_astrasim_hierarchical | full_astrasim_flattened - collective_override: - # Allowed values per op: auto | ring | direct | halvingDoubling | doubleBinaryTree + collectives: + # Default collective algorithms. Per-dimension collective_override entries can override these. + # Allowed values per op: auto | ring | direct | halvingDoubling | doubleBinaryTree | mesh | mesh2d | hypercube | torus2d # Note: leave to auto unless you have a specific reason to change it. all_gather: auto all_reduce: auto diff --git a/configs/hardware-config/a100_80GB.yaml b/configs/hardware-config/a100_80GB.yaml index 0ac5223c..d48cf3f2 100644 --- a/configs/hardware-config/a100_80GB.yaml +++ b/configs/hardware-config/a100_80GB.yaml @@ -93,7 +93,7 @@ network: latency: 5e-6 energy_per_bit: 8e-12 # (J/bit) util: 1.0 - collective_override: {} # leave empty to use 'auto' for all collectives unless an override is needed + collective_override: {} # overrides execution_backend.collectives when provided parallelisms: [tp, cp] - id: dim1 label: inter_wafer_ring_infiniband @@ -104,7 +104,7 @@ network: latency: 5e-6 energy_per_bit: 8e-12 # (J/bit) util: 1.0 - collective_override: {} # defaults to 'auto' per collective + collective_override: {} # overrides execution_backend.collectives when provided parallelisms: [lp] - id: dim2 label: wafer_dp_chains_ring_infiniband @@ -115,7 +115,7 @@ network: latency: 5e-6 energy_per_bit: 8e-12 # (J/bit) util: 0.96 - collective_override: {} # defaults to 'auto' per collective + collective_override: {} # overrides execution_backend.collectives when provided parallelisms: [dp] @@ -124,7 +124,8 @@ execution_backend: astra: backend: astra mode: full_astrasim_flattened # hybrid | full_astrasim_hierarchical | full_astrasim_flattened - collective_override: + collectives: + # Default collective algorithms. Per-dimension collective_override entries can override these. # Allowed values per op: auto | ring | direct | halvingDoubling | doubleBinaryTree | mesh | mesh2d | hypercube | torus2d # Note: leave to auto unless you have a specific reason to change it. all_gather: auto diff --git a/configs/hardware-config/a100_80GB_no_parallelism.yaml b/configs/hardware-config/a100_80GB_no_parallelism.yaml index a1cee857..ef10ba61 100644 --- a/configs/hardware-config/a100_80GB_no_parallelism.yaml +++ b/configs/hardware-config/a100_80GB_no_parallelism.yaml @@ -104,7 +104,7 @@ network: latency: 5e-6 energy_per_bit: 8e-12 # (J/bit) util: 1.0 - collective_override: {} + collective_override: {} # overrides execution_backend.collectives when provided parallelisms: [tp, cp, lp, dp, kp1, kp2] @@ -114,8 +114,9 @@ execution_backend: astra: backend: astra mode: full_astrasim_flattened # hybrid | full_astrasim_hierarchical | full_astrasim_flattened - collective_override: - # Allowed values per op: auto | ring | direct | halvingDoubling | doubleBinaryTree + collectives: + # Default collective algorithms. Per-dimension collective_override entries can override these. + # Allowed values per op: auto | ring | direct | halvingDoubling | doubleBinaryTree | mesh | mesh2d | hypercube | torus2d # Note: leave to auto unless you have a specific reason to change it. all_gather: auto all_reduce: auto diff --git a/configs/hardware-config/v100.yaml b/configs/hardware-config/v100.yaml index 9cc82ade..5f7837cf 100644 --- a/configs/hardware-config/v100.yaml +++ b/configs/hardware-config/v100.yaml @@ -104,7 +104,7 @@ network: latency: 5e-6 energy_per_bit: 8e-12 # (J/bit) util: 1.0 - collective_override: {} + collective_override: {} # overrides execution_backend.collectives when provided parallelisms: [tp, cp, lp, dp, kp1, kp2] @@ -113,8 +113,9 @@ execution_backend: astra: backend: astra mode: hybrid # hybrid | full_astrasim_hierarchical | full_astrasim_flattened - collective_override: - # Allowed values per op: auto | ring | direct | halvingDoubling | doubleBinaryTree + collectives: + # Default collective algorithms. Per-dimension collective_override entries can override these. + # Allowed values per op: auto | ring | direct | halvingDoubling | doubleBinaryTree | mesh | mesh2d | hypercube | torus2d # Note: leave to auto unless you have a specific reason to change it. all_gather: auto all_reduce: auto diff --git a/configs/hardware-config/waferscale_20v100_80hbm.yaml b/configs/hardware-config/waferscale_20v100_80hbm.yaml index c9f6876e..b7d00522 100644 --- a/configs/hardware-config/waferscale_20v100_80hbm.yaml +++ b/configs/hardware-config/waferscale_20v100_80hbm.yaml @@ -121,7 +121,7 @@ network: latency: 5e-6 energy_per_bit: 8e-12 # (J/bit) util: 0.95 - collective_override: {} + collective_override: {} # overrides execution_backend.collectives when provided parallelisms: [tp, cp, lp, dp] execution_backend: @@ -129,8 +129,9 @@ execution_backend: astra: backend: astra mode: hybrid # hybrid | full_astrasim_hierarchical | full_astrasim_flattened - collective_override: - # Allowed values per op: auto | ring | direct | halvingDoubling | doubleBinaryTree + collectives: + # Default collective algorithms. Per-dimension collective_override entries can override these. + # Allowed values per op: auto | ring | direct | halvingDoubling | doubleBinaryTree | mesh | mesh2d | hypercube | torus2d # Note: leave to auto unless you have a specific reason to change it. all_gather: auto all_reduce: auto From 9dc7656366b4a7805cbb1f57a57ca928954fb3e5 Mon Sep 17 00:00:00 2001 From: gkarfakis19 Date: Mon, 3 Nov 2025 00:04:06 -0800 Subject: [PATCH 10/77] Major reformat of network configs to enable N-dimensional astrasim exec.. Includes some flattener bug fixes, cleanup and documentation in executor.py --- LLM_execution.py | 197 +++++++++++- astrasim_lib/config_generation.py | 246 +++++++++++--- astrasim_lib/executor.py | 303 ++++++++++++++---- astrasim_lib/integration.py | 11 +- .../hardware-config/a100_80GB_problem.yaml | 141 ++++++++ configs/model-config/LLM_problem.yaml | 19 ++ examples/problem.sh | 7 + simulate_LLM.py | 13 +- 8 files changed, 817 insertions(+), 120 deletions(-) create mode 100644 configs/hardware-config/a100_80GB_problem.yaml create mode 100644 configs/model-config/LLM_problem.yaml create mode 100755 examples/problem.sh diff --git a/LLM_execution.py b/LLM_execution.py index fa3e8b73..ed4d15c5 100644 --- a/LLM_execution.py +++ b/LLM_execution.py @@ -48,6 +48,8 @@ def __init__( self, pipeline_graph: Graph, transformer_graph: Graph, + *, + rank_layout: Optional[Dict[str, Any]] = None, ) -> None: if transformer_graph is None: raise ValueError("Transformer graph is required for flattening") @@ -63,11 +65,50 @@ def __init__( par_degree = transformer_graph.tp * transformer_graph.cp self._par_degree = max(1, int(par_degree)) self._zero_stage = int(getattr(transformer_graph, "misc_metadata", {}).get("dp_zero_stage", 0)) + self._layout_axis_order: Optional[List[str]] = None + self._layout_axis_sizes: Dict[str, int] = {} + self._layout_axis_strides: Dict[str, int] = {} + self._tp_size = int(getattr(transformer_graph, "tp", 1)) + self._cp_size = int(getattr(transformer_graph, "cp", 1)) + self._lp_size = int(getattr(pipeline_graph, "lp", 1)) + if rank_layout is not None: + self._configure_rank_layout(rank_layout) # Track original ZeRO-3 transformer gather edges -> per-rank clones self._clone_cache: Dict[int, Any] = {} self._op_id_counter: int = 0 - + self._stage_span = 1 + + def _configure_rank_layout(self, descriptor: Dict[str, Any]) -> None: + axis_order = list(descriptor.get("axis_order", [])) + axis_sizes = dict(descriptor.get("axis_sizes", {})) + axis_strides = dict(descriptor.get("axis_strides", {})) + if not axis_order: + self._layout_axis_order = None + return + self._layout_axis_order = axis_order + self._layout_axis_sizes = axis_sizes + if not axis_strides: + span = 1 + for axis in axis_order: + axis_strides[axis] = span + span *= axis_sizes.get(axis, 1) + self._layout_axis_strides = axis_strides + self._tp_size = max(1, axis_sizes.get("tp", self._tp_size)) + self._cp_size = max(1, axis_sizes.get("cp", self._cp_size)) + self._lp_size = max(1, axis_sizes.get("lp", self._lp_size)) + if self._par_degree != self._tp_size * self._cp_size: + raise ValueError( + f"Inconsistent tensor/context parallel factors: tp={self._tp_size}, cp={self._cp_size}, " + f"product does not equal par_degree={self._par_degree}" + ) + stage_span = 1 + for axis in axis_order: + if axis == "dp": + continue + stage_span *= axis_sizes.get(axis, 1) + self._stage_span = stage_span + def _should_shard_zero3_transformer(self, edge: Any) -> bool: if self._par_degree <= 1 or self._zero_stage < 3: return False @@ -240,6 +281,52 @@ def _clone(self, obj: Any) -> Any: raise TypeError(f"Unsupported graph element type: {type(obj)!r}") + def _propagate_local_hw_ids(self, roots: Any) -> None: + stack: List[Any] = [] + if isinstance(roots, (list, tuple)): + stack.extend(list(roots)) + elif roots is not None: + stack.append(roots) + visited: Set[int] = set() + while stack: + obj = stack.pop() + obj_id = id(obj) + if obj_id in visited: + continue + visited.add(obj_id) + + children = getattr(obj, "children", []) + if isinstance(children, (list, tuple)): + stack.extend(children) + elif children is not None: + stack.append(children) + + parents = getattr(obj, "parents", None) + if parents: + if isinstance(parents, (list, tuple)): + stack.extend(parents) + else: + stack.append(parents) + + if isinstance(obj, simulate_LLM.Edge): + new_hw = None + for parent in getattr(obj, "parents", []) or []: + parent_hw = getattr(parent, "hw_id", None) + if parent_hw is None: + parent_hw = getattr(parent, "local_hw_id", None) + if parent_hw is not None and parent_hw >= 0: + new_hw = parent_hw + break + if new_hw is None: + stage_id = getattr(obj, "stage_id", None) + if stage_id is not None: + try: + new_hw = self._hw_id_for_rank(stage_id, 0) + except Exception: + new_hw = None + if new_hw is not None and new_hw >= 0: + obj.local_hw_id = new_hw + def _expand_transformer_node(self, node: simulate_LLM.Node) -> Tuple[Any, ...]: node_id = id(node) @@ -501,7 +588,33 @@ def _next_op_id(self) -> int: def _hw_id_for_rank(self, stage_id: int, tp_rank: int) -> int: stage_int = int(stage_id) if stage_id is not None else 0 - return stage_int * self._par_degree + tp_rank + tp_rank_int = int(tp_rank) + if tp_rank_int < 0 or tp_rank_int >= self._par_degree: + raise ValueError(f"tp_rank {tp_rank_int} is out of range for par_degree {self._par_degree}") + if not self._layout_axis_order: + return stage_int * self._par_degree + tp_rank_int + + coords: Dict[str, int] = {} + if "tp" in self._layout_axis_order: + coords["tp"] = tp_rank_int % self._tp_size + if "cp" in self._layout_axis_order: + coords["cp"] = (tp_rank_int // self._tp_size) % self._cp_size + if "lp" in self._layout_axis_order: + if stage_int < 0 or stage_int >= self._lp_size: + raise ValueError(f"stage_id {stage_int} is out of range for lp={self._lp_size}") + coords["lp"] = stage_int % self._lp_size + + linear_rank = 0 + for axis in self._layout_axis_order: + coord = coords.get(axis, 0) + size = self._layout_axis_sizes.get(axis, 1) + if coord < 0 or coord >= size: + raise ValueError(f"Coordinate {coord} for axis '{axis}' is out of range <{size}") + stride = self._layout_axis_strides.get(axis) + if stride is None: + raise KeyError(f"Rank layout stride missing for axis '{axis}'") + linear_rank += coord * stride + return linear_rank class LLMExecutionDispatcher: @@ -523,6 +636,80 @@ def __init__( self.transformer_forward_root = transformer_forward_root self.transformer_backward_root = transformer_backward_root self.flattened_root: Optional[Any] = None + self._rank_layout = self._build_rank_layout_descriptor() + + def _build_rank_layout_descriptor(self) -> Dict[str, Any]: + hw_config = getattr(self.time_calc, "hw_config", None) + layout = getattr(hw_config, "network_layout", None) + dimensions = getattr(layout, "dimensions", None) if layout is not None else None + if not dimensions: + return {} + + def _safe_int(value: Any, default: int = 1) -> int: + try: + candidate = int(value) + except (TypeError, ValueError): + candidate = default + return max(1, candidate) + + tp_size = _safe_int(getattr(self.pipeline_graph, "tp", getattr(self.time_calc, "tp", 1))) + cp_size = _safe_int(getattr(self.pipeline_graph, "cp", getattr(self.time_calc, "cp", 1))) + lp_size = _safe_int(getattr(self.pipeline_graph, "lp", getattr(self.time_calc, "lp", 1))) + dp_size = _safe_int(getattr(self.time_calc, "dp", 1)) + + axis_sizes: Dict[str, int] = {"tp": tp_size, "cp": cp_size, "lp": lp_size, "dp": dp_size} + axis_order: List[str] = [] + for dim in dimensions: + dim_axes = [str(axis).strip().lower() for axis in getattr(dim, "parallelisms", ())] + for name in dim_axes: + if name == "dp": + continue + if name not in axis_sizes: + raise ValueError( + f"Unsupported parallelism axis '{name}' in network layout. " + "Supported axes for AstraSim integration are: tp, cp, lp, dp." + ) + if name not in axis_order: + axis_order.append(name) + + declared = int(getattr(dim, "size", 1)) + expected = 1 + for axis_name in dim_axes: + expected *= axis_sizes.get(axis_name, 1) + if expected != declared: + raise ValueError( + f"Network dimension '{getattr(dim, 'label', getattr(dim, 'id', ''))}' " + f"size mismatch: declared {declared}, but parallelism factors imply {expected}." + ) + + # Ensure the layout covers active parallel axes + if tp_size > 1 and "tp" not in axis_order: + raise ValueError("Network layout must include 'tp' when tensor parallelism > 1.") + if cp_size > 1 and "cp" not in axis_order: + raise ValueError("Network layout must include 'cp' when context parallelism > 1.") + if lp_size > 1 and "lp" not in axis_order: + raise ValueError("Network layout must include 'lp' when pipeline parallelism > 1.") + + axis_strides: Dict[str, int] = {} + span = 1 + for axis in axis_order: + axis_strides[axis] = span + span *= axis_sizes[axis] + + descriptor = { + "axis_order": axis_order, + "axis_sizes": axis_sizes, + "axis_strides": axis_strides, + "stage_span": span, + } + pretty_sizes = ", ".join(f"{axis}={axis_sizes[axis]}" for axis in axis_order) + dp_msg = f", dp={dp_size}" if dp_size > 1 else "" + print( + "[DeepFlow][hw-id] Flattening axes for hw_id assignment: " + f"{axis_order or ['']} ({pretty_sizes or 'single device'}{dp_msg}) | " + "dp ranks are appended separately." + ) + return descriptor def run(self, mode: ExecutionMode) -> ExecutionResult: if mode == ExecutionMode.ANALYTICAL: @@ -608,6 +795,7 @@ def _run_full_astrasim_hierarchical(self) -> ExecutionResult: "/pipeline_graph_hierarchical", ) + setattr(self.pipeline_root, "_astrasim_rank_layout", self._rank_layout) per_rank_sec, max_sec = run_astra_simulation_only_onepath( self.pipeline_root, self.time_calc, @@ -629,6 +817,7 @@ def _run_full_astrasim_flattened(self) -> ExecutionResult: flattener = PipelineGraphFlattener( pipeline_graph=self.pipeline_graph, transformer_graph=self.transformer_graph, + rank_layout=self._rank_layout, ) if _env_flag("DEEPFLOW_VISUALIZE_GRAPHS") and self.pipeline_root is not None: @@ -641,6 +830,8 @@ def _run_full_astrasim_flattened(self) -> ExecutionResult: if flattened_root is None: raise RuntimeError("Pipeline flattening produced an empty graph") + flattener._propagate_local_hw_ids(flattened_root) + setattr(flattened_root, "_astrasim_rank_layout", self._rank_layout) self.time_calc.flattened_pipeline_root = flattened_root if _env_flag("DEEPFLOW_VISUALIZE_GRAPHS") and self.pipeline_root is not None: self.pipeline_graph.save_graph( @@ -654,7 +845,7 @@ def _run_full_astrasim_flattened(self) -> ExecutionResult: # base_path = os.path.join(output_dir, "pipeline_flattened") # dot = visualize_graph(flattened_root, filename=base_path) # try: - # dot.render(base_path, format="png", cleanup=True) + # dot.render(base_path, format="svg", cleanup=True) # except Exception as exc: # pragma: no cover - visualization best-effort # print(f"[WARN] Failed to render flattened pipeline graph: {exc}") diff --git a/astrasim_lib/config_generation.py b/astrasim_lib/config_generation.py index 92e1d432..699af7cc 100644 --- a/astrasim_lib/config_generation.py +++ b/astrasim_lib/config_generation.py @@ -3,7 +3,7 @@ from __future__ import annotations import os -from typing import Any, Dict, Optional, Tuple +from typing import Any, Dict, List, Optional, Tuple from hw_component import Network @@ -16,14 +16,14 @@ ASTRA_DEBUG = False -_NET_YAML_CACHE: set[tuple[str, int, float, float, str]] = set() +_NET_YAML_CACHE: Dict[tuple, str] = {} _JSON_WRITTEN_BY_NPUS: set[object] = set() def _save_json(path: str, data: Dict[str, Any], npus_key: Optional[int] = None) -> None: """Write ``data`` to ``path`` once per ``npus_key`` per process.""" key = npus_key if npus_key is not None else path - if key in _JSON_WRITTEN_BY_NPUS: + if key in _JSON_WRITTEN_BY_NPUS and os.path.exists(path): return os.makedirs(os.path.dirname(path), exist_ok=True) tmp_path = path + ".tmp" @@ -81,6 +81,10 @@ def derive_topology_from_hw(hw_obj) -> str: layout = getattr(hw_obj, "network_layout", None) primary = layout.primary_dimension() if layout else None topo = (primary.topology_type if primary else "ring") or "ring" + return _normalize_topology_name(topo) + + +def _normalize_topology_name(topo: str) -> str: topo_str = str(topo).lower() if topo_str in ("fc", "fullyconnected", "fully_connected", "fully-connected"): return "FullyConnected" @@ -110,72 +114,202 @@ def generate_astrasim_configs_from_hw( if npus_count is None: raise ValueError("npus_count must be provided explicitly when generating AstraSim configs.") - net_yaml = os.path.join(out_dir, f"network_analytical_{int(npus_count)}.yml") - sys_json = os.path.join(out_dir, "system_native_collectives.json") - - topo = derive_topology_from_hw(hw_obj) - if topo == "Ring" and int(npus_count) <= 2: - topo = "FullyConnected" + layout = getattr(hw_obj, "network_layout", None) + dimensions = list(getattr(layout, "dimensions", [])) if layout else [] + if not dimensions: + raise ValueError("Hardware config is missing network dimensions required for AstraSim integration.") - (intra_ib_bps, intra_ll_s), _ = compute_intra_inter_ib_ll_from_hw(hw_obj) - if not intra_ib_bps or intra_ib_bps <= 0: - raise ValueError("Intra-node bandwidth computed as 0. Check hardware config/network model.") - if not intra_ll_s or intra_ll_s <= 0: - raise ValueError("Intra-node latency computed as 0. Check hardware config/network model.") - ib_gbps = round(_gbps_from_bps(intra_ib_bps), 6) - ll_ns = round(_ns_from_s(intra_ll_s), 3) + sch_config = getattr(hw_obj, "sch_config", None) + def _safe_int(value: Any, default: int = 1) -> int: + try: + candidate = int(value) + except (TypeError, ValueError): + candidate = default + return max(1, candidate) + + axis_sizes: Dict[str, int] = { + "tp": _safe_int(getattr(sch_config, "tp", 1)) if sch_config else 1, + "cp": _safe_int(getattr(sch_config, "cp", 1)) if sch_config else 1, + "lp": _safe_int(getattr(sch_config, "lp", 1)) if sch_config else 1, + "dp": _safe_int(getattr(sch_config, "dp", 1)) if sch_config else 1, + } + axis_order_preference = ["tp", "cp", "lp", "dp"] + + dim_infos: List[Tuple[Any, List[str], int]] = [] + for dim in dimensions: + axes = [str(axis).strip().lower() for axis in getattr(dim, "parallelisms", ())] + effective = 1 + for axis in axes: + if axis not in axis_sizes: + raise ValueError( + f"Unsupported parallelism axis '{axis}' referenced by network dimension '{dim.label}'. " + "Supported axes are tp, cp, lp, dp." + ) + effective *= axis_sizes[axis] + if effective != int(getattr(dim, "size", effective)): + raise ValueError( + f"Network dimension '{dim.label}' size mismatch: declared {dim.size}, " + f"but parallelism factors imply {effective}." + ) + dim_infos.append((dim, axes, effective)) + + active_dim_indices = [idx for idx, (dim, _, _) in enumerate(dim_infos) if int(getattr(dim, "size", 0)) > 1] + if active_dim_indices: + last_active_idx = active_dim_indices[-1] + dp_in_active = False + for idx, (dim, axes, _) in enumerate(dim_infos): + if "dp" in axes: + if idx != last_active_idx: + raise ValueError( + f"Data-parallel axis must be assigned to the last active network dimension. " + f"Dimension '{dim.label}' (index {idx}) includes 'dp', but last active dimension index is {last_active_idx}." + ) + dp_in_active = True + if axis_sizes.get("dp", 1) > 1 and not dp_in_active: + raise ValueError( + "Hardware config declares dp > 1, but no network dimension with size > 1 includes the 'dp' axis." + ) + + target = int(npus_count) + axes_needed: List[str] = [] + remaining = target + for axis in axis_order_preference: + size = axis_sizes.get(axis, 1) + if size <= 1: + continue + if remaining % size == 0: + axes_needed.append(axis) + remaining //= size + if remaining != 1: + raise ValueError( + f"Unable to map requested npus_count={target} to network axes {axis_sizes}." + ) + + axes_needed_set = set(axes_needed) + selected_dims: List[Tuple[Any, List[str], int]] = [] + accumulated = 1 + for dim, axes, _effective in dim_infos: + selected_axes: List[str] = [] + size_contrib = 1 + for axis in axes: + if axis in axes_needed_set: + selected_axes.append(axis) + size_contrib *= axis_sizes[axis] + if not selected_axes: + continue + axes_needed_set.difference_update(selected_axes) + accumulated *= size_contrib + selected_dims.append((dim, selected_axes, size_contrib)) + if axes_needed_set: + raise ValueError( + f"Requested npus_count={target} requires axes {axes_needed} but no network dimension covers {axes_needed_set}." + ) + if accumulated != target: + raise ValueError( + f"Selected network dimensions {[d.label for d, _, _ in selected_dims]} " + f"describe {accumulated} NPUs but simulation expects {target} ranks." + ) + + topo_list: List[str] = [] + npus_list: List[int] = [] + bw_list: List[float] = [] + lat_list: List[float] = [] + + for dim, axes_selected, product_size in selected_dims: + topo = _normalize_topology_name(dim.topology_type) + size = product_size + if topo == "Ring" and size <= 2: + topo = "FullyConnected" + effective_bw = float(getattr(dim, "effective_bandwidth", dim.bandwidth)) + latency_s = float(dim.latency) + topo_list.append(topo) + npus_list.append(size) + bw_list.append(round(_gbps_from_bps(effective_bw), 6)) + lat_list.append(round(_ns_from_s(latency_s), 3)) + + signature_parts = [str(size) for _dim, _axes, size in selected_dims] + dim_signature = "_".join(signature_parts) if signature_parts else f"{target}" + + net_yaml = os.path.join(out_dir, f"network_analytical_{dim_signature}.yml") + sys_json = os.path.join(out_dir, f"system_native_collectives_{dim_signature}.json") + + topo_str = ", ".join(topo_list) + npus_str = ", ".join(str(v) for v in npus_list) + bw_str = ", ".join(str(v) for v in bw_list) + lat_str = ", ".join(str(v) for v in lat_list) net_content = ( - f"topology: [ {topo} ]\n" - f"npus_count: [ {int(npus_count)} ]\n" - f"bandwidth: [ {ib_gbps} ] # GB/s\n" - f"latency: [ {ll_ns} ] # ns\n" + f"topology: [ {topo_str} ]\n" + f"npus_count: [ {npus_str} ]\n" + f"bandwidth: [ {bw_str} ] # GB/s\n" + f"latency: [ {lat_str} ] # ns\n" ) os.makedirs(os.path.dirname(net_yaml), exist_ok=True) - cache_key = (net_yaml, int(npus_count), ib_gbps, ll_ns, topo) + cache_key = (tuple(topo_list), tuple(npus_list), tuple(bw_list), tuple(lat_list)) + cached_path = _NET_YAML_CACHE.get(cache_key) need_write = True - if cache_key in _NET_YAML_CACHE: - need_write = False - elif os.path.exists(net_yaml): + if cached_path and os.path.exists(cached_path): try: - with open(net_yaml, "r", encoding="utf-8") as handle: + with open(cached_path, "r", encoding="utf-8") as handle: existing = handle.read() if existing == net_content: + net_yaml = cached_path need_write = False - except Exception: # pragma: no cover - defensive - pass + except Exception: + need_write = True if need_write: tmp_path = net_yaml + ".tmp" with open(tmp_path, "w", encoding="utf-8") as handle: handle.write(net_content) os.replace(tmp_path, net_yaml) - _NET_YAML_CACHE.add(cache_key) + _NET_YAML_CACHE[cache_key] = net_yaml exec_backend = getattr(hw_obj, "execution_backend", None) if exec_backend and exec_backend.astra: coll = exec_backend.astra.collectives sys_opts = getattr(exec_backend.astra, "sys_options", None) - ag = choose_collective(coll.all_gather, topo, "all-gather") - ar = choose_collective(coll.all_reduce, topo, "all-reduce") - rs = choose_collective(coll.reduce_scatter, topo, "reduce-scatter") - a2a = choose_collective(coll.all_to_all, topo, "all-to-all") + default_ag = coll.all_gather + default_ar = coll.all_reduce + default_rs = coll.reduce_scatter + default_a2a = coll.all_to_all else: - ag = choose_collective("auto", topo, "all-gather") - ar = choose_collective("auto", topo, "all-reduce") - rs = choose_collective("auto", topo, "reduce-scatter") - a2a = choose_collective("auto", topo, "all-to-all") + default_ag = "auto" + default_ar = "auto" + default_rs = "auto" + default_a2a = "auto" sys_opts = None + def _collective_for_dimension(dim, topo_name: str, op: str, default_alg: str) -> str: + override = ( + dim.collective_override.get(op) + or dim.collective_override.get(op.replace("-", "_")) + or dim.collective_override.get(op.replace("_", "-")) + ) + if override: + return override + return choose_collective(default_alg, topo_name, op) + + ag_impl: List[str] = [] + ar_impl: List[str] = [] + rs_impl: List[str] = [] + a2a_impl: List[str] = [] + + for (dim, axes, _), topo_name in zip(selected_dims, topo_list): + ag_impl.append(_collective_for_dimension(dim, topo_name, "all-gather", default_ag)) + ar_impl.append(_collective_for_dimension(dim, topo_name, "all-reduce", default_ar)) + rs_impl.append(_collective_for_dimension(dim, topo_name, "reduce-scatter", default_rs)) + a2a_impl.append(_collective_for_dimension(dim, topo_name, "all-to-all", default_a2a)) + system = { "scheduling-policy": "LIFO", "endpoint-delay": 10, "active-chunks-per-dimension": 1, "preferred-dataset-splits": 1, - "all-reduce-implementation": [ar], - "all-gather-implementation": [ag], - "reduce-scatter-implementation": [rs], - "all-to-all-implementation": [a2a], + "all-reduce-implementation": ar_impl, + "all-gather-implementation": ag_impl, + "reduce-scatter-implementation": rs_impl, + "all-to-all-implementation": a2a_impl, "collective-optimization": "localBWAware", "local-mem-bw": 1600, "boost-mode": 0, @@ -186,13 +320,35 @@ def generate_astrasim_configs_from_hw( if getattr(sys_opts, "endpoint_delay", None) is not None: system["endpoint-delay"] = sys_opts.endpoint_delay if getattr(sys_opts, "active_chunks_per_dimension", None) is not None: - system["active-chunks-per-dimension"] = sys_opts.active_chunks_per_dimension + acpd = sys_opts.active_chunks_per_dimension + if isinstance(acpd, (list, tuple)): + acpd = acpd[0] if acpd else 1 + system["active-chunks-per-dimension"] = int(acpd) if getattr(sys_opts, "preferred_dataset_splits", None) is not None: - system["preferred-dataset-splits"] = sys_opts.preferred_dataset_splits - - _save_json(sys_json, system, npus_key=int(npus_count)) + pds = sys_opts.preferred_dataset_splits + if isinstance(pds, (list, tuple)): + pds = pds[0] if pds else 1 + system["preferred-dataset-splits"] = int(pds) + + _save_json( + sys_json, + system, + npus_key=( + int(npus_count), + dim_signature, + tuple(ar_impl), + tuple(ag_impl), + tuple(rs_impl), + tuple(a2a_impl), + ), + ) - return {"network_yaml": net_yaml, "system_json": sys_json} + return { + "network_yaml": net_yaml, + "system_json": sys_json, + "topology_list": topo_list, + "npus_per_dim": npus_list, + } __all__ = [ diff --git a/astrasim_lib/executor.py b/astrasim_lib/executor.py index 57f71ebe..1932bbfa 100644 --- a/astrasim_lib/executor.py +++ b/astrasim_lib/executor.py @@ -20,7 +20,7 @@ import time from collections import defaultdict, deque from itertools import count as _it_count -from typing import Any, Dict, List, Optional, Set, Tuple +from typing import Any, Dict, List, Mapping, Optional, Sequence, Set, Tuple sys.setrecursionlimit(10000) @@ -64,11 +64,17 @@ def _clean_astrasim_artifacts(directory: str) -> None: os.remove(path) except OSError: pass - elif entry in { - "manifest.json", - "system_native_collectives.json", - "comm_groups.json", - }: + elif entry == "manifest.json": + try: + os.remove(path) + except OSError: + pass + elif entry == "comm_groups.json" or entry.startswith("comm_groups_"): + try: + os.remove(path) + except OSError: + pass + elif entry.startswith("system_native_collectives"): try: os.remove(path) except OSError: @@ -86,7 +92,137 @@ def _clean_astrasim_artifacts(directory: str) -> None: _TP_GROUP_BASE_ID = 1000 _LAST_TP_GROUPS: Optional[Dict[str, List[int]]] = None +_LAST_STAGE_GROUPS: Optional[List[List[int]]] = None _TP_LABEL_DP_TO_ID: Dict[Tuple[str, int], str] = {} +_TP_MEMBERS_TO_GID: Dict[Tuple[int, ...], str] = {} + + +def _extract_axis_layout(rank_layout: Optional[Dict[str, Any]]) -> Tuple[List[str], Dict[str, int]]: + if not isinstance(rank_layout, dict): + return [], {} + axis_order = list(rank_layout.get("axis_order", [])) + raw_sizes = rank_layout.get("axis_sizes", {}) + if isinstance(axis_order, tuple): + axis_order = list(axis_order) + axis_sizes: Dict[str, int] = {} + if isinstance(raw_sizes, dict): + for key, value in raw_sizes.items(): + try: + axis_sizes[str(key)] = max(1, int(value)) + except Exception: + axis_sizes[str(key)] = 1 + return axis_order, axis_sizes + + +def _compute_stage_axis_coords( + stage_ids: Sequence[int], + axis_order: Sequence[str], + axis_sizes: Mapping[str, int], +) -> Dict[int, Dict[str, int]]: + coords: Dict[int, Dict[str, int]] = {} + if not axis_order: + return coords + for stage in stage_ids: + remaining = int(stage) + stage_coords: Dict[str, int] = {} + for axis in axis_order: + size = max(1, int(axis_sizes.get(axis, 1))) + if size > 1: + stage_coords[axis] = remaining % size + remaining //= size + else: + stage_coords[axis] = 0 + coords[stage] = stage_coords + return coords + + +def _build_axis_groups( + axis_order: Sequence[str], + axis_sizes: Mapping[str, int], + stage_axis_coords: Mapping[int, Dict[str, int]], + stage_to_ranks: Mapping[int, List[int]], +) -> Dict[str, Dict[Tuple[int, Tuple[Tuple[str, int], ...]], List[int]]]: + axis_groups: Dict[ + str, Dict[Tuple[int, Tuple[Tuple[str, int], ...]], List[int]] + ] = {} + if not axis_order: + return axis_groups + for axis in axis_order: + size = max(1, int(axis_sizes.get(axis, 1))) + if size <= 1: + continue + groups: Dict[ + Tuple[int, Tuple[Tuple[str, int], ...]], List[int] + ] = defaultdict(list) + for stage, coords in stage_axis_coords.items(): + key_base = tuple((ax, coords.get(ax, 0)) for ax in axis_order if ax != axis) + ranks = stage_to_ranks.get(stage, []) + for dp_idx, rank in enumerate(ranks): + groups[(dp_idx, key_base)].append(rank) + axis_groups[axis] = { + key: sorted(set(values)) for key, values in groups.items() if values + } + return axis_groups + + +def _assign_collective_labels( + tp_collective_groups: Mapping[str, List[Any]], + collective_info: Mapping[Any, Dict[str, Any]], + axis_order: Sequence[str], + axis_sizes: Mapping[str, int], + stage_axis_coords: Mapping[int, Dict[str, int]], + axis_groups: Mapping[str, Dict[Tuple[int, Tuple[Tuple[str, int], ...]], List[int]]], + stage_to_ranks: Mapping[int, List[int]], + dp_count: int, +) -> Tuple[Dict[Any, str], Dict[str, List[Any]], Dict[str, Set[Tuple[str, int, Tuple[int, ...]]]]]: + tp_collective_labels: Dict[Any, str] = {} + label_to_edges: Dict[str, List[Any]] = {} + label_group_tokens: Dict[str, Set[Tuple[str, int, Tuple[int, ...]]]] = defaultdict(set) + group_key_to_label: Dict[Tuple[str, Tuple[int, ...]], str] = {} + label_suffix_counter: Dict[str, int] = defaultdict(int) + + for base_name, edges in tp_collective_groups.items(): + if not edges: + continue + for edge in sorted(edges, key=lambda e: getattr(e, "op_id", 0)): + info = collective_info[edge] + axis = str(info.get("interconnect_type", "tp")).lower() + stage = info["stage"] + coords = stage_axis_coords.get(stage, {}) + key_base = tuple((ax, coords.get(ax, 0)) for ax in axis_order if ax != axis) + + members_per_dp: List[Tuple[int, Tuple[int, ...]]] = [] + for dp_idx in range(dp_count): + members = None + axis_map = axis_groups.get(axis) + if axis_map is not None: + members = axis_map.get((dp_idx, key_base)) + if not members: + ranks = stage_to_ranks.get(stage, []) + if dp_idx < len(ranks): + members = [ranks[dp_idx]] + if members: + members_tuple = tuple(sorted(members)) + members_per_dp.append((dp_idx, members_tuple)) + + if not members_per_dp: + continue + + primary_members = members_per_dp[0][1] + label_key = (base_name, primary_members) + label = group_key_to_label.get(label_key) + if label is None: + suffix = label_suffix_counter[base_name] + label = base_name if suffix == 0 else f"{base_name}_{suffix}" + label_suffix_counter[base_name] += 1 + group_key_to_label[label_key] = label + + tp_collective_labels[edge] = label + label_to_edges.setdefault(label, []).append(edge) + for dp_idx, members_tuple in members_per_dp: + label_group_tokens[label].add((axis, dp_idx, members_tuple)) + + return tp_collective_labels, label_to_edges, label_group_tokens class _RankTrace: @@ -203,6 +339,9 @@ def _visualize_et_files(et_paths: List[str]) -> None: if not et_paths: return + # prune et_paths to only include the first 10 + et_paths = et_paths[:20] + def _render_et_file(et_path: str) -> None: try: fh = chakra_open(et_path) @@ -224,8 +363,7 @@ def _render_et_file(et_path: str) -> None: nodes.append(node) fh.close() - dot = Digraph(comment=os.path.basename(et_path)) - dot.graph_attr.update({"rankdir": "TB", "fontsize": "10"}) + dot = Digraph(comment=os.path.basename(et_path), format="svg",graph_attr={"rankdir": "TB", "fontsize": "10"}) id_to_node = {int(node.id): node for node in nodes} type_color = { @@ -271,11 +409,11 @@ def _render_et_file(et_path: str) -> None: dir_name, base_name = os.path.split(et_path) viz_base = base_name + ".viz" try: - output_path = dot.render(viz_base, directory=dir_name or None, format="png", cleanup=True) - final_png = et_path + ".png" - if output_path and output_path != final_png: + output_path = dot.render(viz_base, directory=dir_name or None, format="svg", cleanup=True) + final_svg = et_path + ".svg" + if output_path and output_path != final_svg: try: - os.replace(output_path, final_png) + os.replace(output_path, final_svg) except FileNotFoundError: # Some graphviz versions return path without creating file when empty graph pass @@ -283,7 +421,7 @@ def _render_et_file(et_path: str) -> None: print(f"[WARN] Failed to render Graphviz graph for {et_path}: {exc}") for et_path in et_paths: - message = f" | ET Graph saved to {et_path}.png" + message = f" | ET Graph saved to {et_path}.svg" graphviz_async.submit( f"et:{os.path.basename(et_path)}", _render_et_file, @@ -367,7 +505,7 @@ def fmt_attr_map(node: pb.Node) -> str: try: with open(out_path, "w") as outf: outf.write("\n".join(lines) + "\n") - print(f"[AstraSim] Saved ET text dump to {out_path}") + # print(f"[AstraSim] Saved ET text dump to {out_path}") except OSError as exc: print(f"[WARN] Failed to write ET text dump for {et_path}: {exc}") @@ -423,12 +561,32 @@ def convert_deepflow_graph_to_chakra_et( output_dir: str, ) -> Tuple[str, List[int], str]: """Convert DeepFlow graph to AstraSim ET format by scheduling per stage and DP rank. - This function is black magic. Don't touch it unless you know what you are doing. - Prefer changes *anywhere* else in the codebase if possible.""" + This means converting the single DeepFlow DAG into a set of different AstraSim ET files. + + INPUTS: ``graph_root`` is the DeepFlow DAG (after flattening). ``dp_size`` is the + data-parallel replication degree. ``output_dir`` is where ET traces + metadata + should be written. + + OUTPUTS: returns ``(et_prefix, rank_ids, manifest_path)``. ``et_prefix`` is the + file prefix for ``llm_graph..et``. ``rank_ids`` is the list of produced rank + IDs, and ``manifest_path`` points to the per-rank summary JSON. + + CONCEPTS: a *stage* is a unique ``hw_id`` (flattened hardware coordinate) of the + compute graph. Each (stage, DP replica) pair becomes an AstraSim rank. The function + recovers tensor/context/pipeline axes, schedules tasks per stage and per DP rank, and + emits Chakra ET traces with matching communicator metadata so AstraSim can replay the + workload deterministically. + + This function is very complicated. Don't touch it unless you know what you are doing. + Prefer changes *anywhere* else in the codebase if possible. I have tried to document it + as much as possible, but it definitely needs further cleaning and breaking up. This is TODO.""" os.makedirs(output_dir, exist_ok=True) debug_enabled = _env_truthy("DEEPFLOW_DUMP_CONVERTER_DEBUG") + # Step 1: Traverse the graph once to snapshot every object reachable from the + # provided root(s). Collecting upfront avoids mutating the graph while we + # analyse dependencies. def collect_objects(root) -> List[Any]: visited: Set[int] = set() ordered: List[Any] = [] @@ -459,6 +617,9 @@ def dfs(obj: Any) -> None: if not compute_nodes: raise ValueError("DeepFlow graph did not expose any executable compute nodes (hw_id >= 0).") + # Step 2: Each unique ``hw_id`` defines a logical stage (one flattened hardware + # coordinate). We build a DP-major rank ordering so AstraSim sees one rank per + # (stage, data-parallel replica) pair. stage_ids = sorted({node.hw_id for node in compute_nodes}) dp_count = max(int(dp_size) if dp_size else 1, 1) @@ -481,6 +642,16 @@ def dfs(obj: Any) -> None: def rank_for(stage: int, dp_idx: int) -> int: return stage_to_ranks[stage][dp_idx] + # Step 3: Recover the multi-dimensional axis layout (tp/cp/lp, etc.) so that + # communicator membership can be reconstructed deterministically later on. + axis_order, axis_sizes = _extract_axis_layout( + getattr(graph_root, "_astrasim_rank_layout", None) + ) + stage_axis_coords = _compute_stage_axis_coords(stage_ids, axis_order, axis_sizes) + axis_groups = _build_axis_groups(axis_order, axis_sizes, stage_axis_coords, stage_to_ranks) + + # Step 4: Attribute each comm edge to the stage that owns it. This lets the + # dependency walkers treat collectives the same way they treat compute nodes. edge_stage: Dict[Any, int] = {} for obj in all_objects: comm = getattr(obj, "comm_type", None) @@ -523,6 +694,8 @@ def rank_for(stage: int, dp_idx: int) -> int: pipeline_edge_map: Dict[Tuple[Any, Any], Any] = {} + # Step 5a: Walk backwards from a compute node to classify dependencies into + # stage-local, cross-stage (pipeline), or collective inputs. def analyze_for_compute(node: Any) -> Tuple[Set[Any], Set[Any], Set[Any]]: stage = node.hw_id stage_deps: Set[Any] = set() @@ -685,6 +858,8 @@ def analyze_for_compute(node: Any) -> Tuple[Set[Any], Set[Any], Set[Any]]: return stage_deps, pipeline_deps, collective_deps + # Step 5b: Same classification for collective edges (they often use local_hw_id + # instead of hw_id, so we check both). def analyze_for_collective(edge: Any) -> Tuple[Set[Any], Set[Any], Set[Any]]: stage = edge_stage[edge] stage_deps: Set[Any] = set() @@ -744,6 +919,7 @@ def analyze_for_collective(edge: Any) -> Tuple[Set[Any], Set[Any], Set[Any]]: return stage_deps, pipeline_deps, collective_deps + # Aggregate the analysis results so later stages can fetch dependencies by node. compute_info: Dict[Any, Dict[str, Any]] = {} for node in compute_nodes: stage_deps, pipeline_deps, collective_deps = analyze_for_compute(node) @@ -792,36 +968,32 @@ def analyze_for_collective(edge: Any) -> Tuple[Set[Any], Set[Any], Set[Any]]: f"[ConverterDebug] collective {name}: stage={info['stage']}" f" stage_deps={stage_names} pipeline_deps={pipe_names} collectives={coll_names}" ) + # Step 6: Group collectives by label/axis so we can reconstruct communicator + # memberships (tp/cp/lp) without relying on the original ordering. tp_collective_groups: Dict[str, List[Any]] = defaultdict(list) for edge, info in collective_info.items(): interconnect_type = info.get("interconnect_type") if interconnect_type and interconnect_type not in {"dp", "lp", "pipeline"}: tp_collective_groups[info["name"]].append(edge) - tp_collective_labels: Dict[Any, str] = {} - # Map final labels (with suffix) to their member edges, preserving order - label_to_edges: Dict[str, List[Any]] = {} - for base_name, edges in tp_collective_groups.items(): - if not edges: - continue - sorted_edges = sorted(edges, key=lambda edge: getattr(edge, "op_id", 0)) - group_index = 0 - idx = 0 - while idx < len(sorted_edges): - edge = sorted_edges[idx] - info = collective_info[edge] - participants = max(1, int(info.get("participants") or 1)) - label = base_name if group_index == 0 else f"{base_name}_{group_index}" - group_members = sorted_edges[idx: idx + participants] - for member in group_members: - tp_collective_labels[member] = label - # Record members for this final label in insertion order - if label not in label_to_edges: - label_to_edges[label] = [] - label_to_edges[label].extend(group_members) - idx += len(group_members) - group_index += 1 + ( + tp_collective_labels, + label_to_edges, + label_group_tokens, + ) = _assign_collective_labels( + tp_collective_groups, + collective_info, + axis_order, + axis_sizes, + stage_axis_coords, + axis_groups, + stage_to_ranks, + dp_count, + ) + # Step 7: Build a per-stage task list (compute + collectives). If a collective + # introduces a new stage we extend ``stage_to_ranks`` so every stage has the + # full (dp-major) rank list. stage_tasks: Dict[int, Set[Any]] = {stage: set() for stage in stage_ids} for node in compute_nodes: stage_tasks[node.hw_id].add(node) @@ -958,32 +1130,30 @@ def analyze_for_collective(edge: Any) -> Tuple[Set[Any], Set[Any], Set[Any]]: tag_counter = itertools.count(start=1) # Build TP communicator groups (ids start at 1000) per label and dp - global _LAST_TP_GROUPS, _TP_LABEL_DP_TO_ID + global _LAST_TP_GROUPS, _TP_LABEL_DP_TO_ID, _TP_MEMBERS_TO_GID _LAST_TP_GROUPS = {} _TP_LABEL_DP_TO_ID = {} + _TP_MEMBERS_TO_GID = {} tp_gid_counter = _it_count(start=_TP_GROUP_BASE_ID) # Collect per-label TP members per dp index (use prebuilt mapping to avoid rescans) - for label in sorted(label_to_edges.keys()): - edges_for_label = label_to_edges.get(label, []) - if not edges_for_label: - continue - # Determine stage for each edge - stages = [collective_info[e]["stage"] for e in edges_for_label] - # For each dp index, map to ranks - for dp_idx in range(dp_count): - members: List[int] = [] - for e, st in zip(edges_for_label, stages): - try: - members.append(stage_to_ranks[st][dp_idx]) - except Exception: - continue - if not members: - continue - gid = str(next(tp_gid_counter)) + # Step 9: Convert axis membership tokens into communicator IDs (reused across + # labels/dp replicas) so later we can write ``comm_groups.json``. + for label in sorted(label_group_tokens.keys()): + tokens = sorted(label_group_tokens[label]) + for axis, dp_idx, members_tuple in tokens: + members_list = list(members_tuple) + members_key = tuple(members_list) + gid = _TP_MEMBERS_TO_GID.get(members_key) + if gid is None: + gid = str(next(tp_gid_counter)) + _TP_MEMBERS_TO_GID[members_key] = gid + if gid not in _LAST_TP_GROUPS: + _LAST_TP_GROUPS[gid] = members_list _TP_LABEL_DP_TO_ID[(label, dp_idx)] = gid - _LAST_TP_GROUPS[gid] = sorted(members) + # Step 8 helper: materialise SEND/RECV control edges across stages whenever a + # pipeline dependency exists between ``parent`` and ``target``. def ensure_pipeline(parent: Any, target: Any, dp_idx: int) -> int: parent_stage = getattr(parent, "hw_id", None) if parent_stage == None: @@ -1066,6 +1236,8 @@ def ensure_pipeline(parent: Any, target: Any, dp_idx: int) -> int: return recv_id + # Step 10: Emit ET nodes per stage (collectives first, then compute nodes), + # iterating in the topological order computed above. for stage in stage_order: order = stage_order[stage] for task in order: @@ -1153,7 +1325,8 @@ def ensure_pipeline(parent: Any, target: Any, dp_idx: int) -> int: # above to avoid duplication. Here we only reference them # via dependencies when needed. - # Attach cross-stage (pipeline) dependencies after all nodes are created. + # Step 11: Attach cross-stage (pipeline) dependencies by wiring the RECV + # nodes created via ``ensure_pipeline`` into the ET nodes’ control deps. for stage, order in stage_order.items(): for task in order: if task in collective_info: @@ -1196,7 +1369,8 @@ def ensure_pipeline(parent: Any, target: Any, dp_idx: int) -> int: node.ctrl_deps.append(rid) - # Build manifest ops directly from in-memory traces (graph-derived), not ET files + # Step 12: Build manifest ops directly from the in-memory traces. This lets the + # caller reuse the same metadata that will be written into the ET files. manifest_ranks: Dict[str, List[List]] = {} def _manifest_op_key(op: List) -> tuple: # Sort primarily by size/duration, with stable type ordering @@ -1257,7 +1431,8 @@ def _manifest_op_key(op: List) -> tuple: rank_ops_sorted = sorted(rank_ops, key=_manifest_op_key) manifest_ranks[str(int(rank))] = rank_ops_sorted - # Now actually write ET files + # Step 13: Write the Chakra ET traces (with control-send renumbering) and + # persist the manifest/metadata alongside them. for trace in rank_traces.values(): trace.close() @@ -1338,8 +1513,6 @@ def run_astra_simulation_only_onepath( if persist and _env_truthy("DEEPFLOW_PERSIST_ARTIFACT_VIZ"): et_paths = [] for rank in rank_ids: - if rank > 10: - break et_path = os.path.join(work_dir, f"llm_graph.{rank}.et") if os.path.exists(et_path): et_paths.append(et_path) @@ -1408,8 +1581,8 @@ def run_astra_simulation_only_onepath( def my_save_graph(roots, output_folder = "output_graph/", filename="graph"): dot_fw = visualize_graph(roots, filename=output_folder + filename) - dot_fw.render(output_folder + filename , format="png", cleanup=True) - print("graph saved to %s%s.png" % (output_folder , filename )) + dot_fw.render(output_folder + filename , format="svg", cleanup=True) + print("graph saved to %s%s.svg" % (output_folder , filename )) import config import pickle diff --git a/astrasim_lib/integration.py b/astrasim_lib/integration.py index d70a8317..15ba19d8 100644 --- a/astrasim_lib/integration.py +++ b/astrasim_lib/integration.py @@ -357,9 +357,16 @@ def run_cache_astrasim( sig["sys_options"] = sys_opts_sig if manifest_json_path: sig["multinode"] = True - canonical = _canonical_sig(sig) - files = generate_astrasim_configs_from_hw(hw_obj, out_dir=astra_config_dir, npus_count=npus_count) + topo_list = files.get("topology_list") + npus_per_dim = files.get("npus_per_dim") + if topo_list: + sig["topology"] = list(topo_list) + sig["physical_topology"] = list(topo_list) + if npus_per_dim: + sig["npus_per_dim"] = list(npus_per_dim) + + canonical = _canonical_sig(sig) if workload_prefix: pass diff --git a/configs/hardware-config/a100_80GB_problem.yaml b/configs/hardware-config/a100_80GB_problem.yaml new file mode 100644 index 00000000..333a61af --- /dev/null +++ b/configs/hardware-config/a100_80GB_problem.yaml @@ -0,0 +1,141 @@ +sw_param: + kernel_launch_overhead: 9e-6 # (s) + # dp_zero_stage selects ZeRO-style data-parallel sharding: 0=DDP, 1=ZeRO-1, 2=ZeRO-2, 3=ZeRO-3 + # only used if data parallelism (dp) > 1 + dp_zero_stage: 2 + precision: + # tensor_format accepts numeric byte counts or dtype strings (fp16, bf16, fp32). + # Take care that the number of flops defined below has been carefully tuned for FP16/BF16 formats. + # For FP32 or other formats they need to be manually adjusted. + # The following are equivalent: (tensor_format: 2.0, tensor_format: fp16, tensor_format: bf16). + tensor_format: bf16 # baseline storage width for activations/gradients. + # mixed_precision: true | false + # -> false forces all activations/weights/gradients/optimizer states/stats to use tensor_format bytes. + # -> true enables AMP-style FP32 accumulations/optimizer states depending on param_storage_mode. + # if true, optimizer states and stats are always FP32. + mixed_precision: false + # param_storage_mode: as_tensor_format | tensor_plus_fp32_master | fp32_params. + # -> as_tensor_format makes weights/gradients use tensor_format bytes. + # -> tensor_plus_fp32_master makes weights/gradients use tensor_format bytes and + # keeps a separate FP32 master copy for weights. + # -> fp32_params makes weights *and* gradients use FP32 bytes without a master copy. + param_storage_mode: as_tensor_format + # kv_cache precision may be set as bytes, dtype string, or "as_tensor_format". + kv_cache: as_tensor_format + h2d_bandwidth: -1 # Bytes/s; set -1 to disable H2D timing +tech_param: + core: + operating_frequency: 1.41e9 # (Hz) + num_bundles: 108 + nominal_flop_rate_per_mcu: 512 + nominal_energy_per_flop: 1.8e-13 # (J) + num_mcu_per_bundle: 4 # number of tensorcores per SM + FMA_d1: 8 # tensor core can compute the results for an d1 x d2 x d1 mixed-precision matrix multiplication per clock + FMA_d2: 4 + dataflow: best # {wst | ast | ost | best}: wst: weight stationary, ast: activation stationary, ost: output stationary, best: best of all 3 + util: 100.0 + DRAM: + size: 80 GB + bandwidth: 198600 GB # (/s) + dynamic_energy_per_bit: 4.5e-12 # (J) + latency: 100e-9 # (s) + util: 1 + SRAM-L2: + size: 40 MB + bandwidth: 7050 GB + dynamic_energy_per_bit: 130e-15 # (J) + latency: 0 + util: 1 + SRAM-L1: + size: 20736 KB + bandwidth: 18 TB + dynamic_energy_per_bit: 130e-15 # (J) + latency: 0 + util: 1 + SRAM-R: + size: 27 MB + bandwidth: 122 TB + dynamic_energy_per_bit: 110e-15 # (J) + latency: 0 + util: 1 + +memory_hierarchy: + l0: #Register Memory + type: SRAM-R + scope: mcu + l1: #Shared Memory + type: SRAM-L1 + scope: mcu-bundle + l2: #L2 + type: SRAM-L2 + scope: global + l3: #Global Memory + type: DRAM + scope: global + +parallelism: + dp: 2 # data parallelism. Acts as replica count during inference + tp: 3 # tensor parallelism + tp_sp: true # tensor-sequence parallelism + cp: 2 # context parallelism. Not supported for inference. + lp: 3 # pipeline parallelism. Number of pipeline stages + # mb is not a true parallelism parameter, does not affect num of GPUs. + mb: 2 # micro-batches (for best results set to >= lp). + +network: + # currently only support astrasim integration with all devices in one node. + - id: dim0 + label: intra_wafer_mesh_silicon + size: 6 + topology: + type: Ring + bandwidth: 1000 Gb + latency: 5e-6 + energy_per_bit: 8e-12 # (J/bit) + util: 1.0 + collective_override: {} # overrides execution_backend.collectives when provided + parallelisms: [tp, cp] + - id: dim1 + label: inter_wafer_ring_infiniband + size: 3 + topology: + type: Ring + bandwidth: 200 Gb + latency: 5e-6 + energy_per_bit: 8e-12 # (J/bit) + util: 1.0 + collective_override: {} # overrides execution_backend.collectives when provided + parallelisms: [lp] + - id: dim2 + label: wafer_dp_chains_ring_infiniband + size: 2 + topology: + type: ring + bandwidth: 200 Gb + latency: 5e-6 + energy_per_bit: 8e-12 # (J/bit) + util: 0.96 + collective_override: {} # overrides execution_backend.collectives when provided + parallelisms: [dp] + + +execution_backend: + model: astra # analytical | astra + astra: + backend: astra + mode: full_astrasim_flattened # hybrid | full_astrasim_hierarchical | full_astrasim_flattened + collectives: + # Default collective algorithms. Per-dimension collective_override entries can override these. + # Allowed values per op: auto | ring | direct | halvingDoubling | doubleBinaryTree | mesh | mesh2d | hypercube | torus2d + # Note: leave to auto unless you have a specific reason to change it. + all_gather: auto + all_reduce: auto + reduce_scatter: auto + all_to_all: auto + sys_options: + endpoint_delay: 5 + active_chunks_per_dimension: 32 # Higher value means significantly slower astrasim runtime, but more accurate especially for tensor parallelism + preferred_dataset_splits: 4 # Should be same as active_chunks_per_dimension + +inference: + kvcache_type: hbm_only # Only supported mode for now. diff --git a/configs/model-config/LLM_problem.yaml b/configs/model-config/LLM_problem.yaml new file mode 100644 index 00000000..9ab1e9c6 --- /dev/null +++ b/configs/model-config/LLM_problem.yaml @@ -0,0 +1,19 @@ +model_param: + mode: "LLM" # Mode of operation + run_type: "training" # Type of run: "training" or "inference" + tied_embeddings: false # Share input and output embedding weights (WIP: affects memory estimation only) + model_type: "gpt" # gpt assumes dense-style architecture with GELU; llama assumes llama-style architecture with SwiGLU + batch_size: 1024 # Sequences per global batch + seq_len: 16384 # Sequence length (tokens per sequence) + hidden_dim: 8192 # Transformer hidden dimension + attention: # Attention-related parameters + attention_type: "mha" # Supported attention types: "mha" or "gqa" ("mla" not implemented) + num_heads: 64 # Total number of attention heads + kv_heads: 8 # Number of key/value heads (used only when attention_type = "gqa") + use_flashattention: false # Enable FlashAttention optimization + attention_tile_size: 256 # Tile size used in attention computation + num_experts: 1 # Number of experts in MoE layers. Set to 1 to disable MoE. + top_k: 1 # Number of active experts per token in MoE + intermediate_size: 32768 # Output dimension of the first FFN layer (typically 4×hidden_dim) + vocab_size: 800000 # Vocabulary size + num_layers: 3 # Number of transformer layers diff --git a/examples/problem.sh b/examples/problem.sh new file mode 100755 index 00000000..01bb28fa --- /dev/null +++ b/examples/problem.sh @@ -0,0 +1,7 @@ +#!/bin/bash + +# Runs performance model in LLM mode +# Setup model parameter in configs/model-config/Llama2-7B.yaml +# Setup hardware parameters (incl. parallelism/network) in configs/hardware-config/a100_80GB.yaml + +DEEPFLOW_PERSIST_ASTRASIM_ARTIFACTS=1 DEEPFLOW_VISUALIZE_GRAPHS=1 DEEPFLOW_PERSIST_ARTIFACT_VIZ=1 python run_perf.py --hardware_config configs/hardware-config/a100_80GB_problem.yaml --model_config configs/model-config/LLM_problem.yaml \ No newline at end of file diff --git a/simulate_LLM.py b/simulate_LLM.py index 8ee46c0a..46a5c52c 100644 --- a/simulate_LLM.py +++ b/simulate_LLM.py @@ -98,7 +98,10 @@ def __init__( self.num_batch = self.misc_metadata.get("num_batch", 0) self.num_layer = self.misc_metadata.get("num_layer", 0) - self.layer_per_device = max(1, math.ceil(self.num_layer / self.lp)) if self.lp else self.num_layer + self.layer_per_device = self.num_layer / self.lp if self.lp else self.num_layer + # if self layers per device is not an integer, raise Exception complaining about divisors! + if self.layer_per_device != int(self.layer_per_device): + raise Exception(f"The number of layers {self.num_layer} is not divisible by the degree of pipeline parallelism {self.lp}") self.all_reduce = self.misc_metadata.get("all_reduce", "every layer") self.transformer_cfg = self.comp_times.get("transformer", {}) @@ -1484,18 +1487,18 @@ def _count_transformer_layers_per_device(root_obj: Any, num_devices: int) -> Lis def save_graph(self, roots, output_folder = "output/LLM/", filename="graph"): os.makedirs(output_folder, exist_ok=True) - printstr = " | Graph saved to %s%s.png" % (output_folder, filename) + printstr = " | Graph saved to %s%s.svg" % (output_folder, filename) def _render_graph() -> None: dot_fw = visualize_graph(roots, filename=output_folder + filename) - dot_fw.render(output_folder + filename, format="png", cleanup=True) + dot_fw.render(output_folder + filename, format="svg", cleanup=True) - graphviz_async.submit(f"{filename}.png", _render_graph, print_message=printstr) + graphviz_async.submit(f"{filename}.svg", _render_graph, print_message=printstr) def visualize_graph(roots, filename="graph"): _ = filename # unused, kept for backwards compatibility with callers - dot = Digraph(comment="Computation Graph") + dot = Digraph(comment="Computation Graph", format="svg") visited = set() def _format_duration(value: float) -> str: From fd1a137fb65116f7bc0d704cd698479774e39f0e Mon Sep 17 00:00:00 2001 From: gkarfakis19 Date: Mon, 3 Nov 2025 03:15:45 -0800 Subject: [PATCH 11/77] enabled multi dim network for non-flattened exec modes --- LLM_execution.py | 67 +++++++++++++++++++++++++++++-- TODO_11_3_2025.md | 3 ++ astrasim_lib/config_generation.py | 29 ++++++++++--- astrasim_lib/executor.py | 37 +++++++++++++++-- astrasim_lib/integration.py | 20 +++++++-- 5 files changed, 139 insertions(+), 17 deletions(-) create mode 100644 TODO_11_3_2025.md diff --git a/LLM_execution.py b/LLM_execution.py index ed4d15c5..db64c3de 100644 --- a/LLM_execution.py +++ b/LLM_execution.py @@ -4,7 +4,7 @@ import os from dataclasses import dataclass from enum import Enum -from typing import Any, Dict, Tuple, Optional, List, Set +from typing import Any, Dict, Tuple, Optional, List, Set, Sequence import simulate_LLM from astrasim_lib import run_astra_simulation_only_onepath @@ -636,6 +636,8 @@ def __init__( self.transformer_forward_root = transformer_forward_root self.transformer_backward_root = transformer_backward_root self.flattened_root: Optional[Any] = None + self._transformer_rank_layout: Optional[Dict[str, Any]] = None + self._pipeline_rank_layout: Optional[Dict[str, Any]] = None self._rank_layout = self._build_rank_layout_descriptor() def _build_rank_layout_descriptor(self) -> Dict[str, Any]: @@ -659,8 +661,31 @@ def _safe_int(value: Any, default: int = 1) -> int: axis_sizes: Dict[str, int] = {"tp": tp_size, "cp": cp_size, "lp": lp_size, "dp": dp_size} axis_order: List[str] = [] + + # Enforce axis ordering for hierarchical/hybrid modes: the first active + # dimension must contain exactly {'tp','cp'}. Subsequent active + # dimensions may contain 'lp' (optionally combined with 'dp'). This + # matches the assumptions in the hierarchical graphs where a stage is a + # TP/CP cluster replicated across LP (and potentially DP) axes. + enforce_layout = self.time_calc.execution_mode in { + ExecutionMode.HYBRID, + ExecutionMode.FULL_ASTRASIM_HIERARCHICAL, + } + first_active_checked = False + for dim in dimensions: dim_axes = [str(axis).strip().lower() for axis in getattr(dim, "parallelisms", ())] + declared = int(getattr(dim, "size", 1)) + + if enforce_layout and declared > 1 and not first_active_checked: + canon = sorted(dim_axes) + if canon != ["cp", "tp"]: + raise ValueError( + "For hierarchical/hybrid AstraSim modes, the first active network " + "dimension must contain exactly {'tp','cp'} to represent the tensor/context cluster." + ) + first_active_checked = True + for name in dim_axes: if name == "dp": continue @@ -672,7 +697,6 @@ def _safe_int(value: Any, default: int = 1) -> int: if name not in axis_order: axis_order.append(name) - declared = int(getattr(dim, "size", 1)) expected = 1 for axis_name in dim_axes: expected *= axis_sizes.get(axis_name, 1) @@ -709,6 +733,29 @@ def _safe_int(value: Any, default: int = 1) -> int: f"{axis_order or ['']} ({pretty_sizes or 'single device'}{dp_msg}) | " "dp ranks are appended separately." ) + + def _subset_layout(allowed: Sequence[str]) -> Optional[Dict[str, Any]]: + subset = [axis for axis in axis_order if axis in allowed and axis_sizes.get(axis, 1) > 1] + if not subset: + return None + strides: Dict[str, int] = {} + span = 1 + for axis in subset: + strides[axis] = span + span *= axis_sizes[axis] + return { + "axis_order": subset, + "axis_sizes": {axis: axis_sizes[axis] for axis in subset}, + "axis_strides": strides, + "stage_span": span, + } + + # Transformer graphs only encode TP/CP axes; pipeline graphs encode LP + # (DP replicas are handled externally). Store these subsets so callers + # can attach the appropriate layout before invoking AstraSim. + self._transformer_rank_layout = _subset_layout(["tp", "cp"]) + self._pipeline_rank_layout = _subset_layout(["lp"]) + return descriptor def run(self, mode: ExecutionMode) -> ExecutionResult: @@ -795,7 +842,11 @@ def _run_full_astrasim_hierarchical(self) -> ExecutionResult: "/pipeline_graph_hierarchical", ) - setattr(self.pipeline_root, "_astrasim_rank_layout", self._rank_layout) + pipeline_layout = getattr(self, "_pipeline_rank_layout", None) + if pipeline_layout: + setattr(self.pipeline_root, "_astrasim_rank_layout", pipeline_layout) + elif hasattr(self.pipeline_root, "_astrasim_rank_layout"): + delattr(self.pipeline_root, "_astrasim_rank_layout") per_rank_sec, max_sec = run_astra_simulation_only_onepath( self.pipeline_root, self.time_calc, @@ -959,6 +1010,11 @@ def _run_transformer_astrasim(self, mode: ExecutionMode) -> Optional[Transformer fwd_max = 0 bwd_max = 0 if self.transformer_forward_root: + layout = getattr(self, "_transformer_rank_layout", None) + if layout: + setattr(self.transformer_forward_root, "_astrasim_rank_layout", layout) + elif hasattr(self.transformer_forward_root, "_astrasim_rank_layout"): + delattr(self.transformer_forward_root, "_astrasim_rank_layout") fwd_per_rank, fwd_max = run_astra_simulation_only_onepath( self.transformer_forward_root, self.time_calc, @@ -967,6 +1023,11 @@ def _run_transformer_astrasim(self, mode: ExecutionMode) -> Optional[Transformer persist_artifacts=self.time_calc.persist_astrasim_artifacts, ) if self.transformer_backward_root: + layout = getattr(self, "_transformer_rank_layout", None) + if layout: + setattr(self.transformer_backward_root, "_astrasim_rank_layout", layout) + elif hasattr(self.transformer_backward_root, "_astrasim_rank_layout"): + delattr(self.transformer_backward_root, "_astrasim_rank_layout") bwd_per_rank, bwd_max = run_astra_simulation_only_onepath( self.transformer_backward_root, self.time_calc, diff --git a/TODO_11_3_2025.md b/TODO_11_3_2025.md new file mode 100644 index 00000000..a066c6c9 --- /dev/null +++ b/TODO_11_3_2025.md @@ -0,0 +1,3 @@ +a) size=auto for network +b) enable decode len = 0 +c) debug hierarchical dp in last dim (3D) vs dp with lp (2D) \ No newline at end of file diff --git a/astrasim_lib/config_generation.py b/astrasim_lib/config_generation.py index 699af7cc..3b581abf 100644 --- a/astrasim_lib/config_generation.py +++ b/astrasim_lib/config_generation.py @@ -3,7 +3,7 @@ from __future__ import annotations import os -from typing import Any, Dict, List, Optional, Tuple +from typing import Any, Dict, List, Optional, Sequence, Tuple from hw_component import Network @@ -109,6 +109,8 @@ def generate_astrasim_configs_from_hw( hw_obj, out_dir: str = "./astra_cache", npus_count: Optional[int] = None, + *, + axes_filter: Optional[Sequence[str]] = None, ) -> Dict[str, str]: """Write AstraSim network/system configs derived from ``hw_obj``.""" if npus_count is None: @@ -133,11 +135,26 @@ def _safe_int(value: Any, default: int = 1) -> int: "lp": _safe_int(getattr(sch_config, "lp", 1)) if sch_config else 1, "dp": _safe_int(getattr(sch_config, "dp", 1)) if sch_config else 1, } + if axes_filter: + axis_sizes = {axis: axis_sizes.get(axis, 1) for axis in axes_filter} + axis_sizes.setdefault("tp", 1) + axis_sizes.setdefault("cp", 1) + axis_sizes.setdefault("lp", 1) + axis_sizes.setdefault("dp", 1) axis_order_preference = ["tp", "cp", "lp", "dp"] + allowed_axes = set(axis_sizes.keys()) if axes_filter else None + # print(f"Allowed axes: {allowed_axes}") + # print(f"Axes sizes: {axis_sizes}") dim_infos: List[Tuple[Any, List[str], int]] = [] + # print(f"Filter: {axes_filter}") for dim in dimensions: axes = [str(axis).strip().lower() for axis in getattr(dim, "parallelisms", ())] + # print(f"All axes for dimension {dim.label}: {axes}") + if allowed_axes is not None and axes_filter: + filtered_axes = [axis for axis in axes if axis in allowed_axes] + axes = filtered_axes + # print(f"Filtered axes for dimension {dim.label}: {axes}") effective = 1 for axis in axes: if axis not in axis_sizes: @@ -146,11 +163,11 @@ def _safe_int(value: Any, default: int = 1) -> int: "Supported axes are tp, cp, lp, dp." ) effective *= axis_sizes[axis] - if effective != int(getattr(dim, "size", effective)): - raise ValueError( - f"Network dimension '{dim.label}' size mismatch: declared {dim.size}, " - f"but parallelism factors imply {effective}." - ) + # if effective != int(getattr(dim, "size", effective)): + # raise ValueError( + # f"Network dimension '{dim.label}' size mismatch: declared {dim.size}, " + # f"but parallelism factors imply {effective}." + # ) dim_infos.append((dim, axes, effective)) active_dim_indices = [idx for idx, (dim, _, _) in enumerate(dim_infos) if int(getattr(dim, "size", 0)) > 1] diff --git a/astrasim_lib/executor.py b/astrasim_lib/executor.py index 1932bbfa..fd2789f3 100644 --- a/astrasim_lib/executor.py +++ b/astrasim_lib/executor.py @@ -165,6 +165,25 @@ def _build_axis_groups( return axis_groups +def _derive_axes_filter( + axis_order: Sequence[str], + axis_sizes: Mapping[str, int], + dp_count: int, +) -> Optional[List[str]]: + axes: List[str] = [] + for axis in axis_order: + try: + size = max(1, int(axis_sizes.get(axis, 1))) + except Exception: + size = 1 + if size > 1 and axis not in axes: + axes.append(axis) + if dp_count > 1: + axes = [axis for axis in axes if axis != "dp"] + axes.append("dp") + return axes or None + + def _assign_collective_labels( tp_collective_groups: Mapping[str, List[Any]], collective_info: Mapping[Any, Dict[str, Any]], @@ -644,9 +663,8 @@ def rank_for(stage: int, dp_idx: int) -> int: # Step 3: Recover the multi-dimensional axis layout (tp/cp/lp, etc.) so that # communicator membership can be reconstructed deterministically later on. - axis_order, axis_sizes = _extract_axis_layout( - getattr(graph_root, "_astrasim_rank_layout", None) - ) + rank_layout = getattr(graph_root, "_astrasim_rank_layout", None) + axis_order, axis_sizes = _extract_axis_layout(rank_layout) stage_axis_coords = _compute_stage_axis_coords(stage_ids, axis_order, axis_sizes) axis_groups = _build_axis_groups(axis_order, axis_sizes, stage_axis_coords, stage_to_ranks) @@ -1529,7 +1547,16 @@ def run_astra_simulation_only_onepath( if os.environ.get("DEEPFLOW_ASTRA_SKIP_EXEC"): print("[AstraSim] DEEPFLOW_ASTRA_SKIP_EXEC set. Exiting after ET artifact generation.") exit() - astra_configs = generate_astrasim_configs_from_hw(time_calc_obj.hw_config, work_dir, rank_count) + + rank_layout = getattr(fwdbwd_root, "_astrasim_rank_layout", None) + axis_order, axis_sizes = _extract_axis_layout(rank_layout) + axes_filter = _derive_axes_filter(axis_order, axis_sizes, dp_count) + astra_configs = generate_astrasim_configs_from_hw( + time_calc_obj.hw_config, + work_dir, + rank_count, + axes_filter=axes_filter, + ) # Run AstraSim simulation on forward graph (cached via manifest) print(f"[AstraSim] Executing forward simulation with {rank_count} ranks...") @@ -1543,6 +1570,8 @@ def run_astra_simulation_only_onepath( manifest_json_path=fwd_manifest, workload_prefix=fwd_et_prefix, comm_group_json=comm_groups_path, + axes_filter=axes_filter, + files=astra_configs, ) diff --git a/astrasim_lib/integration.py b/astrasim_lib/integration.py index 15ba19d8..851dcba7 100644 --- a/astrasim_lib/integration.py +++ b/astrasim_lib/integration.py @@ -8,7 +8,7 @@ import re import subprocess import time -from typing import Any, Dict, Iterable, List, Optional, Tuple +from typing import Any, Dict, Iterable, List, Optional, Sequence, Tuple from .bootstrap import ensure_chakra_available from .config_generation import ( @@ -315,6 +315,8 @@ def run_cache_astrasim( manifest_json_path: Optional[str] = None, workload_prefix: Optional[str] = None, comm_group_json: Optional[str] = None, + axes_filter: Optional[Sequence[str]] = None, + files = None, ) -> Tuple[List[float], float]: """Run AstraSim with caching to avoid recomputation when inputs match.""" @@ -341,6 +343,7 @@ def run_cache_astrasim( topo = derive_topology_from_hw(hw_obj) colls = _collectives_from_hw(hw_obj, topo) sys_opts_sig = _sys_options_from_hw(hw_obj) + axes_filter_tuple: Optional[Tuple[str, ...]] = tuple(axes_filter) if axes_filter else None sig: Dict[str, Any] = { "comm": comm.lower(), @@ -357,7 +360,15 @@ def run_cache_astrasim( sig["sys_options"] = sys_opts_sig if manifest_json_path: sig["multinode"] = True - files = generate_astrasim_configs_from_hw(hw_obj, out_dir=astra_config_dir, npus_count=npus_count) + if axes_filter_tuple is not None: + sig["axes_filter"] = axes_filter_tuple + if not files: + files = generate_astrasim_configs_from_hw( + hw_obj, + out_dir=astra_config_dir, + npus_count=npus_count, + axes_filter=axes_filter_tuple, + ) topo_list = files.get("topology_list") npus_per_dim = files.get("npus_per_dim") if topo_list: @@ -434,6 +445,8 @@ def run_cache_astrasim( debug_cmd.append(f"--remote-memory-configuration={remote_mem_path}") if comm_group_json and os.path.exists(comm_group_json): debug_cmd.append(f"--comm-group-configuration={comm_group_json}") + cmd_str = " ".join(debug_cmd) + # print(f"Full AstraSim command:\n{cmd_str}") while attempts < 5: attempts += 1 @@ -448,9 +461,8 @@ def run_cache_astrasim( last_outcome_zero = False break last_outcome_zero = True - time.sleep(0.5) + time.sleep(0.1) if last_outcome_zero: - cmd_str = " ".join(debug_cmd) raise RuntimeError( f"AstraSim returned zero time for {comm} size={size_bytes} npus={npus_count} after 5 retries.\n" f"Full AstraSim command:\n{cmd_str}" From 88eba5bdd43494c33c48cbbdae03657c767cb5ae Mon Sep 17 00:00:00 2001 From: gkarfakis19 Date: Mon, 3 Nov 2025 16:10:30 -0800 Subject: [PATCH 12/77] added network prints, fixed hierarchical context parallelism --- .gitignore | 3 +- astrasim_lib/config_generation.py | 5 - config.py | 63 ++++++++++--- run_perf.py | 17 +++- simulate_LLM.py | 149 ++++++++++++++++-------------- util.py | 94 +++++++++++++++++++ 6 files changed, 242 insertions(+), 89 deletions(-) diff --git a/.gitignore b/.gitignore index 0d6ea66b..82e51f56 100644 --- a/.gitignore +++ b/.gitignore @@ -5,4 +5,5 @@ output/ venv/ astra_cache/ log/*.log -tmp/ \ No newline at end of file +tmp/ +output*/ \ No newline at end of file diff --git a/astrasim_lib/config_generation.py b/astrasim_lib/config_generation.py index 3b581abf..4bb56439 100644 --- a/astrasim_lib/config_generation.py +++ b/astrasim_lib/config_generation.py @@ -163,11 +163,6 @@ def _safe_int(value: Any, default: int = 1) -> int: "Supported axes are tp, cp, lp, dp." ) effective *= axis_sizes[axis] - # if effective != int(getattr(dim, "size", effective)): - # raise ValueError( - # f"Network dimension '{dim.label}' size mismatch: declared {dim.size}, " - # f"but parallelism factors imply {effective}." - # ) dim_infos.append((dim, axes, effective)) active_dim_indices = [idx for idx, (dim, _, _) in enumerate(dim_infos) if int(getattr(dim, "size", 0)) > 1] diff --git a/config.py b/config.py index b2b691ff..f3fae6ea 100644 --- a/config.py +++ b/config.py @@ -372,12 +372,18 @@ def from_raw( if "size" not in raw: raise ValueError(f"network dimension '{label}' is missing required field 'size'") - try: - size = int(raw["size"]) - except (TypeError, ValueError) as exc: - raise ValueError(f"network dimension '{label}' size must be an integer") from exc - if size < 1: - raise ValueError(f"network dimension '{label}' size must be >= 1") + size_raw = raw["size"] + size_is_auto = False + if isinstance(size_raw, str) and size_raw.strip().lower() == "auto": + size_is_auto = True + size = None + else: + try: + size = int(size_raw) + except (TypeError, ValueError) as exc: + raise ValueError(f"network dimension '{label}' size must be an integer") from exc + if size < 1: + raise ValueError(f"network dimension '{label}' size must be >= 1") topo_dict = raw.get("topology") if not isinstance(topo_dict, dict): @@ -448,12 +454,25 @@ def from_raw( normalized_parallelisms.append(normalized) alias_map[normalized] = name + computed_product = None + if size_is_auto: + computed_product = _compute_dimension_parallelism_product( + dimension_label=label, + normalized_names=tuple(normalized_parallelisms), + alias_map=alias_map, + parallelism_params=parallelism_params, + ) + size = computed_product + if size < 1: + raise ValueError(f"network dimension '{label}' inferred size must be >= 1") + _validate_dimension_parallelisms( dimension_label=label, dimension_size=size, normalized_names=tuple(normalized_parallelisms), alias_map=alias_map, parallelism_params=parallelism_params, + expected_product=computed_product, ) return cls( @@ -481,10 +500,35 @@ def _validate_dimension_parallelisms( normalized_names: Tuple[str, ...], alias_map: Dict[str, str], parallelism_params: Dict[str, object], + expected_product: Optional[int] = None, ) -> None: if not normalized_names: return + product = expected_product + if product is None: + product = _compute_dimension_parallelism_product( + dimension_label=dimension_label, + normalized_names=normalized_names, + alias_map=alias_map, + parallelism_params=parallelism_params, + ) + + if product != dimension_size: + readable = [alias_map.get(name, name) for name in normalized_names] + raise ValueError( + f"Network dimension '{dimension_label}' size mismatch: declared size {dimension_size} " + f"but parallelism factors ({readable}) imply {product}" + ) + + +def _compute_dimension_parallelism_product( + *, + dimension_label: str, + normalized_names: Tuple[str, ...], + alias_map: Dict[str, str], + parallelism_params: Dict[str, object], +) -> int: product = 1 for name in normalized_names: if name not in parallelism_params: @@ -514,12 +558,7 @@ def _validate_dimension_parallelisms( ) product *= factor - if product != dimension_size: - readable = [alias_map.get(name, name) for name in normalized_names] - raise ValueError( - f"network dimension '{dimension_label}' size mismatch: declared size {dimension_size} " - f"but product of parallelisms {readable} is {product}" - ) + return product def _parse_network_layout( diff --git a/run_perf.py b/run_perf.py index 7086868b..265fac13 100755 --- a/run_perf.py +++ b/run_perf.py @@ -10,6 +10,7 @@ import pandas as pd import yaml import shutil +import util import graphviz_async from tile import TiledGEMM, formatBytes @@ -96,6 +97,7 @@ def _validate_network_topology(hw_config) -> None: "Only execution_backend.model='astra' (requires a valid AstraSim install) supports non-ring networks." ) + def run_LSTM( exp_hw_config_path, exp_model_config_path, @@ -237,6 +239,7 @@ def _run_llm_training(exp_hw_config, exp_model_config, exp_dir, mode): output_file = os.path.join(exp_dir, "LLM_training_results.txt") tc_llm = TimeCalculationLLM(exp_hw_config, exp_model_config, mode, output_dir=exp_dir) total_time = tc_llm.calc_time_llm() + topology_lines = util.network_topology_summary_training(exp_hw_config) with open(output_file, "a+") as handle: handle.write("\n\n==============================================\n") @@ -245,9 +248,13 @@ def _run_llm_training(exp_hw_config, exp_model_config, exp_dir, mode): handle.write("Execution Mode: {}\n".format(tc_llm.execution_mode.value)) handle.write("Total Time: {0:.8f}\n".format(total_time)) handle.write("\n") - handle.write("For more info, turn on debug flags. See examples/llm_astra_inference_debug_graphviz.sh") + handle.write("For more info, turn on debug flags. See examples/llm_astra_inference_debug_graphviz.sh\n") + handle.write("\n".join(topology_lines)) + handle.write("\n") print("Training time for batch: {:.2f}s".format(tc_llm.get_time())) + for line in topology_lines: + print(line) print("LLM training results written to {}".format(output_file)) warning_message = tc_llm.memory_capacity_warning() if warning_message: @@ -306,6 +313,8 @@ def _run_llm_inference(exp_hw_config, exp_model_config, exp_dir, mode): ) ) + topology_lines = util.network_topology_summary_inference(exp_hw_config) + output_path = os.path.join(exp_dir, "LLM_inference_results.txt") os.makedirs(exp_dir, exist_ok=True) with open(output_path, "w") as handle: @@ -328,8 +337,12 @@ def _run_llm_inference(exp_hw_config, exp_model_config, exp_dir, mode): handle.write(f"Decode Generations per Second: start={start_gen_rate:.2f}, mid(token {mid_step})={mid_gen_rate:.2f}, end={end_gen_rate:.2f}\n") handle.write(f"Aggregate Decode Throughput Tok/s (batch={batch_size}, dp={dp_replicas}): start={start_gen_rate * batch_size * dp_replicas:.2f}, mid(token {mid_step})={mid_gen_rate * batch_size * dp_replicas:.2f}, end={end_gen_rate * batch_size * dp_replicas:.2f}\n") handle.write("\n") - handle.write("For more info, turn on debug flags. See examples/llm_astra_inference_debug_graphviz.sh") + handle.write("For more info, turn on debug flags. See examples/llm_astra_inference_debug_graphviz.sh\n") + handle.write("\n".join(topology_lines)) + handle.write("\n") + for line in topology_lines: + print(line) print("LLM inference results written to {}".format(output_path)) warning_message = tc_inf.memory_capacity_warning() if warning_message: diff --git a/simulate_LLM.py b/simulate_LLM.py index 46a5c52c..bb071f16 100644 --- a/simulate_LLM.py +++ b/simulate_LLM.py @@ -973,85 +973,96 @@ def construct_transformer_graph(self, direction: str = "both"): if not gemm_entries: raise ValueError("Transformer GEMM times not provided") - tp_degree = self.tp + tp_degree = max(1, int(self.tp)) + cp_degree = max(1, int(self.cp)) root = Data_batch("transformer_root", 0, 0) op_id = 0 - for rank in range(tp_degree): - previous = root - - if direction in {"forward", "both"}: - for idx, entry in enumerate(gemm_entries): - entry_name = entry.get("name", f"g{idx}") - forward_cfg = entry.get("forward", {}) - fwd_duration = forward_cfg.get("duration") - if fwd_duration is None: - raise ValueError("Transformer GEMM entry missing forward duration") - - node = Node( - name=f"{entry_name}_fwd_rank{rank}", - op_id=op_id, - hw_id=rank, - duration=fwd_duration, - fwd=True, - ) - op_id += 1 - previous.add_child(node) - previous = node - comm_key = forward_cfg.get("comm_keys",None) - if comm_key: - if len(comm_key) > 1: - print(f"*CONSTRUCT INFO: Multiple comm keys for {entry_name}") - print(f"*CONSTRUCT INFO: Comm keys: {comm_key}") - raise ValueError(f"Multiple comm keys for {entry_name}") - - comm_key = comm_key[0] - if comm_key not in self.comm_metadata: - raise KeyError(f"Missing transformer comm metadata for key '{comm_key}'") - comm_edge = self.create_comm_edge( - name=comm_key, + for cp_idx in range(cp_degree): + for tp_idx in range(tp_degree): + rank = tp_idx + cp_idx * tp_degree + previous = root + + if direction in {"forward", "both"}: + for idx, entry in enumerate(gemm_entries): + entry_name = entry.get("name", f"g{idx}") + forward_cfg = entry.get("forward", {}) + fwd_duration = forward_cfg.get("duration") + if fwd_duration is None: + raise ValueError("Transformer GEMM entry missing forward duration") + + node = Node( + name=f"{entry_name}_fwd_rank{rank}", op_id=op_id, - comm_key=comm_key, - is_dp=False, - local_hw_id=rank, + hw_id=rank, + duration=fwd_duration, + fwd=True, ) + node.tp_rank = tp_idx + node.cp_rank = cp_idx op_id += 1 - previous.add_child(comm_edge) - previous = comm_edge - - if direction in {"backward", "both"}: - for idx, entry in enumerate(reversed(gemm_entries)): - entry_name = entry.get("name", f"g{idx}") - backward_cfg = entry.get("backward", {}) - bwd_duration = backward_cfg.get("duration") - if bwd_duration is None: - raise ValueError("Transformer GEMM entry missing backward duration") - - node = Node( - name=f"{entry_name}_bwd_rank{rank}", - op_id=op_id, - hw_id=rank, - duration=bwd_duration, - fwd=False, - ) - op_id += 1 - previous.add_child(node) - previous = node - - for comm_idx, comm_key in enumerate(backward_cfg.get("comm_keys", [])): - if comm_key not in self.comm_metadata: - raise KeyError(f"Missing transformer comm metadata for key '{comm_key}'") - comm_edge = self.create_comm_edge( - name=comm_key, + previous.add_child(node) + previous = node + comm_key = forward_cfg.get("comm_keys", None) + if comm_key: + if len(comm_key) > 1: + print(f"*CONSTRUCT INFO: Multiple comm keys for {entry_name}") + print(f"*CONSTRUCT INFO: Comm keys: {comm_key}") + raise ValueError(f"Multiple comm keys for {entry_name}") + + comm_key = comm_key[0] + if comm_key not in self.comm_metadata: + raise KeyError(f"Missing transformer comm metadata for key '{comm_key}'") + comm_edge = self.create_comm_edge( + name=comm_key, + op_id=op_id, + comm_key=comm_key, + is_dp=False, + local_hw_id=rank, + ) + comm_edge.tp_rank = tp_idx + comm_edge.cp_rank = cp_idx + op_id += 1 + previous.add_child(comm_edge) + previous = comm_edge + + if direction in {"backward", "both"}: + for idx, entry in enumerate(reversed(gemm_entries)): + entry_name = entry.get("name", f"g{idx}") + backward_cfg = entry.get("backward", {}) + bwd_duration = backward_cfg.get("duration") + if bwd_duration is None: + raise ValueError("Transformer GEMM entry missing backward duration") + + node = Node( + name=f"{entry_name}_bwd_rank{rank}", op_id=op_id, - comm_key=comm_key, - is_dp=False, - local_hw_id=rank, + hw_id=rank, + duration=bwd_duration, + fwd=False, ) + node.tp_rank = tp_idx + node.cp_rank = cp_idx op_id += 1 - previous.add_child(comm_edge) - previous = comm_edge + previous.add_child(node) + previous = node + + for comm_idx, comm_key in enumerate(backward_cfg.get("comm_keys", [])): + if comm_key not in self.comm_metadata: + raise KeyError(f"Missing transformer comm metadata for key '{comm_key}'") + comm_edge = self.create_comm_edge( + name=comm_key, + op_id=op_id, + comm_key=comm_key, + is_dp=False, + local_hw_id=rank, + ) + comm_edge.tp_rank = tp_idx + comm_edge.cp_rank = cp_idx + op_id += 1 + previous.add_child(comm_edge) + previous = comm_edge return root diff --git a/util.py b/util.py index 3afcf49e..772b28dd 100644 --- a/util.py +++ b/util.py @@ -12,6 +12,100 @@ DF_CACHE_DIR = "./cache" + +def _collect_parallelism_values(hw_config): + sch_config = getattr(hw_config, "sch_config", None) + if sch_config is None: + return {} + + values = {} + for name in getattr(sch_config, "_fields", []): + values[str(name).lower()] = getattr(sch_config, name) + return values + + +def _format_parallelism_terms(dim, parallelism_values): + terms = [] + for axis in getattr(dim, "parallelisms", ()): + if axis not in parallelism_values: + continue + factor = parallelism_values[axis] + try: + factor_value = int(factor) + except (TypeError, ValueError): + factor_value = factor + terms.append(f"{axis} {factor_value}") + return terms + + +def network_topology_summary_training(hw_config): + parallelism_values = _collect_parallelism_values(hw_config) + dimensions = list(getattr(hw_config.network_layout, "dimensions", ())) + lines = [f"Network Topology [dims={len(dimensions)}]"] + aggregate = 1 + + for dim in dimensions: + terms = _format_parallelism_terms(dim, parallelism_values) + axis_repr = " × ".join(terms) if terms else "(none)" + size_value = getattr(dim, "size", 1) + try: + size_int = int(size_value) + except (TypeError, ValueError): + size_int = None + size_display = size_value + else: + aggregate *= size_int + size_display = size_int + lines.append( + f" • {dim.id} {dim.label} : {axis_repr} ⇒ size {size_display}" + ) + + lines.append(f" => aggregate = {aggregate} GPUs") + return lines + + +def network_topology_summary_inference(hw_config): + parallelism_values = _collect_parallelism_values(hw_config) + all_dimensions = list(getattr(hw_config.network_layout, "dimensions", ())) + filtered_dimensions = [ + dim + for dim in all_dimensions + if any(axis != "dp" for axis in getattr(dim, "parallelisms", ())) or not dim.parallelisms + ] + lines = [f"Network Topology [dims={len(filtered_dimensions)}]"] + + aggregate_per_replica = 1 + for dim in filtered_dimensions: + terms = _format_parallelism_terms(dim, parallelism_values) + axis_repr = " × ".join(terms) if terms else "(none)" + size_value = getattr(dim, "size", 1) + try: + size_int = int(size_value) + except (TypeError, ValueError): + size_int = None + size_display = size_value + else: + aggregate_per_replica *= size_int + size_display = size_int + lines.append( + f" • {dim.id} {dim.label} : {axis_repr} ⇒ size {size_display}" + ) + + dp_factor = parallelism_values.get("dp", 1) + try: + dp_replicas = max(1, int(dp_factor)) + except (TypeError, ValueError): + dp_replicas = dp_factor if dp_factor else 1 + if isinstance(aggregate_per_replica, (int, float)) and isinstance(dp_replicas, (int, float)): + total_aggregate = aggregate_per_replica * dp_replicas + else: + total_aggregate = aggregate_per_replica + lines.append(f" replicas (dp): {dp_replicas}") + lines.append( + f" => aggregate = {aggregate_per_replica} GPUs per replica ({total_aggregate} total)" + ) + return lines + def printError(message): sys.exit(message) From bfdd90051cc442d3cf096fbc30103c39700ca3c0 Mon Sep 17 00:00:00 2001 From: gkarfakis19 Date: Mon, 3 Nov 2025 16:26:05 -0800 Subject: [PATCH 13/77] remap send and recv nodes for more astrasim stability --- astrasim_lib/executor.py | 7 +++++++ 1 file changed, 7 insertions(+) diff --git a/astrasim_lib/executor.py b/astrasim_lib/executor.py index fd2789f3..feb5eb22 100644 --- a/astrasim_lib/executor.py +++ b/astrasim_lib/executor.py @@ -283,12 +283,19 @@ def _is_control_send(node: pb.Node) -> bool: node.type == pb.COMM_SEND_NODE and (node.name or "").endswith("_send_control") ) + def _is_control_recv(node: pb.Node) -> bool: + return ( + node.type == pb.COMM_RECV_NODE + and (node.name or "").endswith("_recv_control") + ) control_nodes: List[pb.Node] = [] regular_nodes: List[pb.Node] = [] for node in self.nodes: if _is_control_send(node): control_nodes.append(node) + elif _is_control_recv(node): + control_nodes.append(node) else: regular_nodes.append(node) From 3557dcfd680abf545c54fc351c1687d05272745b Mon Sep 17 00:00:00 2001 From: gkarfakis19 Date: Tue, 4 Nov 2025 14:17:39 -0800 Subject: [PATCH 14/77] bugfix --- README.md | 1 + simulate_LLM.py | 20 ++++++++++---------- 2 files changed, 11 insertions(+), 10 deletions(-) diff --git a/README.md b/README.md index 978018e6..20631f03 100644 --- a/README.md +++ b/README.md @@ -33,6 +33,7 @@ If uv is used, make sure to use `uv run` instead of `python` to execute python s - `git submodule update --remote -- astra-sim` - `ASTRA_SIM=$(realpath ./astra-sim)` - `cd ${ASTRA_SIM}` +- `git submodule update --remote -- extern/network_backend/analytical` - `./build/astra_analytical/build.sh` - `cd ..` diff --git a/simulate_LLM.py b/simulate_LLM.py index bb071f16..0afd6a84 100644 --- a/simulate_LLM.py +++ b/simulate_LLM.py @@ -1135,7 +1135,7 @@ def simulate(self, root): print("child {} ready at time {} ".format(child.name, time)) if isinstance(event, Node): - GPU_list[event.hw_id] = True + GPU_list[int(event.hw_id)] = True @@ -1157,7 +1157,7 @@ def simulate(self, root): ready_list.remove(event) elif isinstance(event, Node): - if GPU_list[event.hw_id] == True: + if GPU_list[int(event.hw_id)] == True: new_time = time + event.duration heappush(event_queue, (new_time, counter, event)) event.scheduled = True @@ -1165,7 +1165,7 @@ def simulate(self, root): if debug: print("{}.{} enqueued at time {} at device {}".format(event.name, event.op_id, time, event.hw_id)) counter = counter + 1 - GPU_list[event.hw_id] = False + GPU_list[int(event.hw_id)] = False ready_list.remove(event) elif isinstance(event, Edge): new_time = time + event.duration @@ -1436,20 +1436,20 @@ def _count_transformer_layers_per_device(root_obj: Any, num_devices: int) -> Lis print("child {} ready at time {} ".format(child.name, time)) if isinstance(event, Node): - GPU_list[event.hw_id] = True + GPU_list[int(event.hw_id)] = True is_transformer_block = _is_transformer_block(event) #TODO: embedding and softmax layers if mode == "training": if not event.fwd: if is_transformer_block: - memory_snapshot.release_activation(event.hw_id, event, time) + memory_snapshot.release_activation(int(event.hw_id), event, time) elif event.fwd: if is_transformer_block: - memory_snapshot.allocate_activation(event.hw_id, event, time) + memory_snapshot.allocate_activation(int(event.hw_id), event, time) elif mode == "inference": if event.fwd: if is_transformer_block: - memory_snapshot.allocate_activation(event.hw_id, event, time) - memory_snapshot.release_activation(event.hw_id, event, time) + memory_snapshot.allocate_activation(int(event.hw_id), event, time) + memory_snapshot.release_activation(int(event.hw_id), event, time) for event in ready_list[:]: @@ -1468,7 +1468,7 @@ def _count_transformer_layers_per_device(root_obj: Any, num_devices: int) -> Lis ready_list.remove(event) elif isinstance(event, Node): - if GPU_list[event.hw_id] == True: + if GPU_list[int(event.hw_id)] == True: new_time = time + event.duration heappush(event_queue, (new_time, counter, event)) event.scheduled = True @@ -1476,7 +1476,7 @@ def _count_transformer_layers_per_device(root_obj: Any, num_devices: int) -> Lis if debug: print("{}.{} enqueued at time {} at device {}".format(event.name, event.op_id, time, event.hw_id)) counter = counter + 1 - GPU_list[event.hw_id] = False + GPU_list[int(event.hw_id)] = False ready_list.remove(event) elif isinstance(event, Edge): new_time = time + event.duration From 35eb723f871c5517359a3a9a2734c5282b3f0cfb Mon Sep 17 00:00:00 2001 From: gkarfakis19 Date: Tue, 4 Nov 2025 14:20:59 -0800 Subject: [PATCH 15/77] faulty links enabled w/ new astra ver --- TODO_11_3_2025.md | 3 - astrasim_lib/config_generation.py | 107 ++++-- astrasim_lib/executor.py | 33 +- astrasim_lib/integration.py | 317 ++++++++-------- config.py | 47 +++ pyproject.toml | 1 + time_calculation.py | 10 +- time_calculation_LLM.py | 26 +- tools/parallelism_sweep.py | 589 ++++++++++++++++++++++++++++++ uv.lock | 15 +- 10 files changed, 948 insertions(+), 200 deletions(-) delete mode 100644 TODO_11_3_2025.md create mode 100644 tools/parallelism_sweep.py diff --git a/TODO_11_3_2025.md b/TODO_11_3_2025.md deleted file mode 100644 index a066c6c9..00000000 --- a/TODO_11_3_2025.md +++ /dev/null @@ -1,3 +0,0 @@ -a) size=auto for network -b) enable decode len = 0 -c) debug hierarchical dp in last dim (3D) vs dp with lp (2D) \ No newline at end of file diff --git a/astrasim_lib/config_generation.py b/astrasim_lib/config_generation.py index 4bb56439..78d9cb3d 100644 --- a/astrasim_lib/config_generation.py +++ b/astrasim_lib/config_generation.py @@ -121,6 +121,10 @@ def generate_astrasim_configs_from_hw( if not dimensions: raise ValueError("Hardware config is missing network dimensions required for AstraSim integration.") + exec_backend = getattr(hw_obj, "execution_backend", None) + astra_cfg = getattr(exec_backend, "astra", None) if exec_backend else None + astra_mode = str(getattr(astra_cfg, "mode", "") or "").strip().lower() + sch_config = getattr(hw_obj, "sch_config", None) def _safe_int(value: Any, default: int = 1) -> int: try: @@ -146,9 +150,9 @@ def _safe_int(value: Any, default: int = 1) -> int: allowed_axes = set(axis_sizes.keys()) if axes_filter else None # print(f"Allowed axes: {allowed_axes}") # print(f"Axes sizes: {axis_sizes}") - dim_infos: List[Tuple[Any, List[str], int]] = [] + dim_infos: List[Tuple[Any, List[str], int, int]] = [] # print(f"Filter: {axes_filter}") - for dim in dimensions: + for dim_idx, dim in enumerate(dimensions): axes = [str(axis).strip().lower() for axis in getattr(dim, "parallelisms", ())] # print(f"All axes for dimension {dim.label}: {axes}") if allowed_axes is not None and axes_filter: @@ -163,44 +167,74 @@ def _safe_int(value: Any, default: int = 1) -> int: "Supported axes are tp, cp, lp, dp." ) effective *= axis_sizes[axis] - dim_infos.append((dim, axes, effective)) - - active_dim_indices = [idx for idx, (dim, _, _) in enumerate(dim_infos) if int(getattr(dim, "size", 0)) > 1] + dim_infos.append((dim, axes, effective, dim_idx)) + + active_dim_indices = [ + dim_idx + for dim, _, _, dim_idx in dim_infos + if int(getattr(dim, "size", 0) or 0) > 1 + ] + allowed_fault_dim_indices: set[int] = set() if active_dim_indices: last_active_idx = active_dim_indices[-1] + if astra_mode in {"full_astrasim_hierarchical", "hybrid"}: + allowed_fault_dim_indices.add(0) + allowed_fault_dim_indices.add(last_active_idx) + else: + allowed_fault_dim_indices.add(last_active_idx) dp_in_active = False - for idx, (dim, axes, _) in enumerate(dim_infos): - if "dp" in axes: - if idx != last_active_idx: - raise ValueError( - f"Data-parallel axis must be assigned to the last active network dimension. " - f"Dimension '{dim.label}' (index {idx}) includes 'dp', but last active dimension index is {last_active_idx}." - ) - dp_in_active = True + if axis_sizes.get("dp", 1) > 1: + for dim, axes, _, dim_idx in dim_infos: + if "dp" in axes: + if dim_idx != last_active_idx: + raise ValueError( + f"Data-parallel axis must be assigned to the last active network dimension. " + f"Dimension '{dim.label}' (index {dim_idx}) includes 'dp', but last active dimension index is {last_active_idx}." + ) + dp_in_active = True if axis_sizes.get("dp", 1) > 1 and not dp_in_active: raise ValueError( "Hardware config declares dp > 1, but no network dimension with size > 1 includes the 'dp' axis." ) + else: + if any(getattr(dim, "faulty_links", ()) for dim, _, _, _ in dim_infos): + raise ValueError( + "Faulty links require at least one active (size > 1) network dimension." + ) + if allowed_fault_dim_indices: + valid_indices = {idx for idx in allowed_fault_dim_indices if 0 <= idx < len(dimensions)} + for dim, _, _, dim_idx in dim_infos: + if getattr(dim, "faulty_links", ()): + if dim_idx not in valid_indices: + raise ValueError( + f"Faulty links for dimension '{dim.label}' (index {dim_idx}) are only permitted on " + f"dimension indices {sorted(valid_indices)} for mode '{astra_mode or 'unspecified'}'." + ) target = int(npus_count) + axes_needed: List[str] = [] remaining = target + for axis in axis_order_preference: size = axis_sizes.get(axis, 1) + if size <= 1: continue + if remaining % size == 0: axes_needed.append(axis) remaining //= size + if remaining != 1: raise ValueError( f"Unable to map requested npus_count={target} to network axes {axis_sizes}." ) axes_needed_set = set(axes_needed) - selected_dims: List[Tuple[Any, List[str], int]] = [] + selected_dims: List[Tuple[Any, List[str], int, int]] = [] accumulated = 1 - for dim, axes, _effective in dim_infos: + for dim, axes, _effective, dim_idx in dim_infos: selected_axes: List[str] = [] size_contrib = 1 for axis in axes: @@ -211,14 +245,14 @@ def _safe_int(value: Any, default: int = 1) -> int: continue axes_needed_set.difference_update(selected_axes) accumulated *= size_contrib - selected_dims.append((dim, selected_axes, size_contrib)) + selected_dims.append((dim, selected_axes, size_contrib, dim_idx)) if axes_needed_set: raise ValueError( f"Requested npus_count={target} requires axes {axes_needed} but no network dimension covers {axes_needed_set}." ) if accumulated != target: raise ValueError( - f"Selected network dimensions {[d.label for d, _, _ in selected_dims]} " + f"Selected network dimensions {[d.label for d, _, _, _ in selected_dims]} " f"describe {accumulated} NPUs but simulation expects {target} ranks." ) @@ -227,7 +261,7 @@ def _safe_int(value: Any, default: int = 1) -> int: bw_list: List[float] = [] lat_list: List[float] = [] - for dim, axes_selected, product_size in selected_dims: + for dim, axes_selected, product_size, _ in selected_dims: topo = _normalize_topology_name(dim.topology_type) size = product_size if topo == "Ring" and size <= 2: @@ -239,7 +273,7 @@ def _safe_int(value: Any, default: int = 1) -> int: bw_list.append(round(_gbps_from_bps(effective_bw), 6)) lat_list.append(round(_ns_from_s(latency_s), 3)) - signature_parts = [str(size) for _dim, _axes, size in selected_dims] + signature_parts = [str(size) for _dim, _axes, size, _ in selected_dims] dim_signature = "_".join(signature_parts) if signature_parts else f"{target}" net_yaml = os.path.join(out_dir, f"network_analytical_{dim_signature}.yml") @@ -257,8 +291,32 @@ def _safe_int(value: Any, default: int = 1) -> int: f"latency: [ {lat_str} ] # ns\n" ) + faulty_links_tuple: Tuple[Tuple[int, int, float], ...] = tuple() + if selected_dims: + dims_with_faults = [ + (dim_idx, getattr(dim, "faulty_links", ())) + for dim, _axes, _size, dim_idx in selected_dims + if getattr(dim, "faulty_links", ()) + ] + if dims_with_faults: + dims_with_faults.sort(key=lambda item: item[0]) + _, faulty_links_tuple = dims_with_faults[-1] + faulty_links_tuple = tuple(faulty_links_tuple) + fault_entries = [ + f"[{src}, {dst}, {float(weight):g}]" + for src, dst, weight in faulty_links_tuple + ] + faulty_links_str = f"[{', '.join(fault_entries)}]" + net_content += f"faulty_links: {faulty_links_str}\n" + os.makedirs(os.path.dirname(net_yaml), exist_ok=True) - cache_key = (tuple(topo_list), tuple(npus_list), tuple(bw_list), tuple(lat_list)) + cache_key = ( + tuple(topo_list), + tuple(npus_list), + tuple(bw_list), + tuple(lat_list), + faulty_links_tuple, + ) cached_path = _NET_YAML_CACHE.get(cache_key) need_write = True if cached_path and os.path.exists(cached_path): @@ -277,10 +335,9 @@ def _safe_int(value: Any, default: int = 1) -> int: os.replace(tmp_path, net_yaml) _NET_YAML_CACHE[cache_key] = net_yaml - exec_backend = getattr(hw_obj, "execution_backend", None) - if exec_backend and exec_backend.astra: - coll = exec_backend.astra.collectives - sys_opts = getattr(exec_backend.astra, "sys_options", None) + if astra_cfg: + coll = astra_cfg.collectives + sys_opts = getattr(astra_cfg, "sys_options", None) default_ag = coll.all_gather default_ar = coll.all_reduce default_rs = coll.reduce_scatter @@ -307,7 +364,7 @@ def _collective_for_dimension(dim, topo_name: str, op: str, default_alg: str) -> rs_impl: List[str] = [] a2a_impl: List[str] = [] - for (dim, axes, _), topo_name in zip(selected_dims, topo_list): + for (dim, axes, _, _), topo_name in zip(selected_dims, topo_list): ag_impl.append(_collective_for_dimension(dim, topo_name, "all-gather", default_ag)) ar_impl.append(_collective_for_dimension(dim, topo_name, "all-reduce", default_ar)) rs_impl.append(_collective_for_dimension(dim, topo_name, "reduce-scatter", default_rs)) diff --git a/astrasim_lib/executor.py b/astrasim_lib/executor.py index feb5eb22..80e92ba8 100644 --- a/astrasim_lib/executor.py +++ b/astrasim_lib/executor.py @@ -1567,19 +1567,26 @@ def run_astra_simulation_only_onepath( # Run AstraSim simulation on forward graph (cached via manifest) print(f"[AstraSim] Executing forward simulation with {rank_count} ranks...") - fwd_times, fwd_total = run_cache_astrasim( - time_calc_obj.hw_config, - comm="graph", - npus_count=rank_count, - size_bytes=0, - astra_config_dir="./astra_cache", - cache_path="./astra_cache/cache.json", - manifest_json_path=fwd_manifest, - workload_prefix=fwd_et_prefix, - comm_group_json=comm_groups_path, - axes_filter=axes_filter, - files=astra_configs, - ) + cache_override_env = os.environ.pop("ASTRA_CACHE_DIR", None) + local_cache_dir = work_dir + local_cache_path = os.path.join(work_dir, "cache.json") + try: + fwd_times, fwd_total = run_cache_astrasim( + time_calc_obj.hw_config, + comm="graph", + npus_count=rank_count, + size_bytes=0, + astra_config_dir=local_cache_dir, + cache_path=local_cache_path, + manifest_json_path=fwd_manifest, + workload_prefix=fwd_et_prefix, + comm_group_json=comm_groups_path, + axes_filter=axes_filter, + files=astra_configs, + ) + finally: + if cache_override_env is not None: + os.environ["ASTRA_CACHE_DIR"] = cache_override_env conversion_and_sim_time = time.time() - astrasim_start diff --git a/astrasim_lib/integration.py b/astrasim_lib/integration.py index 851dcba7..6b942d79 100644 --- a/astrasim_lib/integration.py +++ b/astrasim_lib/integration.py @@ -6,7 +6,9 @@ import json import os import re +import shutil import subprocess +import tempfile import time from typing import Any, Dict, Iterable, List, Optional, Sequence, Tuple @@ -323,169 +325,178 @@ def run_cache_astrasim( if isinstance(hw_obj, str): raise TypeError("run_cache_astrasim expects a parsed HWConfig object, not a path") - override_dir = os.environ.get("ASTRA_CACHE_DIR") - if override_dir: - astra_config_dir = override_dir - cache_path = os.path.join(override_dir, "cache.json") - cache_mode = _cache_mode() allow_read = cache_mode in {"CACHE_READONLY", "CACHE_READWRITE"} allow_write = cache_mode == "CACHE_READWRITE" - if allow_read or allow_write: - ensure_cache_file_exists(cache_path) - - (intra_ib_bps, intra_ll_s), _ = compute_intra_inter_ib_ll_from_hw(hw_obj) - if not intra_ib_bps or intra_ib_bps <= 0 or not intra_ll_s or intra_ll_s <= 0: - raise ValueError("Invalid intra-node IB/LL computed from HW config") - ib_gbps = round((intra_ib_bps) / float(1 << 30), 6) - ll_ns = round(intra_ll_s * 1e9, 3) - topo = derive_topology_from_hw(hw_obj) - colls = _collectives_from_hw(hw_obj, topo) - sys_opts_sig = _sys_options_from_hw(hw_obj) - axes_filter_tuple: Optional[Tuple[str, ...]] = tuple(axes_filter) if axes_filter else None - - sig: Dict[str, Any] = { - "comm": comm.lower(), - "npus": int(npus_count), - "size_bytes": int(size_bytes), - "topology": topo, - "physical_topology": topo, - "ib_gbps": ib_gbps, - "ll_ns": ll_ns, - "collectives": colls, - "backend": "analytical", - } - if sys_opts_sig is not None: - sig["sys_options"] = sys_opts_sig - if manifest_json_path: - sig["multinode"] = True - if axes_filter_tuple is not None: - sig["axes_filter"] = axes_filter_tuple - if not files: - files = generate_astrasim_configs_from_hw( - hw_obj, - out_dir=astra_config_dir, - npus_count=npus_count, - axes_filter=axes_filter_tuple, - ) - topo_list = files.get("topology_list") - npus_per_dim = files.get("npus_per_dim") - if topo_list: - sig["topology"] = list(topo_list) - sig["physical_topology"] = list(topo_list) - if npus_per_dim: - sig["npus_per_dim"] = list(npus_per_dim) - - canonical = _canonical_sig(sig) - - if workload_prefix: - pass - elif not manifest_json_path: - label = size_label(size_bytes) - base_dir = os.path.join(astra_config_dir, "workload", comm.lower(), f"{npus_count}npus_{label}") - workload_prefix = os.path.join(base_dir, f"{comm.lower()}_{label}") - expected = [f"{workload_prefix}.{r}.et" for r in range(npus_count)] - if not all(os.path.exists(p) for p in expected): - generate_workload_et(comm, npus_count, size_bytes, astra_config_dir=astra_config_dir) + temp_cache_dir: Optional[str] = None + if cache_mode == "NO_CACHE": + temp_cache_dir = tempfile.mkdtemp(prefix="astrasim_cache_") + astra_config_dir = temp_cache_dir + cache_path = os.path.join(temp_cache_dir, "cache.json") else: - workload_prefix = os.path.splitext(manifest_json_path)[0] + override_dir = os.environ.get("ASTRA_CACHE_DIR") + if override_dir: + astra_config_dir = override_dir + cache_path = os.path.join(override_dir, "cache.json") + if allow_read or allow_write: + ensure_cache_file_exists(cache_path) - remote_mem_path = get_remote_memory_path() - - if manifest_json_path: - bundle_list: List[str] = [manifest_json_path] - for path in (files["system_json"], files["network_yaml"]): - if path and os.path.exists(path): - bundle_list.append(path) - if remote_mem_path and os.path.exists(remote_mem_path): - bundle_list.append(remote_mem_path) - if comm_group_json and os.path.exists(comm_group_json): - bundle_list.append(comm_group_json) - cache_key = _hash_file_bundle(bundle_list) - else: - cache_key = _hash_sig(canonical) - - if allow_read: - cache = _load_cache(cache_path) - entry = cache.get(cache_key) if cache else None - if entry: - cached_per = entry.get("per_node_sec", []) - cached_max = float(entry.get("max_sec", 0.0)) - if cached_per and cached_max > 0 and len(cached_per) == int(npus_count) and all((t > 0 for t in cached_per)): - if str(comm).lower() == "graph": + try: + (intra_ib_bps, intra_ll_s), _ = compute_intra_inter_ib_ll_from_hw(hw_obj) + if not intra_ib_bps or intra_ib_bps <= 0 or not intra_ll_s or intra_ll_s <= 0: + raise ValueError("Invalid intra-node IB/LL computed from HW config") + ib_gbps = round((intra_ib_bps) / float(1 << 30), 6) + ll_ns = round(intra_ll_s * 1e9, 3) + topo = derive_topology_from_hw(hw_obj) + colls = _collectives_from_hw(hw_obj, topo) + sys_opts_sig = _sys_options_from_hw(hw_obj) + axes_filter_tuple: Optional[Tuple[str, ...]] = tuple(axes_filter) if axes_filter else None + + sig: Dict[str, Any] = { + "comm": comm.lower(), + "npus": int(npus_count), + "size_bytes": int(size_bytes), + "topology": topo, + "physical_topology": topo, + "ib_gbps": ib_gbps, + "ll_ns": ll_ns, + "collectives": colls, + "backend": "analytical", + } + if sys_opts_sig is not None: + sig["sys_options"] = sys_opts_sig + if manifest_json_path: + sig["multinode"] = True + if axes_filter_tuple is not None: + sig["axes_filter"] = axes_filter_tuple + if not files: + files = generate_astrasim_configs_from_hw( + hw_obj, + out_dir=astra_config_dir, + npus_count=npus_count, + axes_filter=axes_filter_tuple, + ) + topo_list = files.get("topology_list") + npus_per_dim = files.get("npus_per_dim") + if topo_list: + sig["topology"] = list(topo_list) + sig["physical_topology"] = list(topo_list) + if npus_per_dim: + sig["npus_per_dim"] = list(npus_per_dim) + + canonical = _canonical_sig(sig) + + if workload_prefix: + pass + elif not manifest_json_path: + label = size_label(size_bytes) + base_dir = os.path.join(astra_config_dir, "workload", comm.lower(), f"{npus_count}npus_{label}") + workload_prefix = os.path.join(base_dir, f"{comm.lower()}_{label}") + expected = [f"{workload_prefix}.{r}.et" for r in range(npus_count)] + if not all(os.path.exists(p) for p in expected): + generate_workload_et(comm, npus_count, size_bytes, astra_config_dir=astra_config_dir) + else: + workload_prefix = os.path.splitext(manifest_json_path)[0] + + remote_mem_path = get_remote_memory_path() + + if manifest_json_path: + bundle_list: List[str] = [manifest_json_path] + for path in (files["system_json"], files["network_yaml"]): + if path and os.path.exists(path): + bundle_list.append(path) + if remote_mem_path and os.path.exists(remote_mem_path): + bundle_list.append(remote_mem_path) + if comm_group_json and os.path.exists(comm_group_json): + bundle_list.append(comm_group_json) + cache_key = _hash_file_bundle(bundle_list) + else: + cache_key = _hash_sig(canonical) + + if allow_read: + cache = _load_cache(cache_path) + entry = cache.get(cache_key) if cache else None + if entry: + cached_per = entry.get("per_node_sec", []) + cached_max = float(entry.get("max_sec", 0.0)) + if cached_per and cached_max > 0 and len(cached_per) == int(npus_count) and all((t > 0 for t in cached_per)): + if str(comm).lower() == "graph": + try: + print( + f"[AstraSim] Cache HIT: comm={comm}, npus={npus_count}, size={size_bytes}, total={cached_max:.6f}s" + ) + except Exception: + pass + return cached_per, cached_max + if allow_write and cache_key in cache: try: - print( - f"[AstraSim] Cache HIT: comm={comm}, npus={npus_count}, size={size_bytes}, total={cached_max:.6f}s" - ) + cache.pop(cache_key) + _save_cache(cache_path, cache) except Exception: pass - return cached_per, cached_max - if allow_write and cache_key in cache: - try: - cache.pop(cache_key) - _save_cache(cache_path, cache) - except Exception: - pass - - attempts = 0 - per_node_sec: List[float] = [] - max_sec: float = 0.0 - last_outcome_zero = False - - # Capture command details for debugging - bin_path = _astrasim_binary_path() - debug_cmd = [ - bin_path, - f"--workload-configuration={workload_prefix}", - f"--system-configuration={files['system_json']}", - f"--network-configuration={files['network_yaml']}", - ] - if remote_mem_path: - debug_cmd.append(f"--remote-memory-configuration={remote_mem_path}") - if comm_group_json and os.path.exists(comm_group_json): - debug_cmd.append(f"--comm-group-configuration={comm_group_json}") - cmd_str = " ".join(debug_cmd) - # print(f"Full AstraSim command:\n{cmd_str}") - - while attempts < 5: - attempts += 1 - per_node_sec, max_sec = run_astrasim_analytical( - workload_prefix=workload_prefix, - system_json=files["system_json"], - network_yaml=files["network_yaml"], - remote_memory_json=remote_mem_path, - comm_group_json=comm_group_json, - ) - if per_node_sec and max_sec > 0 and len(per_node_sec) == int(npus_count) and all((t > 0 for t in per_node_sec)): - last_outcome_zero = False - break - last_outcome_zero = True - time.sleep(0.1) - if last_outcome_zero: - raise RuntimeError( - f"AstraSim returned zero time for {comm} size={size_bytes} npus={npus_count} after 5 retries.\n" - f"Full AstraSim command:\n{cmd_str}" - ) - - cache_entry = { - "signature": sig, - "canonical": canonical, - "per_node_sec": per_node_sec, - "max_sec": max_sec, - "workload_prefix": workload_prefix, - "system_json": files["system_json"], - "network_yaml": files["network_yaml"], - } - if manifest_json_path: - cache_entry["manifest_path"] = manifest_json_path - if allow_write: - cache = _load_cache(cache_path) - cache[cache_key] = cache_entry - _save_cache(cache_path, cache) - - return per_node_sec, max_sec + attempts = 0 + per_node_sec: List[float] = [] + max_sec: float = 0.0 + last_outcome_zero = False + + # Capture command details for debugging + bin_path = _astrasim_binary_path() + debug_cmd = [ + bin_path, + f"--workload-configuration={workload_prefix}", + f"--system-configuration={files['system_json']}", + f"--network-configuration={files['network_yaml']}", + ] + if remote_mem_path: + debug_cmd.append(f"--remote-memory-configuration={remote_mem_path}") + if comm_group_json and os.path.exists(comm_group_json): + debug_cmd.append(f"--comm-group-configuration={comm_group_json}") + cmd_str = " ".join(debug_cmd) + # print(f"Full AstraSim command:\n{cmd_str}") + + while attempts < 5: + attempts += 1 + per_node_sec, max_sec = run_astrasim_analytical( + workload_prefix=workload_prefix, + system_json=files["system_json"], + network_yaml=files["network_yaml"], + remote_memory_json=remote_mem_path, + comm_group_json=comm_group_json, + ) + if per_node_sec and max_sec > 0 and len(per_node_sec) == int(npus_count) and all((t > 0 for t in per_node_sec)): + last_outcome_zero = False + break + last_outcome_zero = True + time.sleep(0.1) + if last_outcome_zero: + raise RuntimeError( + f"AstraSim returned zero time for {comm} size={size_bytes} npus={npus_count} after 5 retries.\n" + f"Full AstraSim command:\n{cmd_str}" + ) + + cache_entry = { + "signature": sig, + "canonical": canonical, + "per_node_sec": per_node_sec, + "max_sec": max_sec, + "workload_prefix": workload_prefix, + "system_json": files["system_json"], + "network_yaml": files["network_yaml"], + } + if manifest_json_path: + cache_entry["manifest_path"] = manifest_json_path + + if allow_write: + cache = _load_cache(cache_path) + cache[cache_key] = cache_entry + _save_cache(cache_path, cache) + + return per_node_sec, max_sec + finally: + if temp_cache_dir is not None: + shutil.rmtree(temp_cache_dir, ignore_errors=True) __all__ = [ diff --git a/config.py b/config.py index f3fae6ea..00a0743e 100644 --- a/config.py +++ b/config.py @@ -354,6 +354,7 @@ class NetworkDimensionLayout: collective_override: Dict[str, str] = field(default_factory=dict) parallelisms: Tuple[str, ...] = field(default_factory=tuple) energy_per_bit: float = 0.0 + faulty_links: Tuple[Tuple[int, int, float], ...] = field(default_factory=tuple) @classmethod def from_raw( @@ -475,6 +476,51 @@ def from_raw( expected_product=computed_product, ) + faulty_links_raw = raw.get("faulty_links", []) + faulty_links: List[Tuple[int, int, float]] = [] + if faulty_links_raw: + if not isinstance(faulty_links_raw, Sequence) or isinstance(faulty_links_raw, (str, bytes)): + raise ValueError( + f"network dimension '{label}' faulty_links must be a sequence of [src, dst, weight] entries" + ) + for idx, entry in enumerate(faulty_links_raw): + if ( + not isinstance(entry, Sequence) + or isinstance(entry, (str, bytes)) + or len(entry) != 3 + ): + raise ValueError( + f"network dimension '{label}' faulty_links[{idx}] must be a three-item sequence [src, dst, weight]" + ) + src_raw, dst_raw, weight_raw = entry + try: + src = int(src_raw) + except (TypeError, ValueError) as exc: + raise ValueError( + f"network dimension '{label}' faulty_links[{idx}][0] must be an integer endpoint" + ) from exc + try: + dst = int(dst_raw) + except (TypeError, ValueError) as exc: + raise ValueError( + f"network dimension '{label}' faulty_links[{idx}][1] must be an integer endpoint" + ) from exc + if src < 0 or dst < 0: + raise ValueError( + f"network dimension '{label}' faulty_links[{idx}] endpoints must be >= 0" + ) + try: + weight = float(weight_raw) + except (TypeError, ValueError) as exc: + raise ValueError( + f"network dimension '{label}' faulty_links[{idx}][2] must be a numeric reliability weight" + ) from exc + if weight < 0.0 or weight > 1.0: + raise ValueError( + f"network dimension '{label}' faulty_links[{idx}] weight must be between 0.0 and 1.0" + ) + faulty_links.append((src, dst, weight)) + return cls( id=dim_id, label=label, @@ -486,6 +532,7 @@ def from_raw( collective_override=collective_override, parallelisms=tuple(normalized_parallelisms), energy_per_bit=energy_per_bit, + faulty_links=tuple(faulty_links), ) @property diff --git a/pyproject.toml b/pyproject.toml index 4580b16f..cde7ab12 100644 --- a/pyproject.toml +++ b/pyproject.toml @@ -17,5 +17,6 @@ dependencies = [ "pyyaml==6.0.2", "seaborn>=0.13.2", "six==1.17.0", + "tqdm>=4.67.1", "tzdata==2025.2", ] diff --git a/time_calculation.py b/time_calculation.py index 9fae111b..cac9dbf0 100644 --- a/time_calculation.py +++ b/time_calculation.py @@ -27,18 +27,18 @@ def __init__(self, hw_config, precision, kernel_overhead, roofline_cb, astra_pol self._roofline = roofline_cb self._astra_policy = astra_policy or "analytical" - def _astra_collective(self, kind: str, participants: int, size_bytes: int) -> float: + def _astra_collective(self, kind: str, participants: int, size_bytes: int, axis: Optional[str] = None) -> float: part = int(participants) if part <= 1 or size_bytes <= 0: return 0.0 byte_count = int(math.ceil(size_bytes)) + axes_filter = [str(axis).lower()] if axis else None _, max_sec = run_cache_astrasim( self.hw_config, comm=kind, npus_count=part, size_bytes=byte_count, - astra_config_dir="./astra_cache", - cache_path="./astra_cache/cache.json", + axes_filter=axes_filter, ) return float(max_sec) @@ -54,6 +54,7 @@ def collective( local_bytes: float = 0.0, local_ops: float = 0.0, debug_label: str = "", + axis: Optional[str] = None, ) -> float: if size_bytes <= 0: return 0.0 @@ -76,7 +77,8 @@ def collective( if kind in collective_ops or kind == "pipeline": # Pipeline uses 2 NPUs for point-to-point, others use part npus = 2 if kind == "pipeline" else part - network_time = self._astra_collective(kind, npus, network_bytes) + axis_filter = None if kind == "pipeline" else axis + network_time = self._astra_collective(kind, npus, network_bytes, axis_filter) else: raise ValueError(f"Unsupported collective operation: {kind}") else: diff --git a/time_calculation_LLM.py b/time_calculation_LLM.py index 89b41017..c152fffe 100644 --- a/time_calculation_LLM.py +++ b/time_calculation_LLM.py @@ -299,6 +299,7 @@ def get_tensor_reduction_time(self, total_bytes: int, kind: str, name: str, part ll=self.LLTP, local_bytes=0, debug_label=name or "comm", + axis="tp", ) return reduction_time @@ -494,6 +495,8 @@ def _tensor_context_hybrid_gemm_forward(self, gemm: Tuple[int, ...], name: str, if gemm_type is None: raise ValueError("gemm_type is required for tensor-context hybrid forward GEMM") + axis_hint = None + if gemm_type == GemmType.ATTENTION_SCORE: # attention gemm gemm_time = self.getGEMMTime(shard_m, k, n, name, disable_overhead=True)[0] * batch / max(1, self.tp) + self.O elif gemm_type == GemmType.ATTENTION_OUTPUT: # attention gemm @@ -503,16 +506,19 @@ def _tensor_context_hybrid_gemm_forward(self, gemm: Tuple[int, ...], name: str, total_bytes = self.get_kv_size_bytes() / self.tp # each tp group holds a tp shard of kv for each cp group kind = "all_gather" participants = self.cp # all gather K V for each cp group + axis_hint = "cp" elif gemm_type == GemmType.OUT_PROJ: gemm_time = self.getGEMMTime(shard_m, shard_k, n, name)[0] total_bytes = math.ceil(self.precision.activations * shard_m * n) kind = "reduce_scatter" participants = self.tp # reduce scatter output activation for each tp group + axis_hint = "tp" elif gemm_type == GemmType.FFN2: # row wise gemm_time = self.getGEMMTime(shard_m, shard_k, n, name)[0] total_bytes = math.ceil(self.precision.activations * shard_m * n) kind = "reduce_scatter" participants = self.tp # reduce scatter output activation for each tp group + axis_hint = "tp" elif gemm_type == GemmType.FFN1: # column wise gemm_time = self.getGEMMTime(shard_m, k, shard_n, name)[0] else: @@ -527,6 +533,7 @@ def _tensor_context_hybrid_gemm_forward(self, gemm: Tuple[int, ...], name: str, ll=self.LLTP, local_bytes=0, debug_label=name or "comm", + axis=axis_hint, ) return gemm_time, reduction_time, total_bytes def _tensor_context_hybrid_gemm_backward(self, gemm: Tuple[int, ...], name: str, gemm_type: Optional[GemmType] = None) -> Tuple[float, float]: @@ -565,12 +572,15 @@ def _tensor_context_hybrid_gemm_backward(self, gemm: Tuple[int, ...], name: str, if gemm_type is None: raise ValueError("gemm_type is required for tensor-context hybrid backward GEMM") + axis_hint = None + if gemm_type == GemmType.ATTENTION_SCORE: # attention gemm grad_time_act = self.getGEMMTime(shard_m, n, k, name, disable_overhead=True)[0] * batch / max(1, self.tp) + self.O grad_time_wt = self.getGEMMTime(k, shard_m, n, name, disable_overhead=True)[0] * batch / max(1, self.tp) + self.O total_bytes = self.precision.grad_communication * k * n * batch * 2 / self.tp # weight gradient of K V need to be reduce scattered *2 account for both attn key and value kind = "reduce_scatter" participants = self.cp + axis_hint = "cp" elif gemm_type == GemmType.ATTENTION_OUTPUT: # attention gemm grad_time_act = self.getGEMMTime(shard_m, n, k, name)[0] * batch / max(1, self.tp) grad_time_wt = self.getGEMMTime(k, shard_m, n, name)[0] * batch / max(1, self.tp) @@ -580,22 +590,25 @@ def _tensor_context_hybrid_gemm_backward(self, gemm: Tuple[int, ...], name: str, total_bytes = math.ceil(self.precision.grad_communication * shard_m * k) kind = "reduce_scatter" participants = self.tp + axis_hint = "tp" elif gemm_type == GemmType.OUT_PROJ: grad_time_act = self.getGEMMTime(shard_m, n, shard_k, name)[0] grad_time_wt = self.getGEMMTime(shard_k, shard_m, n, name)[0] total_bytes = self.get_kv_size_bytes() / self.tp kind = "all_gather" participants = self.cp + axis_hint = "cp" elif gemm_type == GemmType.FFN2: # row wise grad_time_act = self.getGEMMTime(shard_m, n, shard_k, name)[0] grad_time_wt = self.getGEMMTime(shard_k, shard_m, n, name)[0] elif gemm_type == GemmType.FFN1: # column wise - + grad_time_act = self.getGEMMTime(shard_m, shard_n, k, name)[0] grad_time_wt = self.getGEMMTime(k, shard_m, shard_n, name)[0] total_bytes = math.ceil(self.precision.grad_communication * shard_m * k) kind = "reduce_scatter" participants = self.tp + axis_hint = "tp" elif gemm_type == GemmType.LINEAR_SOFTMAX: grad_time_act = self.getGEMMTime(shard_m, n, shard_k, name)[0] grad_time_wt = self.getGEMMTime(shard_k, shard_m, n, name)[0] @@ -614,6 +627,7 @@ def _tensor_context_hybrid_gemm_backward(self, gemm: Tuple[int, ...], name: str, ll=self.LLTP, local_bytes=0, debug_label=name or "comm", + axis=axis_hint, ) return gemm_time, reduction_time, total_bytes def _tensor_parallelism_gemm_forward(self, gemm: Tuple[int, ...], name: str, gemm_type: Optional[GemmType] = None) -> Tuple[float, float]: @@ -793,6 +807,7 @@ def _context_parallelism_gemm_forward(self, gemm: Tuple[int, ...], name: str, ge ll=self.LLTP, local_bytes=0, debug_label=name or "comm", + axis="cp", ) return gemm_time, reduction_time, total_bytes @@ -845,6 +860,7 @@ def _context_parallelism_gemm_backward(self, gemm: Tuple[int, ...], name: str, g ll=self.LLTP, local_bytes=0, debug_label=name or "comm", + axis="cp", ) return gemm_time, reduction_time, total_bytes @@ -1128,6 +1144,7 @@ def get_layernorm_f(self, batch, seq_len, d_model, comm_after=False): ll=self.LLTP, local_bytes=0, debug_label="layernorm_f_all_gather", + axis="tp", ) else: reduction_time = 0.0 @@ -1164,6 +1181,7 @@ def get_layernorm_b(self, batch, seq_len, d_model, comm_after=False): ll=self.LLTP, local_bytes=0, debug_label="layernorm_b_all_gather", + axis="tp", ) else: @@ -1286,6 +1304,7 @@ def get_data_parallel_reduction_llm(self, d, intermediate_size): ll=self.LLD, local_bytes=0.0, debug_label="dp_zero_gradients", + axis="dp", ) gather_time = 0.0 @@ -1299,6 +1318,7 @@ def get_data_parallel_reduction_llm(self, d, intermediate_size): ll=self.LLD, local_bytes=0.0, debug_label="dp_zero_param_gather_zero2", + axis="dp", ) if self.debug: print(f"apply_grad_time (ZeRO-{self.zero_stage}): {apply_grad_time}") @@ -1317,6 +1337,7 @@ def get_data_parallel_reduction_llm(self, d, intermediate_size): ll=self.LLD, local_bytes=0.0, debug_label="qkv_proj reduction", + axis="dp", ) apply_grad_time += self.apply_grad(Dim0=d, Dim1=3*d, name="qkv_proj reduction") @@ -1329,6 +1350,7 @@ def get_data_parallel_reduction_llm(self, d, intermediate_size): ll=self.LLD, local_bytes=0.0, debug_label="output_proj reduction", + axis="dp", ) apply_grad_time += self.apply_grad(Dim0=d, Dim1=d, name="output_proj reduction") @@ -1342,6 +1364,7 @@ def get_data_parallel_reduction_llm(self, d, intermediate_size): ll=self.LLD, local_bytes=0.0, debug_label="ffn1 reduction", + axis="dp", ) apply_grad_time += self.apply_grad(Dim0=ffn1_dim, Dim1=d, name="ffn1 reduction") @@ -1354,6 +1377,7 @@ def get_data_parallel_reduction_llm(self, d, intermediate_size): ll=self.LLD, local_bytes=0.0, debug_label="ffn2 reduction", + axis="dp", ) apply_grad_time += self.apply_grad(Dim0=intermediate_size, Dim1=d, name="ffn2 reduction") diff --git a/tools/parallelism_sweep.py b/tools/parallelism_sweep.py new file mode 100644 index 00000000..69b09d4c --- /dev/null +++ b/tools/parallelism_sweep.py @@ -0,0 +1,589 @@ +#!/usr/bin/env python3 +""" +Parallelism sweep utility for DeepFlow LLM configurations. + +Update the global configuration section below to point at the desired hardware +and model config files and to tailor the parallelism search space. The tool +enumerates every combination, filters those whose total GPU count falls inside +the configured bounds, evaluates runtime with DeepFlow, and plots a scatter +chart of accuracy (by default, 1 / runtime) versus GPU count with horizontal +jitter to avoid overlap. +""" + +from __future__ import print_function + +import copy +import itertools +import math +import os +import random +import shutil +import sys +import tempfile +from concurrent.futures import ProcessPoolExecutor, as_completed +from contextlib import redirect_stdout, redirect_stderr +from typing import Dict, Iterable, List, Tuple + +import matplotlib.pyplot as plt +import yaml +from tqdm import tqdm + +sys.path.append(os.path.dirname(os.path.dirname(os.path.abspath(__file__)))) +import config +from time_calculation_LLM import TimeCalculationLLM +from LLM_util import process_gemm_shapes + + +# ----------------------------------------------------------------------------- +# Global configuration knobs (no CLI) +# ----------------------------------------------------------------------------- + +# Paths to the baseline configuration files +HARDWARE_CONFIG_PATH = "configs/hardware-config/a100_80GB.yaml" +MODEL_CONFIG_PATH = "configs/model-config/Llama3.1-405B.yaml" + +# Parallelism values to sweep (dense grid). Edit to suit your search space. +# Keys should match the entries under the YAML "parallelism" section. +PARALLELISM_SWEEP = { + "tp": [2**i for i in range(0, 9)], + "cp": [2**i for i in range(0, 9)], + "dp": [2**i for i in range(0, 9)], + "lp": [2**i for i in range(0, 9)], +} + +# Optional knobs that still live inside the parallelism section but do not +# affect GPU counts. Leave empty if you do not want to explore them. +OTHER_PARALLELISM_OPTIONS = { + "tp_sp": [True], +} + +# GPU count filter: only evaluate combinations whose TP*CP*DP*LP fall inside +# this inclusive range. +GPU_COUNT_MIN = 64 +GPU_COUNT_MAX = 1024 + +# Select the y-axis metric for the scatter plot and report. Accepted values: +# "runtime" -> plot raw runtime in seconds (lower is better) +# "performance" -> plot 1 / runtime (higher is better) +PLOT_METRIC = "runtime" + +# Random seed for reproducible jitter in the scatter plot. +PLOT_JITTER_SEED = 1234 +# Maximum absolute horizontal jitter (in GPU units). +PLOT_JITTER_WIDTH = 0.175 + +# Default output artefacts +PLOT_OUTPUT_PATH = "tools/parallelism_sweep.png" +PLOT_MFU_OUTPUT_PATH = "tools/parallelism_sweep_mfu.png" +REPORT_OUTPUT_PATH = "tools/parallelism_sweep.tsv" + +# AstraSim cache handling within DeepFlow (mirrors run_perf default options). +ASTRA_CACHE_MODE = "NO_CACHE" # Options: NO_CACHE, CACHE_READONLY, CACHE_READWRITE + +# Maximum number of parallel worker processes (set <= available CPUs - 1). Set to 1 to disable multiprocessing. +MAX_WORKERS = 96 + + +# ----------------------------------------------------------------------------- +# Helper utilities +# ----------------------------------------------------------------------------- + +def read_yaml(path): + with open(path, "r") as handle: + return yaml.safe_load(handle) + + +def determine_model_mode(model_config_path): + model_dict = read_yaml(model_config_path) + model_param = model_dict.get("model_param") or {} + mode = model_param.get("mode") + if not mode: + raise ValueError("model_param.mode must be defined in {}".format(model_config_path)) + return mode + + +def cartesian_product(option_map): + """Yield dictionaries for every combination inside option_map.""" + if not option_map: + yield {} + return + keys = sorted(option_map.keys()) + value_lists = [option_map[key] for key in keys] + for values in itertools.product(*value_lists): + yield dict(zip(keys, values)) + + +def total_gpu_count(parallel_cfg): + total = 1 + for axis in ("tp", "cp", "dp", "lp"): + value = int(parallel_cfg.get(axis, 1) or 1) + total *= max(1, value) + return total + + +def make_temp_hw_config(base_hw_dict, parallel_settings): + """Return parsed HW config for the given parallelism override.""" + updated = copy.deepcopy(base_hw_dict) + parallel_block = updated.setdefault("parallelism", {}) + for key, value in parallel_settings.items(): + parallel_block[key] = value + + tmp_file = tempfile.NamedTemporaryFile("w", suffix=".yaml", delete=False) + try: + yaml.safe_dump(updated, tmp_file, default_flow_style=False, sort_keys=False) + tmp_file.flush() + tmp_file.close() + return config.parse_config(tmp_file.name, config_type="hardware") + finally: + try: + tmp_file.close() + except Exception: + pass + try: + os.unlink(tmp_file.name) + except Exception: + pass + + +def gpu_peak_flops(hw_config) -> float: + """Return the theoretical peak FLOPs/s for a single GPU.""" + core = hw_config.tech_config.core + bundles = core.num_bundles or 1 + per_cycle_flops = float(core.nominal_flop_rate_per_mcu) + mcu_per_bundle = float(core.num_mcu_per_bundle) + frequency = float(core.operating_frequency or core.nominal_frequency or 0.0) + util = float(getattr(core, "util", 1.0) or 1.0) + peak = per_cycle_flops * mcu_per_bundle * float(bundles) + if frequency > 0: + peak *= frequency + return peak * util if peak > 0 else 0.0 + + +def compute_total_flops(calculator: TimeCalculationLLM) -> float: + """Estimate total FLOPs executed for one batch.""" + batch_size = calculator._effective_transformer_batch() + vocab_size = calculator.vocab_size + hidden_dim = calculator.hidden_dim + seq_len = calculator.seq_len + num_heads = calculator.num_heads + kv_heads = calculator.kv_heads + intermediate_size = calculator.intermediate_size + num_SMs = calculator.hw_config.tech_config.core.num_bundles or 1 + + transformer_timings, _ = calculator.compute_all_gemm_and_node_times( + batch_size, + vocab_size, + hidden_dim, + seq_len, + num_heads, + kv_heads, + intermediate_size, + num_SMs, + ) + gemm_shapes = process_gemm_shapes( + calculator, + batch_size, + seq_len, + hidden_dim, + num_heads, + kv_heads, + intermediate_size, + vocab_size, + ) + + def _gemm_flops(shape) -> float: + if shape is None: + return 0.0 + try: + dims = list(shape) + except TypeError: + return 0.0 + if len(dims) == 4: + b, m, k, n = dims + return 2.0 * float(b) * float(m) * float(k) * float(n) + if len(dims) == 3: + m, k, n = dims + return 2.0 * float(m) * float(k) * float(n) + return 0.0 + + fallback_forward = { + "qkv_proj": _gemm_flops(gemm_shapes.get("qkv_proj")), + "attention_score": _gemm_flops(gemm_shapes.get("attention_score")), + "attention_output": _gemm_flops(gemm_shapes.get("attention_output")), + "output_proj": _gemm_flops(gemm_shapes.get("output_proj")), + "ffn1": _gemm_flops(gemm_shapes.get("ffn1")), + "ffn2": _gemm_flops(gemm_shapes.get("ffn2")), + "linear_softmax": _gemm_flops(gemm_shapes.get("linear")), + } + fallback_forward["attention"] = ( + fallback_forward["attention_score"] + fallback_forward["attention_output"] + ) + + def _op_flops(name: str, fallback: float = 0.0, include_backward: bool = True) -> float: + timing = transformer_timings.get(name) + forward = 0.0 + if timing is not None and timing.forward is not None: + forward = float(timing.forward.flops or 0.0) + if forward <= 0 and fallback > 0: + forward = fallback + backward = 0.0 + if include_backward and timing is not None and timing.backward is not None: + backward = float(timing.backward.flops or 0.0) + if include_backward and backward <= 0 and forward > 0: + backward = 2.0 * forward + return forward + backward + + per_layer_ops = ("qkv_proj", "attention", "output_proj", "ffn1", "ffn2") + per_layer_flops = sum( + _op_flops(name, fallback_forward.get(name, 0.0)) for name in per_layer_ops + ) + + total_flops = per_layer_flops * calculator.num_layers + + total_flops += _op_flops("linear_softmax", fallback=fallback_forward.get("linear_softmax", 0.0), include_backward=True) + total_flops += _op_flops("embedding", fallback=0.0, include_backward=True) + + return float(total_flops) + + +def evaluate_parallelism(hw_dict, model_config_obj, mode, parallel_settings): + hw_config = make_temp_hw_config(hw_dict, parallel_settings) + temp_dir = tempfile.mkdtemp(prefix="parallelism_sweep_") + try: + calculator = TimeCalculationLLM(hw_config, model_config_obj, mode, output_dir=temp_dir) + with open(os.devnull, "w") as devnull, redirect_stdout(devnull), redirect_stderr(devnull): + runtime = calculator.calc_time_llm() + total_flops = compute_total_flops(calculator) + performance = (1.0 / runtime) if runtime and runtime > 0.0 else float("nan") + peak_flops = gpu_peak_flops(hw_config) + num_gpus = total_gpu_count(parallel_settings) + achieved_flops = (total_flops / runtime) if runtime > 0 else float("nan") + denom = peak_flops * num_gpus if peak_flops and num_gpus else float("nan") + mfu = (achieved_flops / denom) if denom and denom > 0 else float("nan") + return { + "runtime": runtime, + "performance": performance, + "total_flops": total_flops, + "peak_flops": peak_flops, + "mfu": mfu, + "achieved_flops": achieved_flops, + "memory_exceeded": getattr(calculator, "memory_capacity_exceeded", False), + "memory_violation_gb": getattr(calculator, "memory_capacity_violation_gb", 0.0), + } + finally: + shutil.rmtree(temp_dir, ignore_errors=True) + + +def set_astrasim_cache_mode(mode_str): + mapping = { + "NO_CACHE": "NO_CACHE", + "CACHE_READONLY": "CACHE_READONLY", + "CACHE_READWRITE": "CACHE_READWRITE", + } + env_value = mapping.get(str(mode_str).strip().upper(), "NO_CACHE") + os.environ["DEEPFLOW_ASTRA_CACHE_MODE"] = env_value + + +def write_report(results, path): + header = [ + "num_gpus", + "runtime_s", + "performance_1_over_s", + "total_flops", + "achieved_flops_per_s", + "peak_flops_per_gpu", + "mfu", + "memory_exceeded", + "memory_violation_gb", + "parallelism", + ] + with open(path, "w") as handle: + handle.write("\t".join(header) + "\n") + for entry in results: + row = [ + str(entry["num_gpus"]), + "{:.6f}".format(entry["runtime"]), + ( + "{:.6f}".format(entry["performance"]) + if entry["performance"] == entry["performance"] + else "nan" + ), + "{:.6e}".format(entry["total_flops"]), + "{:.6e}".format(entry["achieved_flops"]) if entry["achieved_flops"] == entry["achieved_flops"] else "nan", + "{:.6e}".format(entry["peak_flops"]), + "{:.6f}".format(entry["mfu"]) if entry["mfu"] == entry["mfu"] else "nan", + str(entry["memory_exceeded"]), + "{:.6f}".format(entry["memory_violation_gb"]), + repr(entry["parallelism"]), + ] + handle.write("\t".join(row) + "\n") + + +def jitter_positions(gpu_counts, jitter_width): + jittered = [] + for count in gpu_counts: + offset = random.uniform(1 - jitter_width, 1 + jitter_width) + jittered.append(max(count * offset, 1e-3)) + return jittered + + +def plot_results(results, output_path): + if not results: + print("No successful configurations to plot.", file=sys.stderr) + return + + random.seed(PLOT_JITTER_SEED) + xs = jitter_positions([item["num_gpus"] for item in results], PLOT_JITTER_WIDTH) + metric_key = "runtime" if PLOT_METRIC.lower() == "runtime" else "performance" + ys = [item[metric_key] for item in results] + + plt.figure(figsize=(10, 6)) + plt.scatter(xs, ys, s=60, alpha=0.7, edgecolors="none") + + best = min(results, key=lambda item: item["runtime"]) + best_x = best["num_gpus"] + best_y = best[metric_key] + plt.scatter([best_x], [best_y], s=180, marker="*", c="red", label="Best runtime") + + plt.xlabel("Number of GPUs") + if metric_key == "runtime": + plt.ylabel("Runtime (s)") + else: + plt.ylabel("Performance (1 / s)") + plt.xscale("log") + if metric_key == "runtime": + plt.yscale("log") + plt.title("Parallelism sweep") + plt.grid(alpha=0.3) + xticks = sorted(set(item["num_gpus"] for item in results)) + plt.xticks(xticks, [str(int(x)) for x in xticks]) + plt.legend(loc="best") + plt.tight_layout() + plt.savefig(output_path, dpi=200) + plt.close() + print("Saved scatter plot to {}".format(output_path)) + + +def plot_mfu(results, output_path): + if not results: + return + random.seed(PLOT_JITTER_SEED) + xs = jitter_positions([item["num_gpus"] for item in results], PLOT_JITTER_WIDTH) + ys = [item["mfu"] for item in results] + plt.figure(figsize=(10, 6)) + plt.scatter(xs, ys, s=60, alpha=0.7, edgecolors="none") + valid = [item for item in results if item["mfu"] == item["mfu"]] + if valid: + best = max(valid, key=lambda item: item["mfu"]) + plt.scatter([best["num_gpus"]], [best["mfu"]], s=180, marker="*", c="green", label="Max MFU") + plt.xlabel("Number of GPUs") + plt.ylabel("MFU") + plt.xscale("log") + plt.ylim(0.0, 1.05) + plt.title("Parallelism sweep - MFU") + plt.grid(alpha=0.3) + xticks = sorted(set(item["num_gpus"] for item in results)) + plt.xticks(xticks, [str(int(x)) for x in xticks]) + if valid: + plt.legend(loc="best") + plt.tight_layout() + plt.savefig(output_path, dpi=200) + plt.close() + print("Saved MFU scatter plot to {}".format(output_path)) + + +_GLOBAL_MODEL_CONFIG = None +_GLOBAL_MODE = None +_GLOBAL_HW_DICT = None + + +def _worker_init(hw_dict, model_config_path, mode): + global _GLOBAL_MODEL_CONFIG, _GLOBAL_MODE, _GLOBAL_HW_DICT + _GLOBAL_HW_DICT = hw_dict + _GLOBAL_MODE = mode + _GLOBAL_MODEL_CONFIG = config.parse_config(model_config_path, config_type=mode) + + +def _worker_task(parallel_items: Tuple[Tuple[str, object], ...]): + parallel_settings = {k: v for k, v in parallel_items} + try: + metrics = evaluate_parallelism( + _GLOBAL_HW_DICT, + _GLOBAL_MODEL_CONFIG, + _GLOBAL_MODE, + parallel_settings, + ) + return { + "status": "ok", + "parallelism": parallel_settings, + "metrics": metrics, + } + except Exception as exc: + return { + "status": "error", + "parallelism": parallel_settings, + "error": str(exc), + } + + +def _build_tasks( + gpu_choices: Iterable[Dict[str, int]], + other_choices: Iterable[Dict[str, object]], +) -> List[Tuple[Tuple[str, object], ...]]: + tasks: List[Tuple[Tuple[str, object], ...]] = [] + for gpu_choice in gpu_choices: + for other_choice in other_choices: + settings: Dict[str, object] = {} + settings.update(gpu_choice) + settings.update(other_choice) + settings["mb"] = settings.get("lp", 1) + tasks.append(tuple(sorted(settings.items()))) + return tasks + + +def main(): + set_astrasim_cache_mode(ASTRA_CACHE_MODE) + + base_hw_dict = read_yaml(HARDWARE_CONFIG_PATH) + mode = determine_model_mode(MODEL_CONFIG_PATH) + + gpu_axes = list(PARALLELISM_SWEEP.keys()) + other_axes = list(OTHER_PARALLELISM_OPTIONS.keys()) + gpu_combos = list(cartesian_product(PARALLELISM_SWEEP)) + other_combos = list(cartesian_product(OTHER_PARALLELISM_OPTIONS)) + task_items = _build_tasks(gpu_combos, other_combos) + print("Enumerating {} parallelism combinations: {}".format(len(task_items), ", ".join(gpu_axes))) + + results = [] + skipped_out_of_range = 0 + skipped_errors = 0 + error_messages: List[str] = [] + evaluated = 0 + + filtered_tasks: List[Tuple[Tuple[str, object], ...]] = [] + for items in task_items: + settings = dict(items) + num_gpus = total_gpu_count(settings) + if GPU_COUNT_MIN <= num_gpus <= GPU_COUNT_MAX: + filtered_tasks.append(items) + else: + skipped_out_of_range += 1 + + if not filtered_tasks: + print("No configurations within GPU count bounds.") + return + + available_cpus = max(1, os.cpu_count() or 1) + if MAX_WORKERS is None or MAX_WORKERS <= 0: + worker_limit = max(1, available_cpus - 1) + else: + worker_limit = min(MAX_WORKERS, max(1, available_cpus - 1)) + worker_count = max(1, worker_limit) + print(f"Using {worker_count} worker process(es) (out of {available_cpus} CPUs).") + + if worker_count > 1 and len(filtered_tasks) > 1: + with ProcessPoolExecutor( + max_workers=worker_count, + initializer=_worker_init, + initargs=(base_hw_dict, MODEL_CONFIG_PATH, mode), + ) as executor: + futures = {executor.submit(_worker_task, items): items for items in filtered_tasks} + with tqdm(total=len(filtered_tasks), desc="Evaluating", unit="config") as progress: + for future in as_completed(futures): + progress.update(1) + result = future.result() + settings = result["parallelism"] + num_gpus = total_gpu_count(settings) + if result["status"] != "ok": + skipped_errors += 1 + msg = result.get("error") or "unknown error" + error_messages.append(f"{settings}: {msg}") + continue + metrics = result["metrics"] + evaluated += 1 + entry = { + "parallelism": settings, + "num_gpus": num_gpus, + "runtime": metrics["runtime"], + "performance": metrics["performance"], + "total_flops": metrics["total_flops"], + "achieved_flops": metrics["achieved_flops"], + "peak_flops": metrics["peak_flops"], + "mfu": metrics["mfu"], + "memory_exceeded": metrics["memory_exceeded"], + "memory_violation_gb": metrics["memory_violation_gb"], + } + results.append(entry) + else: + model_config_obj = config.parse_config(MODEL_CONFIG_PATH, config_type=mode) + with tqdm(total=len(filtered_tasks), desc="Evaluating", unit="config") as progress: + for items in filtered_tasks: + settings = dict(items) + num_gpus = total_gpu_count(settings) + progress.update(1) + try: + metrics = evaluate_parallelism(base_hw_dict, model_config_obj, mode, settings) + except Exception as exc: + skipped_errors += 1 + error_messages.append(f"{settings}: {exc}") + continue + + evaluated += 1 + entry = { + "parallelism": settings, + "num_gpus": num_gpus, + "runtime": metrics["runtime"], + "performance": metrics["performance"], + "total_flops": metrics["total_flops"], + "achieved_flops": metrics["achieved_flops"], + "peak_flops": metrics["peak_flops"], + "mfu": metrics["mfu"], + "memory_exceeded": metrics["memory_exceeded"], + "memory_violation_gb": metrics["memory_violation_gb"], + } + results.append(entry) + + total_skipped = skipped_out_of_range + skipped_errors + if not results: + print( + "No valid configurations evaluated ({} skipped: {} out-of-range, {} errors).".format( + total_skipped, skipped_out_of_range, skipped_errors + ) + ) + if error_messages: + print("Encountered errors for configurations:") + for msg in error_messages: + print(f" {msg}") + return + + best = min(results, key=lambda item: item["runtime"]) + print( + f"Evaluated {evaluated} configuration(s); skipped {total_skipped} " + f"(out_of_range={skipped_out_of_range}, errors={skipped_errors})." + ) + if error_messages: + print("Some configurations failed:") + for msg in error_messages: + print(f" {msg}") + print("\nBest configuration (lowest runtime):") + print(" Parallelism: {}".format(best["parallelism"])) + print(" GPUs: {}".format(best["num_gpus"])) + print(" Runtime: {:.4f} s".format(best["runtime"])) + print(" Performance (1/s): {:.4f}".format(best["performance"])) + print(" Total FLOPs: {:.3e}".format(best["total_flops"])) + print(" MFU: {:.3f}".format(best["mfu"])) + if best["memory_exceeded"]: + print(" Memory capacity exceeded by {:.3f} GB".format(best["memory_violation_gb"])) + + try: + write_report(results, REPORT_OUTPUT_PATH) + print("Wrote detailed report to {}".format(REPORT_OUTPUT_PATH)) + except Exception as exc: + print("Warning: failed to write report: {}".format(exc), file=sys.stderr) + + plot_results(results, PLOT_OUTPUT_PATH) + plot_mfu(results, PLOT_MFU_OUTPUT_PATH) + + +if __name__ == "__main__": + main() diff --git a/uv.lock b/uv.lock index ffb4f894..01fc09c5 100644 --- a/uv.lock +++ b/uv.lock @@ -135,6 +135,7 @@ dependencies = [ { name = "pyyaml" }, { name = "seaborn" }, { name = "six" }, + { name = "tqdm" }, { name = "tzdata" }, ] @@ -152,6 +153,7 @@ requires-dist = [ { name = "pyyaml", specifier = "==6.0.2" }, { name = "seaborn", specifier = ">=0.13.2" }, { name = "six", specifier = "==1.17.0" }, + { name = "tqdm", specifier = ">=4.67.1" }, { name = "tzdata", specifier = "==2025.2" }, ] @@ -685,7 +687,6 @@ wheels = [ { url = "https://files.pythonhosted.org/packages/fa/de/02b54f42487e3d3c6efb3f89428677074ca7bf43aae402517bc7cca949f3/PyYAML-6.0.2-cp313-cp313-win_amd64.whl", hash = "sha256:8388ee1976c416731879ac16da0aff3f63b286ffdd57cdeb95f3f2e085687563", size = 156446 }, ] - [[package]] name = "seaborn" version = "0.13.2" @@ -709,6 +710,18 @@ wheels = [ { url = "https://files.pythonhosted.org/packages/b7/ce/149a00dd41f10bc29e5921b496af8b574d8413afcd5e30dfa0ed46c2cc5e/six-1.17.0-py2.py3-none-any.whl", hash = "sha256:4721f391ed90541fddacab5acf947aa0d3dc7d27b2e1e8eda2be8970586c3274", size = 11050 }, ] +[[package]] +name = "tqdm" +version = "4.67.1" +source = { registry = "https://pypi.org/simple" } +dependencies = [ + { name = "colorama", marker = "sys_platform == 'win32'" }, +] +sdist = { url = "https://files.pythonhosted.org/packages/a8/4b/29b4ef32e036bb34e4ab51796dd745cdba7ed47ad142a9f4a1eb8e0c744d/tqdm-4.67.1.tar.gz", hash = "sha256:f8aef9c52c08c13a65f30ea34f4e5aac3fd1a34959879d7e59e63027286627f2", size = 169737 } +wheels = [ + { url = "https://files.pythonhosted.org/packages/d0/30/dc54f88dd4a2b5dc8a0279bdd7270e735851848b762aeb1c1184ed1f6b14/tqdm-4.67.1-py3-none-any.whl", hash = "sha256:26445eca388f82e72884e0d580d5464cd801a3ea01e63e5601bdff9ba6a48de2", size = 78540 }, +] + [[package]] name = "tzdata" version = "2025.2" From 62aa31c79347ef29962b9217815f47a955fbd34b Mon Sep 17 00:00:00 2001 From: NanoMembers Date: Tue, 4 Nov 2025 15:57:41 -0800 Subject: [PATCH 16/77] memory estimation enabled and flash attention supported --- LLM_util.py | 15 ++++++++++----- time_calculation_LLM.py | 4 +++- 2 files changed, 13 insertions(+), 6 deletions(-) diff --git a/LLM_util.py b/LLM_util.py index b8dd07ee..bfdd1143 100644 --- a/LLM_util.py +++ b/LLM_util.py @@ -122,19 +122,24 @@ def process_gemm_shapes(self, batch_size, seq_len, d_model, num_heads, kv_heads, return processed -def get_transformer_mem_layer( dp, tp, batch_size, hidden_dim, seq_len, intermediate_size, n_heads, precision, zero_stage, model_type="gpt"):#https://www.determined.ai/blog/act-mem-1. https://arxiv.org/pdf/2205.05198. https://shjwudp.github.io/blog/2023/gpt-training-memory-estimation-nemo-training-practice/ +def get_transformer_mem_layer( dp, tp, batch_size, hidden_dim, seq_len, intermediate_size, n_heads, precision, zero_stage, flash_attention, model_type="gpt"):# https://arxiv.org/pdf/2205.05198. https://shjwudp.github.io/blog/2023/gpt-training-memory-estimation-nemo-training-practice/ """ memory estimation of transformer layer for single gpu case in inference mode is supported. other modes or layers are work in progress.""" #Activations refer to output activations that need to be stored - act_memory_layer = seq_len * batch_size * hidden_dim * (34 / tp + 5 * n_heads * seq_len/(hidden_dim * tp) ) * (precision.activations / 2) #from https://arxiv.org/pdf/2205.05198 + + if flash_attention: + act_memory_layer = seq_len * batch_size * hidden_dim * (34 / tp ) * (precision.activations / 2) #flash attention activation saved + else: + act_memory_layer = seq_len * batch_size * hidden_dim * (34 / tp + 5 * n_heads * seq_len/(hidden_dim * tp) ) * (precision.activations / 2) + + act_memory_layer_inf = seq_len * batch_size * intermediate_size / tp * precision.activations #inference max activation memory, no need to store for backpropagation ffn_proj_factor = 3 if str(model_type).lower() == "llama" else 2 transformer_param_layer = 4* hidden_dim * hidden_dim + intermediate_size * ffn_proj_factor * hidden_dim # weights Wq,Wk,Wv,Wo,ffn - optimizer_mem = 10 * transformer_param_layer / dp # zero1 style optimizer memory - #TODO: which optimizer mem to use? - optimizer_mem = 10 * transformer_param_layer * (precision.optimizer_states / 2) / tp # don't divide by dp for DDP + + optimizer_mem = (precision.optimizer_states * 2 + precision.activations) * transformer_param_layer / tp # don't divide by dp for DDP, NeMo ZeRO-1 optimizer style tensor_weight_memory_layer = transformer_param_layer * precision.parameters / tp #weight memory # master_parameters is set to 0 by default, so this works. master_weight_memory_layer = transformer_param_layer * precision.master_parameters / tp diff --git a/time_calculation_LLM.py b/time_calculation_LLM.py index c152fffe..21b2a74e 100644 --- a/time_calculation_LLM.py +++ b/time_calculation_LLM.py @@ -2345,6 +2345,7 @@ def calc_time_llm(self): precision=self.precision, model_type=self.model_type, zero_stage=self.zero_stage, + flash_attention=self.flash_attention, ) ) @@ -2437,7 +2438,8 @@ def calc_time_llm(self): self.pipeline_graph = dispatcher.pipeline_graph self.pipeline_root = pipeline_root self.pipeline_interconnect = dispatcher.interconnect_params - + _, peak_mem = self._simulate_with_memory(graph_root, memory_data, mode="training") + # debug helper. If set, print analytical transformer time and actual transformer time From 03b505efa1867c7372db1cf217eb1652ee0fcec1 Mon Sep 17 00:00:00 2001 From: gkarfakis19 Date: Tue, 4 Nov 2025 16:07:14 -0800 Subject: [PATCH 17/77] sweep updates --- astra-sim | 2 +- tools/parallelism_sweep.py | 704 ++++++++++++++++++++++++++++--------- 2 files changed, 548 insertions(+), 158 deletions(-) diff --git a/astra-sim b/astra-sim index 6a22754b..3f56a2f0 160000 --- a/astra-sim +++ b/astra-sim @@ -1 +1 @@ -Subproject commit 6a22754be1578ccc7116848f84df5d62898f48ce +Subproject commit 3f56a2f016cc9f8718789e9d2a401fb84dc3a847 diff --git a/tools/parallelism_sweep.py b/tools/parallelism_sweep.py index 69b09d4c..efae0f62 100644 --- a/tools/parallelism_sweep.py +++ b/tools/parallelism_sweep.py @@ -12,6 +12,8 @@ from __future__ import print_function +import argparse +import ast import copy import itertools import math @@ -33,6 +35,12 @@ from time_calculation_LLM import TimeCalculationLLM from LLM_util import process_gemm_shapes +import seaborn as sns +import numpy as np +import pandas as pd +import math +import numpy as np +from matplotlib.patches import Polygon # ----------------------------------------------------------------------------- # Global configuration knobs (no CLI) @@ -88,6 +96,68 @@ # Helper utilities # ----------------------------------------------------------------------------- +def add_rgb_ternary_legend(ax, + corner_labels=("R = log2(tp+cp)", + "G = log2(lp)", + "B = log2(dp)"), + gamma=0.85, + inset_xywh=(0.72, 0.58, 0.24, 0.24), # Axes fraction: x, y, w, h + n=120): + """ + Draw a tiny triangular legend that shows how the RGB mix works. + - Uses barycentric blending inside an equilateral triangle. + - 'gamma' should match your color gamma used in the scatter. + - inset_xywh is (x0,y0,w,h) in axes-relative coords. + """ + # Inset axis + iax = ax.inset_axes(inset_xywh, transform=ax.transAxes) + iax.set_aspect("equal") + iax.set_axis_off() + + # Equilateral triangle vertices (R, G, B corners) + A = np.array([0.0, 0.0]) # R corner + B = np.array([1.0, 0.0]) # G corner + C = np.array([0.5, math.sqrt(3)/2.0]) # B corner + + # Sample interior with barycentric coords r,g,b (r+g+b=1) + xs, ys, cols = [], [], [] + for i in range(n + 1): + for j in range(n + 1 - i): + r = i / n + g = j / n + b = 1.0 - r - g + p = r * A + g * B + b * C + # apply same gamma tweak used in your plot colors + rr = r ** gamma + gg = g ** gamma + bb = b ** gamma + xs.append(p[0]); ys.append(p[1]); cols.append((rr, gg, bb)) + + # iax.scatter(xs, ys, s=2, c=cols, edgecolors="none") + iax.add_patch(Polygon([A, B, C], facecolor="none", edgecolor="black", linewidth=1)) + + # Corner labels + iax.text(A[0]-0.06, A[1]-0.04, corner_labels[0], ha="right", va="top", fontsize=8) + iax.text(B[0]+0.06, B[1]-0.04, corner_labels[1], ha="left", va="top", fontsize=8) + iax.text(C[0], C[1]+0.06, corner_labels[2], ha="center", va="bottom", fontsize=8) + + # Optional tick marks (25/50/75%) along edges (comment out if you want it cleaner) + for t in (0.25, 0.5, 0.75): + # Edge AB (vary r vs g, b=0) + p = (1-t) * A + t * B + iax.plot([p[0]], [p[1]], marker="|", color="black", ms=6) + # Edge BC (vary g vs b, r=0) + p = (1-t) * B + t * C + iax.plot([p[0]], [p[1]], marker="_", color="black", ms=6) + # Edge CA (vary b vs r, g=0) + p = (1-t) * C + t * A + iax.plot([p[0]], [p[1]], marker="|", color="black", ms=6) + + # Tiny caption (how channels were normalized) + iax.text(0.5, -0.18, "Each channel min–max normalized\nthen gamma-adjusted", + ha="center", va="top", fontsize=7, transform=iax.transAxes) + + def read_yaml(path): with open(path, "r") as handle: return yaml.safe_load(handle) @@ -144,6 +214,38 @@ def make_temp_hw_config(base_hw_dict, parallel_settings): except Exception: pass +def _rgb_from_parallelism(entries, gamma: float = 0.85): + """ + Build per-point RGBA colors using: + R = log2(tp+cp), G = log2(lp), B = log2(dp) + Each channel is min-max normalized over the dataset, then gamma-adjusted. + """ + tps = np.array([max(1, int(e["parallelism"].get("tp", 1))) for e in entries], dtype=float) + cps = np.array([max(1, int(e["parallelism"].get("cp", 1))) for e in entries], dtype=float) + dps = np.array([max(1, int(e["parallelism"].get("dp", 1))) for e in entries], dtype=float) + lps = np.array([max(1, int(e["parallelism"].get("lp", 1))) for e in entries], dtype=float) + + r_raw = np.log2(tps + cps) + g_raw = np.log2(lps) + b_raw = np.log2(dps) + + def _norm(x): + xmin, xmax = float(np.min(x)), float(np.max(x)) + if not np.isfinite(xmin) or not np.isfinite(xmax) or xmax <= xmin: + return np.zeros_like(x) + y = (x - xmin) / (xmax - xmin) + if gamma and gamma != 1.0: + y = np.power(y, gamma) + return y + + r = _norm(r_raw) + g = _norm(g_raw) + b = _norm(b_raw) + + # RGBA with slight transparency to reduce overdraw + return np.stack([r, g, b, np.full_like(r, 0.9)], axis=1) + + def gpu_peak_flops(hw_config) -> float: """Return the theoretical peak FLOPs/s for a single GPU.""" @@ -319,6 +421,61 @@ def write_report(results, path): handle.write("\t".join(row) + "\n") +def _parse_float(value: str) -> float: + try: + return float(value) + except (TypeError, ValueError): + return float("nan") + + +def load_results_from_report(path: str) -> List[Dict[str, object]]: + if not os.path.exists(path): + raise FileNotFoundError(f"No report found at {path}") + + results: List[Dict[str, object]] = [] + with open(path, "r") as handle: + header = handle.readline() + if not header: + return results + for line in handle: + line = line.strip() + if not line: + continue + parts = line.split("\t") + if len(parts) < 10: + print(f"Warning: skipping malformed row in report: {line}", file=sys.stderr) + continue + try: + parallelism = ast.literal_eval(parts[9]) + except (SyntaxError, ValueError) as exc: + print(f"Warning: failed to parse parallelism '{parts[9]}': {exc}", file=sys.stderr) + continue + entry = { + "num_gpus": int(parts[0]), + "runtime": _parse_float(parts[1]), + "performance": _parse_float(parts[2]), + "total_flops": _parse_float(parts[3]), + "achieved_flops": _parse_float(parts[4]), + "peak_flops": _parse_float(parts[5]), + "mfu": _parse_float(parts[6]), + "memory_exceeded": str(parts[7]).strip().lower() == "true", + "memory_violation_gb": _parse_float(parts[8]), + "parallelism": parallelism, + } + results.append(entry) + return results + + +def parse_args(): + parser = argparse.ArgumentParser(description="DeepFlow parallelism sweep utility.") + parser.add_argument( + "--plot_only", + action="store_true", + help="Skip evaluation and only regenerate plots from the existing TSV report.", + ) + return parser.parse_args() + + def jitter_positions(gpu_counts, jitter_width): jittered = [] for count in gpu_counts: @@ -327,69 +484,274 @@ def jitter_positions(gpu_counts, jitter_width): return jittered +# def plot_results_legacy(results, output_path): +# if not results: +# print("No successful configurations to plot.", file=sys.stderr) +# return + +# random.seed(PLOT_JITTER_SEED) +# xs = jitter_positions([item["num_gpus"] for item in results], PLOT_JITTER_WIDTH) +# metric_key = "runtime" if PLOT_METRIC.lower() == "runtime" else "performance" +# ys = [item[metric_key] for item in results] + +# plt.figure(figsize=(10, 6)) +# plt.scatter(xs, ys, s=60, alpha=0.7, edgecolors="none") + +# best = min(results, key=lambda item: item["runtime"]) +# best_x = best["num_gpus"] +# best_y = best[metric_key] +# plt.scatter([best_x], [best_y], s=180, marker="*", c="red", label="Best runtime") + +# plt.xlabel("Number of GPUs") +# if metric_key == "runtime": +# plt.ylabel("Runtime (s)") +# else: +# plt.ylabel("Performance (1 / s)") +# plt.xscale("log") +# if metric_key == "runtime": +# plt.yscale("log") +# plt.title("Parallelism sweep") +# plt.grid(alpha=0.3) +# xticks = sorted(set(item["num_gpus"] for item in results)) +# plt.xticks(xticks, [str(int(x)) for x in xticks]) +# plt.legend(loc="best") +# plt.tight_layout() +# plt.savefig(output_path, dpi=200) +# plt.close() +# print("Saved scatter plot to {}".format(output_path)) + +def _gpu_exp(n: int) -> float: + """Return log2(num_gpus). Assumes powers of two; will return floats otherwise.""" + return math.log2(float(n)) + + def plot_results(results, output_path): if not results: print("No successful configurations to plot.", file=sys.stderr) return - random.seed(PLOT_JITTER_SEED) - xs = jitter_positions([item["num_gpus"] for item in results], PLOT_JITTER_WIDTH) metric_key = "runtime" if PLOT_METRIC.lower() == "runtime" else "performance" - ys = [item[metric_key] for item in results] - plt.figure(figsize=(10, 6)) - plt.scatter(xs, ys, s=60, alpha=0.7, edgecolors="none") + # Build tidy frame (include GPU exponent) + rows = [] + for i, item in enumerate(results): + p = item["parallelism"] + ng = int(item["num_gpus"]) + rows.append({ + "row_id": i, + "num_gpus": ng, + "gpu_exp": _gpu_exp(ng), + "tp": int(p.get("tp", 1)), + "cp": int(p.get("cp", 1)), + "dp": int(p.get("dp", 1)), + "lp": int(p.get("lp", 1)), + metric_key: float(item[metric_key]), + }) + df = pd.DataFrame(rows) + + # Order categories by exponent (gives even spacing like log2) + order = sorted(df["gpu_exp"].unique()) + df["gpu_exp_cat"] = pd.Categorical(df["gpu_exp"], categories=order, ordered=True) + + # --- Per-point RGB colors: R=log2(tp+cp), G=log2(lp), B=log2(dp) --- + tps = df["tp"].to_numpy(dtype=float) + cps = df["cp"].to_numpy(dtype=float) + dps = df["dp"].to_numpy(dtype=float) + lps = df["lp"].to_numpy(dtype=float) + + r_raw = np.log2(tps + cps) + g_raw = np.log2(lps) + b_raw = np.log2(dps) + + def _norm(x, gamma=0.85): + x = x.astype(float) + xmin, xmax = float(np.min(x)), float(np.max(x)) + if not np.isfinite(xmin) or not np.isfinite(xmax) or xmax <= xmin: + y = np.zeros_like(x) + else: + y = (x - xmin) / (xmax - xmin) + return np.power(y, gamma) if gamma and gamma != 1.0 else y - best = min(results, key=lambda item: item["runtime"]) - best_x = best["num_gpus"] - best_y = best[metric_key] - plt.scatter([best_x], [best_y], s=180, marker="*", c="red", label="Best runtime") + r = _norm(r_raw) + g = _norm(g_raw) + b = _norm(b_raw) + a = np.full_like(r, 0.9) - plt.xlabel("Number of GPUs") + # Palette: unique color per row via hue="row_id" + palette = {rid: (float(r[i]), float(g[i]), float(b[i]), float(a[i])) + for i, rid in enumerate(df["row_id"])} + + plt.figure(figsize=(10, 6)) + ax = sns.swarmplot( + data=df, + x="gpu_exp_cat", + y=metric_key, + hue="row_id", + palette=palette, # per-point RGBA + size=3.75, + linewidth=0, + dodge=False, + legend=False + ) + # add_rgb_ternary_legend( + # ax, + # corner_labels=("R = log2(tp+cp)", "G = log2(lp)", "B = log2(dp)"), + # gamma=0.85, # match your _rgb_from_parallelism gamma + # inset_xywh=(0.70, 0.50, 0.26, 0.26) # tweak position/size to taste + # ) + + + # Best runtime star — place at the matching exponent category index + best = min(results, key=lambda it: it["runtime"]) + best_exp = _gpu_exp(int(best["num_gpus"])) + best_idx = order.index(best_exp) + best_y = float(best[metric_key]) + ax.scatter([best_idx], [best_y], s=180, marker="*", c="red", zorder=5, label="Best runtime") + + # Labels & scales + ax.set_xlabel("Number of GPUs (log2 categories)") if metric_key == "runtime": - plt.ylabel("Runtime (s)") + ax.set_ylabel("Runtime (s)") + ax.set_yscale("log") else: - plt.ylabel("Performance (1 / s)") - plt.xscale("log") - if metric_key == "runtime": - plt.yscale("log") - plt.title("Parallelism sweep") - plt.grid(alpha=0.3) - xticks = sorted(set(item["num_gpus"] for item in results)) - plt.xticks(xticks, [str(int(x)) for x in xticks]) - plt.legend(loc="best") - plt.tight_layout() - plt.savefig(output_path, dpi=200) - plt.close() - print("Saved scatter plot to {}".format(output_path)) + ax.set_ylabel("Performance (1 / s)") + # Tick labels show the REAL GPU count (2**exp) so axis *looks* log2 + ax.set_xticks(range(len(order))) + ax.set_xticklabels([str(int(round(2 ** e))) if float(e).is_integer() else f"2^{e:.2f}" + for e in order]) + + ax.set_title("Parallelism sweep – beeswarm on log2(GPUs) categories") + ax.grid(alpha=0.3, axis="y") + ax.legend(loc="best") + + # Reminder of color encoding + ax.text(0.99, 0.01, + "Color: R=log2(tp+cp), G=log2(lp), B=log2(dp)", + transform=ax.transAxes, ha="right", va="bottom", fontsize=9, alpha=0.8) -def plot_mfu(results, output_path): - if not results: - return - random.seed(PLOT_JITTER_SEED) - xs = jitter_positions([item["num_gpus"] for item in results], PLOT_JITTER_WIDTH) - ys = [item["mfu"] for item in results] - plt.figure(figsize=(10, 6)) - plt.scatter(xs, ys, s=60, alpha=0.7, edgecolors="none") - valid = [item for item in results if item["mfu"] == item["mfu"]] - if valid: - best = max(valid, key=lambda item: item["mfu"]) - plt.scatter([best["num_gpus"]], [best["mfu"]], s=180, marker="*", c="green", label="Max MFU") - plt.xlabel("Number of GPUs") - plt.ylabel("MFU") - plt.xscale("log") - plt.ylim(0.0, 1.05) - plt.title("Parallelism sweep - MFU") - plt.grid(alpha=0.3) - xticks = sorted(set(item["num_gpus"] for item in results)) - plt.xticks(xticks, [str(int(x)) for x in xticks]) - if valid: - plt.legend(loc="best") plt.tight_layout() plt.savefig(output_path, dpi=200) plt.close() - print("Saved MFU scatter plot to {}".format(output_path)) + print(f"Saved swarm plot to {output_path}") + + +# def plot_results_categorical(results, output_path): +# if not results: +# print("No successful configurations to plot.", file=sys.stderr) +# return + +# metric_key = "runtime" if PLOT_METRIC.lower() == "runtime" else "performance" + +# # Build a tidy DataFrame for seaborn (categorical x keeps each GPU count together) +# df_rows = [] +# for item in results: +# p = item["parallelism"] +# df_rows.append({ +# "num_gpus": int(item["num_gpus"]), +# metric_key: float(item[metric_key]), +# "tp": int(p.get("tp", 1)), +# "cp": int(p.get("cp", 1)), +# "dp": int(p.get("dp", 1)), +# "lp": int(p.get("lp", 1)), +# }) +# df = pd.DataFrame(df_rows) + +# # Order GPU categories left->right by numeric value +# order = sorted(df["num_gpus"].unique()) +# df["num_gpus_cat"] = pd.Categorical(df["num_gpus"], categories=order, ordered=True) + +# # --- Beeswarm placement --- +# plt.figure(figsize=(10, 6)) +# ax = sns.swarmplot( +# data=df, +# x="num_gpus_cat", +# y=metric_key, +# size=5, +# color="k", # temporary (we'll recolor with our RGBs after layout) +# linewidth=0, +# alpha=0.0 # invisible; we only want the positions it computes +# ) + +# # Grab the laid-out positions, remove seaborn's collection, and redraw with our colors +# if not ax.collections: +# print("Warning: swarmplot produced no collections.", file=sys.stderr) +# return +# coll = ax.collections[0] +# offsets = coll.get_offsets() # Nx2 array (x_index, y_value) in data coords (x is categorical index) +# coll.remove() + +# # Compute per-point RGBA colors from parallelism +# colors = _rgb_from_parallelism(results, gamma=0.85) + +# # Plot the recolored points at the computed positions +# ax.scatter( +# offsets[:, 0], offsets[:, 1], +# s=60, +# c=colors, +# edgecolors="none", +# zorder=3 +# ) + +# # Global best marker (use category index for x) +# best = min(results, key=lambda item: item["runtime"]) +# best_x_idx = order.index(int(best["num_gpus"])) # categorical index +# best_y = float(best[metric_key]) +# ax.scatter([best_x_idx], [best_y], s=180, marker="*", c="red", zorder=5, label="Best runtime") + +# # Axes/labels +# ax.set_xlabel("Number of GPUs") +# if metric_key == "runtime": +# ax.set_ylabel("Runtime (s)") +# ax.set_yscale("log") +# else: +# ax.set_ylabel("Performance (1 / s)") + +# # Pretty ticks for categorical axis +# ax.set_xticks(range(len(order))) +# ax.set_xticklabels([str(x) for x in order]) + +# ax.set_title("Parallelism sweep (beeswarm with log-RGB coloring)") +# ax.grid(alpha=0.3, axis="y") +# ax.legend(loc="best") + +# # Small caption to recall color encoding +# ax.text(0.99, 0.01, +# "Color channels: R=log2(tp+cp), G=log2(lp), B=log2(dp)", +# transform=ax.transAxes, ha="right", va="bottom", fontsize=9, alpha=0.8) + +# plt.tight_layout() +# plt.savefig(output_path, dpi=200) +# plt.close() +# print(f"Saved swarm plot to {output_path}") + + +# def plot_mfu(results, output_path): +# if not results: +# return +# random.seed(PLOT_JITTER_SEED) +# xs = jitter_positions([item["num_gpus"] for item in results], PLOT_JITTER_WIDTH) +# ys = [item["mfu"] for item in results] +# plt.figure(figsize=(10, 6)) +# plt.scatter(xs, ys, s=60, alpha=0.7, edgecolors="none") +# valid = [item for item in results if item["mfu"] == item["mfu"]] +# if valid: +# best = max(valid, key=lambda item: item["mfu"]) +# plt.scatter([best["num_gpus"]], [best["mfu"]], s=180, marker="*", c="green", label="Max MFU") +# plt.xlabel("Number of GPUs") +# plt.ylabel("MFU") +# plt.xscale("log") +# plt.ylim(0.0, 1.05) +# plt.title("Parallelism sweep - MFU") +# plt.grid(alpha=0.3) +# xticks = sorted(set(item["num_gpus"] for item in results)) +# plt.xticks(xticks, [str(int(x)) for x in xticks]) +# if valid: +# plt.legend(loc="best") +# plt.tight_layout() +# plt.savefig(output_path, dpi=200) +# plt.close() +# print("Saved MFU scatter plot to {}".format(output_path)) _GLOBAL_MODEL_CONFIG = None @@ -442,64 +804,117 @@ def _build_tasks( def main(): + args = parse_args() set_astrasim_cache_mode(ASTRA_CACHE_MODE) - base_hw_dict = read_yaml(HARDWARE_CONFIG_PATH) - mode = determine_model_mode(MODEL_CONFIG_PATH) - - gpu_axes = list(PARALLELISM_SWEEP.keys()) - other_axes = list(OTHER_PARALLELISM_OPTIONS.keys()) - gpu_combos = list(cartesian_product(PARALLELISM_SWEEP)) - other_combos = list(cartesian_product(OTHER_PARALLELISM_OPTIONS)) - task_items = _build_tasks(gpu_combos, other_combos) - print("Enumerating {} parallelism combinations: {}".format(len(task_items), ", ".join(gpu_axes))) - - results = [] - skipped_out_of_range = 0 - skipped_errors = 0 - error_messages: List[str] = [] - evaluated = 0 - - filtered_tasks: List[Tuple[Tuple[str, object], ...]] = [] - for items in task_items: - settings = dict(items) - num_gpus = total_gpu_count(settings) - if GPU_COUNT_MIN <= num_gpus <= GPU_COUNT_MAX: - filtered_tasks.append(items) - else: - skipped_out_of_range += 1 - - if not filtered_tasks: - print("No configurations within GPU count bounds.") - return + results: List[Dict[str, object]] = [] - available_cpus = max(1, os.cpu_count() or 1) - if MAX_WORKERS is None or MAX_WORKERS <= 0: - worker_limit = max(1, available_cpus - 1) + if args.plot_only: + try: + results = load_results_from_report(REPORT_OUTPUT_PATH) + except FileNotFoundError as exc: + print(str(exc)) + return + + if not results: + print(f"No entries found in {REPORT_OUTPUT_PATH}; nothing to plot.") + return + + best = min(results, key=lambda item: item["runtime"]) + print(f"Loaded {len(results)} configuration(s) from {REPORT_OUTPUT_PATH}") + print("\nBest configuration (lowest runtime):") + print(" Parallelism: {}".format(best["parallelism"])) + print(" GPUs: {}".format(best["num_gpus"])) + print(" Runtime: {:.4f} s".format(best["runtime"])) + print(" Performance (1/s): {:.4f}".format(best["performance"])) + print(" Total FLOPs: {:.3e}".format(best["total_flops"])) + print(" MFU: {:.3f}".format(best["mfu"])) + if best["memory_exceeded"]: + print(" Memory capacity exceeded by {:.3f} GB".format(best["memory_violation_gb"])) else: - worker_limit = min(MAX_WORKERS, max(1, available_cpus - 1)) - worker_count = max(1, worker_limit) - print(f"Using {worker_count} worker process(es) (out of {available_cpus} CPUs).") - - if worker_count > 1 and len(filtered_tasks) > 1: - with ProcessPoolExecutor( - max_workers=worker_count, - initializer=_worker_init, - initargs=(base_hw_dict, MODEL_CONFIG_PATH, mode), - ) as executor: - futures = {executor.submit(_worker_task, items): items for items in filtered_tasks} + base_hw_dict = read_yaml(HARDWARE_CONFIG_PATH) + mode = determine_model_mode(MODEL_CONFIG_PATH) + + gpu_axes = list(PARALLELISM_SWEEP.keys()) + other_axes = list(OTHER_PARALLELISM_OPTIONS.keys()) + gpu_combos = list(cartesian_product(PARALLELISM_SWEEP)) + other_combos = list(cartesian_product(OTHER_PARALLELISM_OPTIONS)) + task_items = _build_tasks(gpu_combos, other_combos) + print("Enumerating {} parallelism combinations: {}".format(len(task_items), ", ".join(gpu_axes))) + + skipped_out_of_range = 0 + skipped_errors = 0 + error_messages: List[str] = [] + evaluated = 0 + + filtered_tasks: List[Tuple[Tuple[str, object], ...]] = [] + for items in task_items: + settings = dict(items) + num_gpus = total_gpu_count(settings) + if GPU_COUNT_MIN <= num_gpus <= GPU_COUNT_MAX: + filtered_tasks.append(items) + else: + skipped_out_of_range += 1 + + if not filtered_tasks: + print("No configurations within GPU count bounds.") + return + + available_cpus = max(1, os.cpu_count() or 1) + if MAX_WORKERS is None or MAX_WORKERS <= 0: + worker_limit = max(1, available_cpus - 1) + else: + worker_limit = min(MAX_WORKERS, max(1, available_cpus - 1)) + worker_count = max(1, worker_limit) + print(f"Using {worker_count} worker process(es) (out of {available_cpus} CPUs).") + + if worker_count > 1 and len(filtered_tasks) > 1: + with ProcessPoolExecutor( + max_workers=worker_count, + initializer=_worker_init, + initargs=(base_hw_dict, MODEL_CONFIG_PATH, mode), + ) as executor: + futures = {executor.submit(_worker_task, items): items for items in filtered_tasks} + with tqdm(total=len(filtered_tasks), desc="Evaluating", unit="config") as progress: + for future in as_completed(futures): + progress.update(1) + result = future.result() + settings = result["parallelism"] + num_gpus = total_gpu_count(settings) + if result["status"] != "ok": + skipped_errors += 1 + msg = result.get("error") or "unknown error" + error_messages.append(f"{settings}: {msg}") + continue + metrics = result["metrics"] + evaluated += 1 + entry = { + "parallelism": settings, + "num_gpus": num_gpus, + "runtime": metrics["runtime"], + "performance": metrics["performance"], + "total_flops": metrics["total_flops"], + "achieved_flops": metrics["achieved_flops"], + "peak_flops": metrics["peak_flops"], + "mfu": metrics["mfu"], + "memory_exceeded": metrics["memory_exceeded"], + "memory_violation_gb": metrics["memory_violation_gb"], + } + results.append(entry) + else: + model_config_obj = config.parse_config(MODEL_CONFIG_PATH, config_type=mode) with tqdm(total=len(filtered_tasks), desc="Evaluating", unit="config") as progress: - for future in as_completed(futures): - progress.update(1) - result = future.result() - settings = result["parallelism"] + for items in filtered_tasks: + settings = dict(items) num_gpus = total_gpu_count(settings) - if result["status"] != "ok": + progress.update(1) + try: + metrics = evaluate_parallelism(base_hw_dict, model_config_obj, mode, settings) + except Exception as exc: skipped_errors += 1 - msg = result.get("error") or "unknown error" - error_messages.append(f"{settings}: {msg}") + error_messages.append(f"{settings}: {exc}") continue - metrics = result["metrics"] + evaluated += 1 entry = { "parallelism": settings, @@ -514,75 +929,50 @@ def main(): "memory_violation_gb": metrics["memory_violation_gb"], } results.append(entry) - else: - model_config_obj = config.parse_config(MODEL_CONFIG_PATH, config_type=mode) - with tqdm(total=len(filtered_tasks), desc="Evaluating", unit="config") as progress: - for items in filtered_tasks: - settings = dict(items) - num_gpus = total_gpu_count(settings) - progress.update(1) - try: - metrics = evaluate_parallelism(base_hw_dict, model_config_obj, mode, settings) - except Exception as exc: - skipped_errors += 1 - error_messages.append(f"{settings}: {exc}") - continue - - evaluated += 1 - entry = { - "parallelism": settings, - "num_gpus": num_gpus, - "runtime": metrics["runtime"], - "performance": metrics["performance"], - "total_flops": metrics["total_flops"], - "achieved_flops": metrics["achieved_flops"], - "peak_flops": metrics["peak_flops"], - "mfu": metrics["mfu"], - "memory_exceeded": metrics["memory_exceeded"], - "memory_violation_gb": metrics["memory_violation_gb"], - } - results.append(entry) - - total_skipped = skipped_out_of_range + skipped_errors - if not results: - print( - "No valid configurations evaluated ({} skipped: {} out-of-range, {} errors).".format( - total_skipped, skipped_out_of_range, skipped_errors + + total_skipped = skipped_out_of_range + skipped_errors + if not results: + print( + "No valid configurations evaluated ({} skipped: {} out-of-range, {} errors).".format( + total_skipped, skipped_out_of_range, skipped_errors + ) ) + if error_messages: + print("Encountered errors for configurations:") + for msg in error_messages: + print(f" {msg}") + return + + best = min(results, key=lambda item: item["runtime"]) + print( + f"Evaluated {evaluated} configuration(s); skipped {total_skipped} " + f"(out_of_range={skipped_out_of_range}, errors={skipped_errors})." ) if error_messages: - print("Encountered errors for configurations:") + print("Some configurations failed:") for msg in error_messages: print(f" {msg}") - return + print("\nBest configuration (lowest runtime):") + print(" Parallelism: {}".format(best["parallelism"])) + print(" GPUs: {}".format(best["num_gpus"])) + print(" Runtime: {:.4f} s".format(best["runtime"])) + print(" Performance (1/s): {:.4f}".format(best["performance"])) + print(" Total FLOPs: {:.3e}".format(best["total_flops"])) + print(" MFU: {:.3f}".format(best["mfu"])) + if best["memory_exceeded"]: + print(" Memory capacity exceeded by {:.3f} GB".format(best["memory_violation_gb"])) - best = min(results, key=lambda item: item["runtime"]) - print( - f"Evaluated {evaluated} configuration(s); skipped {total_skipped} " - f"(out_of_range={skipped_out_of_range}, errors={skipped_errors})." - ) - if error_messages: - print("Some configurations failed:") - for msg in error_messages: - print(f" {msg}") - print("\nBest configuration (lowest runtime):") - print(" Parallelism: {}".format(best["parallelism"])) - print(" GPUs: {}".format(best["num_gpus"])) - print(" Runtime: {:.4f} s".format(best["runtime"])) - print(" Performance (1/s): {:.4f}".format(best["performance"])) - print(" Total FLOPs: {:.3e}".format(best["total_flops"])) - print(" MFU: {:.3f}".format(best["mfu"])) - if best["memory_exceeded"]: - print(" Memory capacity exceeded by {:.3f} GB".format(best["memory_violation_gb"])) + try: + write_report(results, REPORT_OUTPUT_PATH) + print("Wrote detailed report to {}".format(REPORT_OUTPUT_PATH)) + except Exception as exc: + print("Warning: failed to write report: {}".format(exc), file=sys.stderr) - try: - write_report(results, REPORT_OUTPUT_PATH) - print("Wrote detailed report to {}".format(REPORT_OUTPUT_PATH)) - except Exception as exc: - print("Warning: failed to write report: {}".format(exc), file=sys.stderr) + if not results: + return plot_results(results, PLOT_OUTPUT_PATH) - plot_mfu(results, PLOT_MFU_OUTPUT_PATH) + # plot_mfu(results, PLOT_MFU_OUTPUT_PATH) if __name__ == "__main__": From a222bd1217e1c0e623d9edc63886ca7f4122c3fe Mon Sep 17 00:00:00 2001 From: gkarfakis19 Date: Tue, 4 Nov 2025 17:02:11 -0800 Subject: [PATCH 18/77] small fix --- LLM_execution.py | 2 +- tools/parallelism_sweep.py | 84 ++++++++++++++------------------------ 2 files changed, 32 insertions(+), 54 deletions(-) diff --git a/LLM_execution.py b/LLM_execution.py index db64c3de..ce1eeb74 100644 --- a/LLM_execution.py +++ b/LLM_execution.py @@ -679,7 +679,7 @@ def _safe_int(value: Any, default: int = 1) -> int: if enforce_layout and declared > 1 and not first_active_checked: canon = sorted(dim_axes) - if canon != ["cp", "tp"]: + if canon != ["cp", "tp"] and canon != ["tp", "cp"] and (axis_sizes["tp"] > 1 or axis_sizes["cp"] > 1): raise ValueError( "For hierarchical/hybrid AstraSim modes, the first active network " "dimension must contain exactly {'tp','cp'} to represent the tensor/context cluster." diff --git a/tools/parallelism_sweep.py b/tools/parallelism_sweep.py index efae0f62..ac4a5997 100644 --- a/tools/parallelism_sweep.py +++ b/tools/parallelism_sweep.py @@ -53,10 +53,10 @@ # Parallelism values to sweep (dense grid). Edit to suit your search space. # Keys should match the entries under the YAML "parallelism" section. PARALLELISM_SWEEP = { - "tp": [2**i for i in range(0, 9)], - "cp": [2**i for i in range(0, 9)], - "dp": [2**i for i in range(0, 9)], - "lp": [2**i for i in range(0, 9)], + "tp": [2**i for i in range(0, 3)], + "cp": [2**i for i in range(0, 3)], + "dp": [2**i for i in range(0, 3)], + "lp": [2**i for i in range(0, 3)], } # Optional knobs that still live inside the parallelism section but do not @@ -67,8 +67,8 @@ # GPU count filter: only evaluate combinations whose TP*CP*DP*LP fall inside # this inclusive range. -GPU_COUNT_MIN = 64 -GPU_COUNT_MAX = 1024 +GPU_COUNT_MIN = 16 +GPU_COUNT_MAX = 64 # Select the y-axis metric for the scatter plot and report. Accepted values: # "runtime" -> plot raw runtime in seconds (lower is better) @@ -89,7 +89,7 @@ ASTRA_CACHE_MODE = "NO_CACHE" # Options: NO_CACHE, CACHE_READONLY, CACHE_READWRITE # Maximum number of parallel worker processes (set <= available CPUs - 1). Set to 1 to disable multiprocessing. -MAX_WORKERS = 96 +MAX_WORKERS = 16 # ----------------------------------------------------------------------------- @@ -262,29 +262,23 @@ def gpu_peak_flops(hw_config) -> float: def compute_total_flops(calculator: TimeCalculationLLM) -> float: - """Estimate total FLOPs executed for one batch.""" - batch_size = calculator._effective_transformer_batch() + """Estimate total FLOPs executed for one global batch (forward + backward).""" + global_batch = getattr(calculator, "batch_size", None) + if global_batch is None or global_batch <= 0: + global_batch = calculator._effective_transformer_batch() + if not global_batch or global_batch <= 0: + return float("nan") + vocab_size = calculator.vocab_size hidden_dim = calculator.hidden_dim seq_len = calculator.seq_len num_heads = calculator.num_heads kv_heads = calculator.kv_heads intermediate_size = calculator.intermediate_size - num_SMs = calculator.hw_config.tech_config.core.num_bundles or 1 - transformer_timings, _ = calculator.compute_all_gemm_and_node_times( - batch_size, - vocab_size, - hidden_dim, - seq_len, - num_heads, - kv_heads, - intermediate_size, - num_SMs, - ) gemm_shapes = process_gemm_shapes( calculator, - batch_size, + global_batch, seq_len, hidden_dim, num_heads, @@ -308,42 +302,26 @@ def _gemm_flops(shape) -> float: return 2.0 * float(m) * float(k) * float(n) return 0.0 - fallback_forward = { - "qkv_proj": _gemm_flops(gemm_shapes.get("qkv_proj")), - "attention_score": _gemm_flops(gemm_shapes.get("attention_score")), - "attention_output": _gemm_flops(gemm_shapes.get("attention_output")), - "output_proj": _gemm_flops(gemm_shapes.get("output_proj")), - "ffn1": _gemm_flops(gemm_shapes.get("ffn1")), - "ffn2": _gemm_flops(gemm_shapes.get("ffn2")), - "linear_softmax": _gemm_flops(gemm_shapes.get("linear")), - } - fallback_forward["attention"] = ( - fallback_forward["attention_score"] + fallback_forward["attention_output"] - ) - - def _op_flops(name: str, fallback: float = 0.0, include_backward: bool = True) -> float: - timing = transformer_timings.get(name) - forward = 0.0 - if timing is not None and timing.forward is not None: - forward = float(timing.forward.flops or 0.0) - if forward <= 0 and fallback > 0: - forward = fallback - backward = 0.0 - if include_backward and timing is not None and timing.backward is not None: - backward = float(timing.backward.flops or 0.0) - if include_backward and backward <= 0 and forward > 0: - backward = 2.0 * forward + def _forward_backward(shape) -> float: + forward = _gemm_flops(shape) + if forward <= 0.0: + return 0.0 + backward = 2.0 * forward return forward + backward - per_layer_ops = ("qkv_proj", "attention", "output_proj", "ffn1", "ffn2") - per_layer_flops = sum( - _op_flops(name, fallback_forward.get(name, 0.0)) for name in per_layer_ops - ) + per_layer_flops = 0.0 + per_layer_flops += _forward_backward(gemm_shapes.get("qkv_proj")) + per_layer_flops += _forward_backward(gemm_shapes.get("attention_score")) + per_layer_flops += _forward_backward(gemm_shapes.get("attention_output")) + per_layer_flops += _forward_backward(gemm_shapes.get("output_proj")) + per_layer_flops += _forward_backward(gemm_shapes.get("ffn1")) + per_layer_flops += _forward_backward(gemm_shapes.get("ffn2")) + + total_flops = per_layer_flops * float(calculator.num_layers) - total_flops = per_layer_flops * calculator.num_layers + total_flops += _forward_backward(gemm_shapes.get("linear")) - total_flops += _op_flops("linear_softmax", fallback=fallback_forward.get("linear_softmax", 0.0), include_backward=True) - total_flops += _op_flops("embedding", fallback=0.0, include_backward=True) + # Embedding forward/backward is comparatively small; keep placeholder for future refinement. return float(total_flops) From 15433892f63cbc8bde00dd4552a30fa5159282f3 Mon Sep 17 00:00:00 2001 From: gkarfakis19 Date: Tue, 4 Nov 2025 17:30:56 -0800 Subject: [PATCH 19/77] synthetic duplication fix --- astrasim_lib/config_generation.py | 59 ++++++++++++++++++++++++------- astrasim_lib/executor.py | 4 +++ time_calculation.py | 2 +- 3 files changed, 51 insertions(+), 14 deletions(-) diff --git a/astrasim_lib/config_generation.py b/astrasim_lib/config_generation.py index 78d9cb3d..03ef3ede 100644 --- a/astrasim_lib/config_generation.py +++ b/astrasim_lib/config_generation.py @@ -4,6 +4,7 @@ import os from typing import Any, Dict, List, Optional, Sequence, Tuple +from types import SimpleNamespace from hw_component import Network @@ -133,21 +134,53 @@ def _safe_int(value: Any, default: int = 1) -> int: candidate = default return max(1, candidate) - axis_sizes: Dict[str, int] = { + axes_filter_original = tuple(axes_filter) if axes_filter else None + axes_filter_normalized: Optional[Tuple[str, ...]] = None + if axes_filter_original: + axes_filter_normalized = tuple(str(axis).strip().lower() for axis in axes_filter_original) + + axis_sizes_full: Dict[str, int] = { "tp": _safe_int(getattr(sch_config, "tp", 1)) if sch_config else 1, "cp": _safe_int(getattr(sch_config, "cp", 1)) if sch_config else 1, "lp": _safe_int(getattr(sch_config, "lp", 1)) if sch_config else 1, "dp": _safe_int(getattr(sch_config, "dp", 1)) if sch_config else 1, } - if axes_filter: - axis_sizes = {axis: axis_sizes.get(axis, 1) for axis in axes_filter} - axis_sizes.setdefault("tp", 1) - axis_sizes.setdefault("cp", 1) - axis_sizes.setdefault("lp", 1) - axis_sizes.setdefault("dp", 1) - axis_order_preference = ["tp", "cp", "lp", "dp"] - - allowed_axes = set(axis_sizes.keys()) if axes_filter else None + synthetic_only = axes_filter_normalized is not None and set(axes_filter_normalized) == {"synthetic2"} + if synthetic_only: + if dimensions: + base_dim = dimensions[0] + base_bw = float(getattr(base_dim, "effective_bandwidth", getattr(base_dim, "bandwidth", None))) + base_latency = float(getattr(base_dim, "latency", None)) + base_topology = getattr(base_dim, "topology_type", None) + base_collectives = getattr(base_dim, "collective_override", {}) or {} + else: + raise ValueError(f"Synthetic dimension requested but no dimensions found in hardware config. Info: Called with filter {axes_filter_original}, npus_count={npus_count}, axes_sizes={axis_sizes_full}.") + synthetic_dim = SimpleNamespace( + label="synthetic2", + parallelisms=("synthetic2",), + size=2, + topology_type=base_topology, + effective_bandwidth=base_bw, + bandwidth=base_bw, + latency=base_latency, + collective_override=base_collectives, + faulty_links=(), + ) + dimensions = [synthetic_dim] + axis_sizes_full = {"synthetic2": 2} + axis_sizes = {"synthetic2": 2} + axis_order_preference = ["synthetic2"] + else: + axis_sizes: Dict[str, int] = dict(axis_sizes_full) + if axes_filter_normalized: + axis_sizes = {axis: axis_sizes_full.get(axis, 1) for axis in axes_filter_normalized} + axis_sizes.setdefault("tp", 1) + axis_sizes.setdefault("cp", 1) + axis_sizes.setdefault("lp", 1) + axis_sizes.setdefault("dp", 1) + axis_order_preference = ["tp", "cp", "lp", "dp"] + + allowed_axes = set(axis_sizes.keys()) if axes_filter_normalized else None # print(f"Allowed axes: {allowed_axes}") # print(f"Axes sizes: {axis_sizes}") dim_infos: List[Tuple[Any, List[str], int, int]] = [] @@ -155,7 +188,7 @@ def _safe_int(value: Any, default: int = 1) -> int: for dim_idx, dim in enumerate(dimensions): axes = [str(axis).strip().lower() for axis in getattr(dim, "parallelisms", ())] # print(f"All axes for dimension {dim.label}: {axes}") - if allowed_axes is not None and axes_filter: + if allowed_axes is not None and axes_filter_normalized: filtered_axes = [axis for axis in axes if axis in allowed_axes] axes = filtered_axes # print(f"Filtered axes for dimension {dim.label}: {axes}") @@ -215,7 +248,7 @@ def _safe_int(value: Any, default: int = 1) -> int: axes_needed: List[str] = [] remaining = target - + for axis in axis_order_preference: size = axis_sizes.get(axis, 1) @@ -228,7 +261,7 @@ def _safe_int(value: Any, default: int = 1) -> int: if remaining != 1: raise ValueError( - f"Unable to map requested npus_count={target} to network axes {axis_sizes}." + f"Unable to map requested npus_count={target} to network axes {axis_sizes}. Info: Called with filter {axes_filter_original}, npus_count={npus_count}, axes_sizes={axis_sizes}." ) axes_needed_set = set(axes_needed) diff --git a/astrasim_lib/executor.py b/astrasim_lib/executor.py index 80e92ba8..8dc8c84a 100644 --- a/astrasim_lib/executor.py +++ b/astrasim_lib/executor.py @@ -1526,12 +1526,14 @@ def run_astra_simulation_only_onepath( rank_count = len(rank_ids) # Astrasim doesn't play well with only 1 rank. # When that happens, let's duplicate to 2 ranks. No collectives exist between the two so this should not have an effect. + synthetic_pair = False if rank_count == 1: # duplicate the .et file src = os.path.join(work_dir, "llm_graph.0.et") dst = os.path.join(work_dir, "llm_graph.1.et") shutil.copy(src, dst) rank_ids = [rank_ids[0], rank_ids[0]+1] + synthetic_pair = True rank_count = len(rank_ids) # Handle artifact visualization if enabled @@ -1558,6 +1560,8 @@ def run_astra_simulation_only_onepath( rank_layout = getattr(fwdbwd_root, "_astrasim_rank_layout", None) axis_order, axis_sizes = _extract_axis_layout(rank_layout) axes_filter = _derive_axes_filter(axis_order, axis_sizes, dp_count) + if synthetic_pair: + axes_filter = ["synthetic2"] astra_configs = generate_astrasim_configs_from_hw( time_calc_obj.hw_config, work_dir, diff --git a/time_calculation.py b/time_calculation.py index cac9dbf0..5b8c6247 100644 --- a/time_calculation.py +++ b/time_calculation.py @@ -77,7 +77,7 @@ def collective( if kind in collective_ops or kind == "pipeline": # Pipeline uses 2 NPUs for point-to-point, others use part npus = 2 if kind == "pipeline" else part - axis_filter = None if kind == "pipeline" else axis + axis_filter = axis if kind != "pipeline" else None network_time = self._astra_collective(kind, npus, network_bytes, axis_filter) else: raise ValueError(f"Unsupported collective operation: {kind}") From 13fc1749b69b83c398a858bcb53f09b8057966fe Mon Sep 17 00:00:00 2001 From: gkarfakis19 Date: Tue, 4 Nov 2025 17:38:48 -0800 Subject: [PATCH 20/77] mem awareness --- tools/parallelism_sweep.py | 14 ++++++++++++-- 1 file changed, 12 insertions(+), 2 deletions(-) diff --git a/tools/parallelism_sweep.py b/tools/parallelism_sweep.py index ac4a5997..29afeb84 100644 --- a/tools/parallelism_sweep.py +++ b/tools/parallelism_sweep.py @@ -88,6 +88,9 @@ # AstraSim cache handling within DeepFlow (mirrors run_perf default options). ASTRA_CACHE_MODE = "NO_CACHE" # Options: NO_CACHE, CACHE_READONLY, CACHE_READWRITE +# Plotting behaviour toggles +MEM_AWARE_FILTER = True # When True, skip memory-violating configurations in plots. + # Maximum number of parallel worker processes (set <= available CPUs - 1). Set to 1 to disable multiprocessing. MAX_WORKERS = 16 @@ -508,11 +511,18 @@ def plot_results(results, output_path): print("No successful configurations to plot.", file=sys.stderr) return + plot_entries = results + if MEM_AWARE_FILTER: + plot_entries = [item for item in results if not item.get("memory_exceeded", False)] + if not plot_entries: + print("Warning: all configurations violate memory limits; skipping plot.", file=sys.stderr) + return + metric_key = "runtime" if PLOT_METRIC.lower() == "runtime" else "performance" # Build tidy frame (include GPU exponent) rows = [] - for i, item in enumerate(results): + for i, item in enumerate(plot_entries): p = item["parallelism"] ng = int(item["num_gpus"]) rows.append({ @@ -580,7 +590,7 @@ def _norm(x, gamma=0.85): # Best runtime star — place at the matching exponent category index - best = min(results, key=lambda it: it["runtime"]) + best = min(plot_entries, key=lambda it: it["runtime"]) best_exp = _gpu_exp(int(best["num_gpus"])) best_idx = order.index(best_exp) best_y = float(best[metric_key]) From fae70d5bae83580b7cf13122702944ffc2ff0bf1 Mon Sep 17 00:00:00 2001 From: NanoMembers Date: Wed, 5 Nov 2025 18:48:41 -0800 Subject: [PATCH 21/77] modified memory estimation to pipedream style and microbatch 1 assumption --- LLM_util.py | 8 +++++--- time_calculation_LLM.py | 2 ++ 2 files changed, 7 insertions(+), 3 deletions(-) diff --git a/LLM_util.py b/LLM_util.py index bfdd1143..626ed480 100644 --- a/LLM_util.py +++ b/LLM_util.py @@ -122,17 +122,19 @@ def process_gemm_shapes(self, batch_size, seq_len, d_model, num_heads, kv_heads, return processed -def get_transformer_mem_layer( dp, tp, batch_size, hidden_dim, seq_len, intermediate_size, n_heads, precision, zero_stage, flash_attention, model_type="gpt"):# https://arxiv.org/pdf/2205.05198. https://shjwudp.github.io/blog/2023/gpt-training-memory-estimation-nemo-training-practice/ +def get_transformer_mem_layer( dp, tp, lp, mb, batch_size, hidden_dim, seq_len, intermediate_size, n_heads, precision, zero_stage, flash_attention, model_type="gpt"):# https://arxiv.org/pdf/2205.05198. https://shjwudp.github.io/blog/2023/gpt-training-memory-estimation-nemo-training-practice/ """ memory estimation of transformer layer for single gpu case in inference mode is supported. other modes or layers are work in progress.""" #Activations refer to output activations that need to be stored if flash_attention: - act_memory_layer = seq_len * batch_size * hidden_dim * (34 / tp ) * (precision.activations / 2) #flash attention activation saved + # act_memory_layer = seq_len * batch_size * hidden_dim * (34 / tp ) * (precision.activations / 2) #flash attention activation saved full recompute + # act_memory_layer = seq_len * batch_size * hidden_dim * 2 * (precision.activations / 2) + act_memory_layer = seq_len * 1 * hidden_dim * (34 / tp ) * (precision.activations / 2) #assuming micro batch size 1 and selective recompute else: act_memory_layer = seq_len * batch_size * hidden_dim * (34 / tp + 5 * n_heads * seq_len/(hidden_dim * tp) ) * (precision.activations / 2) - + act_memory_layer *= lp / mb #pipedream style pipelining activation memory act_memory_layer_inf = seq_len * batch_size * intermediate_size / tp * precision.activations #inference max activation memory, no need to store for backpropagation ffn_proj_factor = 3 if str(model_type).lower() == "llama" else 2 diff --git a/time_calculation_LLM.py b/time_calculation_LLM.py index 21b2a74e..1f1315fb 100644 --- a/time_calculation_LLM.py +++ b/time_calculation_LLM.py @@ -2337,6 +2337,8 @@ def calc_time_llm(self): LLM_util.get_transformer_mem_layer( dp=self.dp, tp=self.tp, + lp=self.lp, + mb=self.mb, batch_size=batch_size, hidden_dim=hidden_dim, seq_len=seq_len / self.cp, From 71736323d21c496076848e733ad07f9ab3f9ccac Mon Sep 17 00:00:00 2001 From: NanoMembers Date: Wed, 5 Nov 2025 19:51:57 -0800 Subject: [PATCH 22/77] enabled memory estimation with gradient accumulation considered enabled memory estimation with gradient accumulation considered enabled memory estimation with gradient accumulation considered --- LLM_util.py | 16 ++++++++--- config.py | 32 +++++++++++++++++++-- configs/model-config/LLM.yaml | 3 +- configs/model-config/LLM_inf.yaml | 3 +- configs/model-config/LLM_problem.yaml | 3 +- configs/model-config/Llama2-7B.yaml | 3 +- configs/model-config/Llama2-7B_inf.yaml | 3 +- configs/model-config/Llama3.1-405B.yaml | 3 +- configs/model-config/Llama3.1-405B_inf.yaml | 3 +- configs/model-config/Llama3.1-70B_inf.yaml | 3 +- configs/model-config/Llama3.1-8B_inf.yaml | 3 +- configs/model-config/adding_new_models.md | 4 +-- configs/model-config/hf_to_config.py | 11 +++++-- model.py | 3 +- time_calculation.py | 9 +++++- time_calculation_LLM.py | 4 +-- 16 files changed, 82 insertions(+), 24 deletions(-) diff --git a/LLM_util.py b/LLM_util.py index 626ed480..6b10f6b2 100644 --- a/LLM_util.py +++ b/LLM_util.py @@ -128,9 +128,7 @@ def get_transformer_mem_layer( dp, tp, lp, mb, batch_size, hidden_dim, seq_len, #Activations refer to output activations that need to be stored if flash_attention: - # act_memory_layer = seq_len * batch_size * hidden_dim * (34 / tp ) * (precision.activations / 2) #flash attention activation saved full recompute - # act_memory_layer = seq_len * batch_size * hidden_dim * 2 * (precision.activations / 2) - act_memory_layer = seq_len * 1 * hidden_dim * (34 / tp ) * (precision.activations / 2) #assuming micro batch size 1 and selective recompute + act_memory_layer = seq_len * batch_size * hidden_dim * (34 / tp ) * (precision.activations / 2) #assuming selective recompute else: act_memory_layer = seq_len * batch_size * hidden_dim * (34 / tp + 5 * n_heads * seq_len/(hidden_dim * tp) ) * (precision.activations / 2) @@ -207,7 +205,17 @@ def get_embedding_weight_mem( def get_tot_mem_req(exp_hw_config, exp_model_config, **kwargs): # Model Params - batch_size = int(kwargs.get('batch_size', exp_model_config.model_config.batch_size)) + default_global_batch = getattr( + exp_model_config.model_config, + "global_batch_size", + getattr(exp_model_config.model_config, "batch_size"), + ) + batch_size = int( + kwargs.get( + "global_batch_size", + kwargs.get("batch_size", default_global_batch), + ) + ) hidden_dim = int(kwargs.get('hidden_dim', exp_model_config.model_config.hidden_dim)) vocab_size = int(kwargs.get('vocab_size', exp_model_config.model_config.vocab_size)) n_layers = int(kwargs.get('num_layer', exp_model_config.model_config.num_layers)) diff --git a/config.py b/config.py index 00a0743e..85560466 100644 --- a/config.py +++ b/config.py @@ -766,7 +766,8 @@ class LLMConfig: tied_embeddings: bool num_layers: int hidden_dim: int - batch_size: int + global_batch_size: int + gradient_accumulation_steps: int seq_len: int decode_len: Optional[int] intermediate_size: Optional[int] @@ -788,6 +789,12 @@ def use_flashattention(self) -> bool: @property def use_moe(self) -> bool: return self.num_experts > 1 + + + @property + def grad_accumulation_steps(self) -> int: + """Backward-compatible alias for gradient accumulation steps.""" + return self.gradient_accumulation_steps LLMInferenceConfig = _namedtuple( "inference_param", [ @@ -1331,7 +1338,25 @@ def _pop_required_int(field: str) -> int: num_layers = _pop_required_int("num_layers") hidden_dim = _pop_required_int("hidden_dim") - batch_size = _pop_required_int("batch_size") + + global_batch_size = _pop_required_int("global_batch_size") + + + grad_accum_field = mp.pop("gradient_accumulation_steps", None) + if grad_accum_field is None: + grad_accum_field = mp.pop("gradient_accumulation_step", None) + if grad_accum_field is None: + gradient_accumulation_steps = 1 + else: + try: + gradient_accumulation_steps = int(grad_accum_field) + except (TypeError, ValueError) as exc: + raise ValueError( + f"model_param.gradient_accumulation_steps must be an integer (got {grad_accum_field!r})" + ) from exc + if gradient_accumulation_steps <= 0: + raise ValueError("model_param.gradient_accumulation_steps must be a positive integer") + seq_len = _pop_required_int("seq_len") vocab_size = _pop_required_int("vocab_size") @@ -1352,7 +1377,8 @@ def _pop_required_int(field: str) -> int: tied_embeddings=tied_embeddings, num_layers=num_layers, hidden_dim=hidden_dim, - batch_size=batch_size, + global_batch_size=global_batch_size, + gradient_accumulation_steps=gradient_accumulation_steps, seq_len=seq_len, decode_len=decode_len, intermediate_size=intermediate_size, diff --git a/configs/model-config/LLM.yaml b/configs/model-config/LLM.yaml index 22072e7c..73ea9b04 100644 --- a/configs/model-config/LLM.yaml +++ b/configs/model-config/LLM.yaml @@ -3,7 +3,8 @@ model_param: run_type: "training" # Type of run: "training" or "inference" tied_embeddings: false # Share input and output embedding weights (WIP: affects memory estimation only) model_type: "gpt" # gpt assumes dense-style architecture with GELU; llama assumes llama-style architecture with SwiGLU - batch_size: 1024 # Sequences per global batch + global_batch_size: 1024 # Sequences per global batch + gradient_accumulation_steps: 1 # Number of gradient accumulation micro-steps per optimizer update seq_len: 16384 # Sequence length (tokens per sequence) hidden_dim: 8192 # Transformer hidden dimension attention: # Attention-related parameters diff --git a/configs/model-config/LLM_inf.yaml b/configs/model-config/LLM_inf.yaml index b16307bb..34e92bcc 100644 --- a/configs/model-config/LLM_inf.yaml +++ b/configs/model-config/LLM_inf.yaml @@ -3,7 +3,8 @@ model_param: run_type: "inference" # Type of run: "training" or "inference" tied_embeddings: true # Share input and output embedding weights (WIP: affects memory estimation only) model_type: "gpt" # gpt assumes dense-style architecture with GELU; llama assumes llama-style architecture with SwiGLU - batch_size: 2048 # Sequences per global batch + global_batch_size: 2048 # Sequences per global batch + gradient_accumulation_steps: 1 # Number of gradient accumulation micro-steps per optimizer update seq_len: 2048 # Total context length (prefill + decode_len for inference) decode_len: 1024 # Number of autoregressive decode steps. hidden_dim: 8192 # Transformer hidden dimension diff --git a/configs/model-config/LLM_problem.yaml b/configs/model-config/LLM_problem.yaml index 9ab1e9c6..146d22e8 100644 --- a/configs/model-config/LLM_problem.yaml +++ b/configs/model-config/LLM_problem.yaml @@ -3,7 +3,8 @@ model_param: run_type: "training" # Type of run: "training" or "inference" tied_embeddings: false # Share input and output embedding weights (WIP: affects memory estimation only) model_type: "gpt" # gpt assumes dense-style architecture with GELU; llama assumes llama-style architecture with SwiGLU - batch_size: 1024 # Sequences per global batch + global_batch_size: 1024 # Sequences per global batch + gradient_accumulation_steps: 1 # Number of gradient accumulation micro-steps per optimizer update seq_len: 16384 # Sequence length (tokens per sequence) hidden_dim: 8192 # Transformer hidden dimension attention: # Attention-related parameters diff --git a/configs/model-config/Llama2-7B.yaml b/configs/model-config/Llama2-7B.yaml index bf26bb5b..d862e9e4 100644 --- a/configs/model-config/Llama2-7B.yaml +++ b/configs/model-config/Llama2-7B.yaml @@ -3,7 +3,8 @@ model_param: run_type: "training" # Type of run: "training" or "inference" tied_embeddings: true # Share input and output embedding weights (WIP: affects memory estimation only) model_type: "llama" # gpt assumes dense-style architecture with GELU; llama assumes llama-style architecture with SwiGLU - batch_size: 1024 # Sequences per global batch + global_batch_size: 1024 # Sequences per global batch + gradient_accumulation_steps: 1 # Number of gradient accumulation micro-steps per optimizer update seq_len: 2048 # Sequence length (tokens per sequence) hidden_dim: 4096 # Transformer hidden dimension attention: # Attention-related parameters diff --git a/configs/model-config/Llama2-7B_inf.yaml b/configs/model-config/Llama2-7B_inf.yaml index 092477c1..0bed4a63 100644 --- a/configs/model-config/Llama2-7B_inf.yaml +++ b/configs/model-config/Llama2-7B_inf.yaml @@ -3,7 +3,8 @@ model_param: run_type: "inference" # Type of run: "training" or "inference" tied_embeddings: true # Share input and output embedding weights (WIP: affects memory estimation only) model_type: "llama" # gpt assumes dense-style architecture with GELU; llama assumes llama-style architecture with SwiGLU - batch_size: 16 # Sequences per global batch + global_batch_size: 16 # Sequences per global batch + gradient_accumulation_steps: 1 # Number of gradient accumulation micro-steps per optimizer update seq_len: 2048 # Total context length (prefill + decode_len for inference) decode_len: 1024 # Number of tokens to decode during inference hidden_dim: 4096 # Transformer hidden dimension diff --git a/configs/model-config/Llama3.1-405B.yaml b/configs/model-config/Llama3.1-405B.yaml index 3b277e7c..6894dea3 100644 --- a/configs/model-config/Llama3.1-405B.yaml +++ b/configs/model-config/Llama3.1-405B.yaml @@ -3,7 +3,8 @@ model_param: run_type: "training" # Type of run: "training" or "inference" tied_embeddings: true # Share input and output embedding weights (WIP: affects memory estimation only) model_type: "llama" # gpt assumes dense-style architecture with GELU; llama assumes llama-style architecture with SwiGLU - batch_size: 1024 # Sequences per global batch + global_batch_size: 1024 # Global batch size + gradient_accumulation_steps: 1 # Gradient accumulation steps (only for training) seq_len: 2048 # Sequence length (tokens per sequence) hidden_dim: 16384 # Transformer hidden dimension attention: # Attention-related parameters diff --git a/configs/model-config/Llama3.1-405B_inf.yaml b/configs/model-config/Llama3.1-405B_inf.yaml index 57094ac2..cbad8a6e 100644 --- a/configs/model-config/Llama3.1-405B_inf.yaml +++ b/configs/model-config/Llama3.1-405B_inf.yaml @@ -3,7 +3,8 @@ model_param: run_type: "inference" # Type of run: "training" or "inference" tied_embeddings: true # Share input and output embedding weights (WIP: affects memory estimation only) model_type: "llama" # gpt assumes dense-style architecture with GELU; llama assumes llama-style architecture with SwiGLU - batch_size: 1 # Sequences per global batch + global_batch_size: 1 # Sequences per global batch + gradient_accumulation_steps: 1 # Number of gradient accumulation micro-steps per optimizer update seq_len: 2048 # Total context length (prefill + decode_len for inference) decode_len: 1024 # Number of tokens to decode during inference hidden_dim: 16384 # Transformer hidden dimension diff --git a/configs/model-config/Llama3.1-70B_inf.yaml b/configs/model-config/Llama3.1-70B_inf.yaml index 4f133ac5..4bee4285 100644 --- a/configs/model-config/Llama3.1-70B_inf.yaml +++ b/configs/model-config/Llama3.1-70B_inf.yaml @@ -3,7 +3,8 @@ model_param: run_type: "inference" # Type of run: "training" or "inference" tied_embeddings: true # Share input and output embedding weights (WIP: affects memory estimation only) model_type: "llama" # gpt assumes dense-style architecture with GELU; llama assumes llama-style architecture with SwiGLU - batch_size: 1 # Sequences per global batch + global_batch_size: 1 # Sequences per global batch + gradient_accumulation_steps: 1 # Number of gradient accumulation micro-steps per optimizer update seq_len: 2048 # Total context length (prefill + decode_len for inference) decode_len: 1024 # Number of tokens to decode during inference hidden_dim: 8192 # Transformer hidden dimension diff --git a/configs/model-config/Llama3.1-8B_inf.yaml b/configs/model-config/Llama3.1-8B_inf.yaml index d69e5f7a..439048b9 100644 --- a/configs/model-config/Llama3.1-8B_inf.yaml +++ b/configs/model-config/Llama3.1-8B_inf.yaml @@ -3,7 +3,8 @@ model_param: run_type: "inference" # Type of run: "training" or "inference" tied_embeddings: true # Share input and output embedding weights (WIP: affects memory estimation only) model_type: "llama" # gpt assumes dense-style architecture with GELU; llama assumes llama-style architecture with SwiGLU - batch_size: 1 # Sequences per global batch + global_batch_size: 1 # Sequences per global batch + gradient_accumulation_steps: 1 # Number of gradient accumulation micro-steps per optimizer update seq_len: 2048 # Total context length (prefill + decode_len for inference) decode_len: 1024 # Number of tokens to decode during inference hidden_dim: 4096 # Transformer hidden dimension diff --git a/configs/model-config/adding_new_models.md b/configs/model-config/adding_new_models.md index 07373397..3305cb67 100644 --- a/configs/model-config/adding_new_models.md +++ b/configs/model-config/adding_new_models.md @@ -2,10 +2,10 @@ The hf_to_config.py utility can be used to convert a Hugging Face model config.j Usage: ``` -python hf_to_config.py --revision --run-type --batch-size --seq-len --decode-len --use-flashattention --flash-tile-size -o +python hf_to_config.py --revision --run-type --global-batch-size --gradient-accumulation-steps --seq-len --decode-len --use-flashattention --flash-tile-size -o ``` Example: ``` -python hf_to_config.py NousResearch/Hermes-3-Llama-3.1-405B --run-type inference --batch-size 32 --seq-len 65536 --decode-len 1024 --use-flashattention true --flash-tile-size 256 -o Hermes_Llama3.1-405B.yaml +python hf_to_config.py NousResearch/Hermes-3-Llama-3.1-405B --run-type inference --global-batch-size 32 --gradient-accumulation-steps 2 --seq-len 65536 --decode-len 1024 --use-flashattention true --flash-tile-size 256 -o Hermes_Llama3.1-405B.yaml ``` \ No newline at end of file diff --git a/configs/model-config/hf_to_config.py b/configs/model-config/hf_to_config.py index 6a2ca66c..8efe902e 100644 --- a/configs/model-config/hf_to_config.py +++ b/configs/model-config/hf_to_config.py @@ -206,7 +206,8 @@ def _build_yaml_config(cfg: dict, args: argparse.Namespace, model_type: str) -> "run_type": args.run_type, "tied_embeddings": tied_embeddings, "model_type": model_type, - "batch_size": args.batch_size, + "global_batch_size": args.global_batch_size, + "gradient_accumulation_steps": args.gradient_accumulation_steps, "seq_len": int(seq_len), "decode_len": args.decode_len, "hidden_dim": int(hidden_dim), @@ -226,7 +227,13 @@ def parse_args(argv: Optional[List[str]]) -> argparse.Namespace: parser = argparse.ArgumentParser(description=__doc__) parser.add_argument("model_id", help="Hugging Face model identifier, e.g. meta-llama/Llama-2-7b-hf") parser.add_argument("--revision", default="main", help="Model revision/branch to use (default: main)") - parser.add_argument("--batch-size", type=int, default=1, help="Training batch size to encode in the config (default: 1)") + parser.add_argument("--global-batch-size", type=int, default=1, help="Training global batch size to encode in the config (default: 1)") + parser.add_argument( + "--gradient-accumulation-steps", + type=int, + default=1, + help="Number of gradient accumulation micro-steps per optimizer update (default: 1)", + ) parser.add_argument("--seq-len", type=int, default=None, help="Override sequence length. Defaults to HF config max position embeddings if not provided") parser.add_argument("--decode-len", type=int, default=0, help="Decode sequence length (default: 0)") parser.add_argument("--run-type", default="training", help="Config run_type field ('training' or 'inference')") diff --git a/model.py b/model.py index f1678270..c395e755 100644 --- a/model.py +++ b/model.py @@ -20,7 +20,8 @@ def __init__(self, exp_config): class Model_LLM: def __init__(self, exp_config): - self.batch_size = exp_config.model_config.batch_size + self.global_batch_size = exp_config.model_config.global_batch_size + self.gradient_accumulation_steps = exp_config.model_config.gradient_accumulation_steps self.vocab_size = exp_config.model_config.vocab_size self.num_layers = exp_config.model_config.num_layers self.hidden_dim = exp_config.model_config.hidden_dim diff --git a/time_calculation.py b/time_calculation.py index 5b8c6247..b108b5ab 100644 --- a/time_calculation.py +++ b/time_calculation.py @@ -329,7 +329,14 @@ def __init__(self, hw_config, model_config, mode, *, astra_policy_override: Opti self.N = self.model.N if mode == "LLM": - self.batch_size = self.model.batch_size + self.global_batch_size = self.model.global_batch_size + self.gradient_accumulation_steps = self.model.gradient_accumulation_steps + + if self.global_batch_size % self.gradient_accumulation_steps != 0: + raise ValueError( + "Global batch size must be divisible by gradient accumulation steps" + ) + self.batch_size = self.global_batch_size // self.gradient_accumulation_steps self.vocab_size = self.model.vocab_size self.num_layers = self.model.num_layers self.hidden_dim = self.model.hidden_dim diff --git a/time_calculation_LLM.py b/time_calculation_LLM.py index 1f1315fb..f3b58f3a 100644 --- a/time_calculation_LLM.py +++ b/time_calculation_LLM.py @@ -2452,9 +2452,9 @@ def calc_time_llm(self): print(f"Actual transformer forward time: {self.transformer_astrasim_time_forward:.4f}s") print(f"Actual transformer backward time: {self.transformer_astrasim_time_backward:.4f}s") - self.tot_time = time_fw_bw + self.tot_time = time_fw_bw * self.gradient_accumulation_steps - return time_fw_bw + return self.tot_time def _simulate_with_memory( self, From 864d1d522b1fcc449aceabcd42fe11b8b5045712 Mon Sep 17 00:00:00 2001 From: NanoMembers Date: Thu, 6 Nov 2025 15:39:52 -0800 Subject: [PATCH 23/77] bug fix for hybrid mode --- LLM_execution.py | 10 +++++----- 1 file changed, 5 insertions(+), 5 deletions(-) diff --git a/LLM_execution.py b/LLM_execution.py index ce1eeb74..cc6bcb66 100644 --- a/LLM_execution.py +++ b/LLM_execution.py @@ -774,11 +774,11 @@ def _run_pipeline_with_analytical_comm(self, declared_mode: ExecutionMode) -> Ex timed_root = self.pipeline_root else: # must be "ANALYTICAL" filename = "/analytical_graph" - timed_root = self.pipeline_graph.convert_comm_sizes_to_times( - self.pipeline_root, - self.time_calc.network_model, - self.interconnect_params, - ) + timed_root = self.pipeline_graph.convert_comm_sizes_to_times( + self.pipeline_root, + self.time_calc.network_model, + self.interconnect_params, + ) generate_graphs = _env_flag("DEEPFLOW_VISUALIZE_GRAPHS") if generate_graphs: From d733edb4c7fa25d974495e6bec0048cd1a54a413 Mon Sep 17 00:00:00 2001 From: gkarfakis19 Date: Thu, 6 Nov 2025 15:40:00 -0800 Subject: [PATCH 24/77] added many tools + step 1 of 'proper' faulty link support. --- .gitignore | 7 +- HIER_BUG.md | 25 + LLM_execution.py | 106 +++ astrasim_lib/config_generation.py | 64 +- astrasim_lib/executor.py | 28 +- astrasim_lib/fault_projection.py | 159 ++++ astrasim_lib/integration.py | 2 + astrasim_lib/layout_utils.py | 154 ++++ config.py | 123 +-- time_calculation.py | 3 + tools/exec_with_fault.py | 212 +++++ tools/fault_meta_sweep.py | 189 +++++ tools/fault_sweep.py | 1209 +++++++++++++++++++++++++++++ tools/parallelism_sweep.py | 79 +- 14 files changed, 2234 insertions(+), 126 deletions(-) create mode 100644 HIER_BUG.md create mode 100644 astrasim_lib/fault_projection.py create mode 100644 astrasim_lib/layout_utils.py create mode 100644 tools/exec_with_fault.py create mode 100644 tools/fault_meta_sweep.py create mode 100644 tools/fault_sweep.py diff --git a/.gitignore b/.gitignore index 82e51f56..d4467d02 100644 --- a/.gitignore +++ b/.gitignore @@ -6,4 +6,9 @@ venv/ astra_cache/ log/*.log tmp/ -output*/ \ No newline at end of file +output*/ +dbg_hw_conf/ +tools/fault_sweep_results* +tools/*.png +tools/*.tsv +tools/*.json \ No newline at end of file diff --git a/HIER_BUG.md b/HIER_BUG.md new file mode 100644 index 00000000..013ac907 --- /dev/null +++ b/HIER_BUG.md @@ -0,0 +1,25 @@ +# Hierarchical Mode Fault Bugs + +## 1. Fault Projection Leak +Hierarchical execution runs AstraSim separately for the transformer (TP/CP axes) and pipeline (LP axis) graphs. Each run filters the hardware layout down to the axes needed for that graph, but `generate_astrasim_configs_from_hw()` still injects the original `network.faulty_links` list without remapping the endpoints into the reduced coordinate space. Any fault whose flattened IDs fall inside the smaller rank range now appears to live on *every* simulated axis subset. This reinterprets a single physical faulty link multiple times (e.g., once in the transformer run and again in the LP run), so hierarchical mode applies the same degradation repeatedly and slows far more than flattened execution. + +## 2. Single-Replica Transformer Timing +The hierarchical pipeline measures AstraSim on exactly one transformer replica and copies that timing to every transformer node in the pipeline graph. That assumption holds only when all stages experience identical network conditions. Faulty links are physical, rank-specific degradations—different pipeline stages map to different hardware coordinates, so only the stages whose TP/CP cluster actually touches the faulty link should be slowed. Because hierarchical mode reuses the single measurement for every stage, *all* stages inherit the worst-case transformer time, even if the fault affects just one cluster. Consequently hierarchical mode always over-penalizes the transformer, regardless of where the fault truly sits. + +## Remediation Plan (hierarchical & hybrid) +1. **Decode Fault Coordinates** + - Reconstruct the global axis layout (`tp`, `cp`, `lp`, `dp`) exactly as flattened mode assigns hw_ids. + - For every `faulty_links` entry, convert `src/dst` into `(tp,cp,lp,dp)` tuples, then re-linearize per axis subset. This projection will emit two remapped lists: one for TP/CP-only graphs (transformer) and one for LP/DP graphs (pipeline). Hybrid mode follows the same flow because it still invokes AstraSim for transformer+pipeline pieces separately. + +2. **Route Faults to the Correct Graph** + - Use the projected coordinates to classify each link by the highest dimension it touches (dim0 → transformer cluster, dim≥1 → pipeline/DP axes). Only pass transformer-dimension faults into the transformer AstraSim jobs; only pass higher-dimension faults into the pipeline job. + +3. **Per-Cluster Transformer Simulations** + - Identify the set of TP/CP clusters impacted by faults. Run AstraSim `(N + 1)` times per direction: one baseline (no transformer faults) plus one run per faulty cluster with its specific remapped fault list attached. Hybrid mode needs the same forward/backward pairs. + - Cache the timing per faulty cluster (forward/backward) alongside the baseline. + +4. **Annotate Pipeline Graph Nodes** + - When applying transformer timings in hierarchical/hybrid modes, assign the baseline duration to transformer nodes whose cluster is fault-free, and assign the corresponding per-cluster duration to nodes mapped to faulty hardware. This ensures only the affected stages slow down. + +5. **Pipeline Fault Injection** + - Feed the LP/DP-projected fault list into the pipeline AstraSim execution. Because the endpoints are already remapped to the LP/DP-only axis ordering, the existing hierarchical pipeline simulation will see the correct degraded links. diff --git a/LLM_execution.py b/LLM_execution.py index cc6bcb66..5f29b94e 100644 --- a/LLM_execution.py +++ b/LLM_execution.py @@ -8,6 +8,8 @@ import simulate_LLM from astrasim_lib import run_astra_simulation_only_onepath +from astrasim_lib.fault_projection import FaultProjectionResult, FaultSpace +from astrasim_lib.layout_utils import axis_layout_from_descriptor from simulate_LLM import Graph @@ -639,6 +641,9 @@ def __init__( self._transformer_rank_layout: Optional[Dict[str, Any]] = None self._pipeline_rank_layout: Optional[Dict[str, Any]] = None self._rank_layout = self._build_rank_layout_descriptor() + self._fault_space: Optional[FaultSpace] = None + self._fault_projections: Dict[str, FaultProjectionResult] = {} + self._initialize_fault_mappings() def _build_rank_layout_descriptor(self) -> Dict[str, Any]: hw_config = getattr(self.time_calc, "hw_config", None) @@ -758,6 +763,102 @@ def _subset_layout(allowed: Sequence[str]) -> Optional[Dict[str, Any]]: return descriptor + def _initialize_fault_mappings(self) -> None: + hw_config = getattr(self.time_calc, "hw_config", None) + network_layout = getattr(hw_config, "network_layout", None) + faulty_links: Tuple[Tuple[int, int, float], ...] = tuple( + getattr(network_layout, "faulty_links", ()) or () + ) + if not faulty_links: + self._fault_space = None + self._fault_projections = {} + return + if not self._rank_layout: + raise ValueError("Faulty links require a populated rank layout descriptor.") + axis_layout = axis_layout_from_descriptor(self._rank_layout) + axis_dim_map = self._axis_to_dimension_map(network_layout) + self._fault_space = FaultSpace( + axis_layout, + faulty_links, + axis_to_dimension=axis_dim_map, + ) + self._fault_projections = self._build_fault_projections_from_space(self._fault_space) + self._validate_fault_coverage() + + def _axis_to_dimension_map(self, network_layout) -> Dict[str, int]: + mapping: Dict[str, int] = {} + if network_layout is None: + return mapping + dimensions = getattr(network_layout, "dimensions", None) + if not dimensions: + return mapping + for idx, dim in enumerate(dimensions): + for axis in getattr(dim, "parallelisms", ()) or (): + normalized = str(axis).strip().lower() + if normalized: + mapping[normalized] = idx + return mapping + + def _build_fault_projections_from_space( + self, + space: FaultSpace, + ) -> Dict[str, FaultProjectionResult]: + projections: Dict[str, FaultProjectionResult] = {} + if not space.entries: + return projections + + global_axes = tuple(space.layout.axis_order) + if global_axes: + projections["global"] = space.project(global_axes) + + transformer_axes: Tuple[str, ...] = tuple( + self._transformer_rank_layout.get("axis_order", ()) + ) if self._transformer_rank_layout else tuple() + if transformer_axes: + projections["transformer"] = space.project(transformer_axes) + + pipeline_axes: Tuple[str, ...] = tuple( + self._pipeline_rank_layout.get("axis_order", ()) + ) if self._pipeline_rank_layout else tuple() + if pipeline_axes: + projections["pipeline"] = space.project(pipeline_axes) + + return projections + + def _validate_fault_coverage(self) -> None: + if self._fault_space is None: + return + covered: Set[Tuple[int, int, float]] = set() + for label in ("transformer", "pipeline"): + proj = self._fault_projections.get(label) + if proj: + covered.update(proj.covered_originals) + uncovered = [ + entry.original + for entry in self._fault_space.entries + if entry.original not in covered + ] + if uncovered: + formatted = ", ".join(str(item) for item in uncovered) + raise ValueError( + "Faulty links do not map to any transformer or pipeline axis subset: " + f"{formatted}" + ) + + def _fault_projection_for(self, label: str) -> Optional[FaultProjectionResult]: + return self._fault_projections.get(label) + + def _fault_links_for(self, label: str) -> Tuple[Tuple[int, int, float], ...]: + projection = self._fault_projection_for(label) + if projection is None: + return tuple() + return projection.remapped_links + + def _fault_override(self, label: str) -> Optional[Tuple[Tuple[int, int, float], ...]]: + if self._fault_space is None: + return None + return self._fault_links_for(label) + def run(self, mode: ExecutionMode) -> ExecutionResult: if mode == ExecutionMode.ANALYTICAL: return self._run_pipeline_with_analytical_comm(ExecutionMode.ANALYTICAL) @@ -827,8 +928,10 @@ def _run_full_astrasim_hierarchical(self) -> ExecutionResult: if self.time_calc.persist_astrasim_artifacts: artifact_dir = os.path.join(self.time_calc.output_dir, "astra_hier") + pipeline_fault_override = self._fault_override("pipeline") run_kwargs = { "persist_artifacts": self.time_calc.persist_astrasim_artifacts, + "faulty_links_override": pipeline_fault_override, } run_type = str(getattr(getattr(self.time_calc, "model", None), "run_type", "training")).lower() effective_dp = 1 if run_type == "inference" else max(1, getattr(self.time_calc, "dp", 1)) @@ -1009,6 +1112,7 @@ def _run_transformer_astrasim(self, mode: ExecutionMode) -> Optional[Transformer bwd_per_rank = None fwd_max = 0 bwd_max = 0 + transformer_fault_override = self._fault_override("transformer") if self.transformer_forward_root: layout = getattr(self, "_transformer_rank_layout", None) if layout: @@ -1021,6 +1125,7 @@ def _run_transformer_astrasim(self, mode: ExecutionMode) -> Optional[Transformer artifact_dir_fwd, dp_override=1, persist_artifacts=self.time_calc.persist_astrasim_artifacts, + faulty_links_override=transformer_fault_override, ) if self.transformer_backward_root: layout = getattr(self, "_transformer_rank_layout", None) @@ -1034,6 +1139,7 @@ def _run_transformer_astrasim(self, mode: ExecutionMode) -> Optional[Transformer artifact_dir_bwd, dp_override=1, persist_artifacts=self.time_calc.persist_astrasim_artifacts, + faulty_links_override=transformer_fault_override, ) self.time_calc.transformer_astrasim_per_rank_forward = fwd_per_rank diff --git a/astrasim_lib/config_generation.py b/astrasim_lib/config_generation.py index 03ef3ede..1a9ccf64 100644 --- a/astrasim_lib/config_generation.py +++ b/astrasim_lib/config_generation.py @@ -112,6 +112,9 @@ def generate_astrasim_configs_from_hw( npus_count: Optional[int] = None, *, axes_filter: Optional[Sequence[str]] = None, + transform_2d_to_1d: bool = False, + faulty_links_override: Optional[Sequence[Tuple[int, int, float]]] = None, + reuse_cached_network: bool = True, ) -> Dict[str, str]: """Write AstraSim network/system configs derived from ``hw_obj``.""" if npus_count is None: @@ -119,6 +122,14 @@ def generate_astrasim_configs_from_hw( layout = getattr(hw_obj, "network_layout", None) dimensions = list(getattr(layout, "dimensions", [])) if layout else [] + layout_faults: Tuple[Tuple[int, int, float], ...] + if faulty_links_override is not None: + layout_faults = tuple( + (int(src), int(dst), float(weight)) + for src, dst, weight in faulty_links_override + ) + else: + layout_faults = tuple(getattr(layout, "faulty_links", ()) or ()) if layout else tuple() if not dimensions: raise ValueError("Hardware config is missing network dimensions required for AstraSim integration.") @@ -207,14 +218,8 @@ def _safe_int(value: Any, default: int = 1) -> int: for dim, _, _, dim_idx in dim_infos if int(getattr(dim, "size", 0) or 0) > 1 ] - allowed_fault_dim_indices: set[int] = set() if active_dim_indices: last_active_idx = active_dim_indices[-1] - if astra_mode in {"full_astrasim_hierarchical", "hybrid"}: - allowed_fault_dim_indices.add(0) - allowed_fault_dim_indices.add(last_active_idx) - else: - allowed_fault_dim_indices.add(last_active_idx) dp_in_active = False if axis_sizes.get("dp", 1) > 1: for dim, axes, _, dim_idx in dim_infos: @@ -230,20 +235,10 @@ def _safe_int(value: Any, default: int = 1) -> int: "Hardware config declares dp > 1, but no network dimension with size > 1 includes the 'dp' axis." ) else: - if any(getattr(dim, "faulty_links", ()) for dim, _, _, _ in dim_infos): + if layout_faults: raise ValueError( "Faulty links require at least one active (size > 1) network dimension." ) - - if allowed_fault_dim_indices: - valid_indices = {idx for idx in allowed_fault_dim_indices if 0 <= idx < len(dimensions)} - for dim, _, _, dim_idx in dim_infos: - if getattr(dim, "faulty_links", ()): - if dim_idx not in valid_indices: - raise ValueError( - f"Faulty links for dimension '{dim.label}' (index {dim_idx}) are only permitted on " - f"dimension indices {sorted(valid_indices)} for mode '{astra_mode or 'unspecified'}'." - ) target = int(npus_count) axes_needed: List[str] = [] @@ -297,6 +292,13 @@ def _safe_int(value: Any, default: int = 1) -> int: for dim, axes_selected, product_size, _ in selected_dims: topo = _normalize_topology_name(dim.topology_type) size = product_size + if transform_2d_to_1d and topo in ["Torus2D", "Mesh2D", "KingMesh2D"]: + if topo == "Torus2D": + topo = "Ring" + if topo == "Mesh2D": + topo = "Mesh" + if topo == "KingMesh2D": + topo = "HyperCube" if topo == "Ring" and size <= 2: topo = "FullyConnected" effective_bw = float(getattr(dim, "effective_bandwidth", dim.bandwidth)) @@ -324,23 +326,14 @@ def _safe_int(value: Any, default: int = 1) -> int: f"latency: [ {lat_str} ] # ns\n" ) - faulty_links_tuple: Tuple[Tuple[int, int, float], ...] = tuple() - if selected_dims: - dims_with_faults = [ - (dim_idx, getattr(dim, "faulty_links", ())) - for dim, _axes, _size, dim_idx in selected_dims - if getattr(dim, "faulty_links", ()) + faulty_links_tuple: Tuple[Tuple[int, int, float], ...] = tuple(layout_faults) + if faulty_links_tuple: + fault_entries = [ + f"[{src}, {dst}, {float(weight):g}]" + for src, dst, weight in faulty_links_tuple ] - if dims_with_faults: - dims_with_faults.sort(key=lambda item: item[0]) - _, faulty_links_tuple = dims_with_faults[-1] - faulty_links_tuple = tuple(faulty_links_tuple) - fault_entries = [ - f"[{src}, {dst}, {float(weight):g}]" - for src, dst, weight in faulty_links_tuple - ] - faulty_links_str = f"[{', '.join(fault_entries)}]" - net_content += f"faulty_links: {faulty_links_str}\n" + faulty_links_str = f"[{', '.join(fault_entries)}]" + net_content += f"faulty_links: {faulty_links_str}\n" os.makedirs(os.path.dirname(net_yaml), exist_ok=True) cache_key = ( @@ -350,7 +343,7 @@ def _safe_int(value: Any, default: int = 1) -> int: tuple(lat_list), faulty_links_tuple, ) - cached_path = _NET_YAML_CACHE.get(cache_key) + cached_path = _NET_YAML_CACHE.get(cache_key) if reuse_cached_network else None need_write = True if cached_path and os.path.exists(cached_path): try: @@ -366,7 +359,8 @@ def _safe_int(value: Any, default: int = 1) -> int: with open(tmp_path, "w", encoding="utf-8") as handle: handle.write(net_content) os.replace(tmp_path, net_yaml) - _NET_YAML_CACHE[cache_key] = net_yaml + if reuse_cached_network: + _NET_YAML_CACHE[cache_key] = net_yaml if astra_cfg: coll = astra_cfg.collectives diff --git a/astrasim_lib/executor.py b/astrasim_lib/executor.py index 8dc8c84a..c03fb2bc 100644 --- a/astrasim_lib/executor.py +++ b/astrasim_lib/executor.py @@ -40,6 +40,7 @@ write_et_node, ) from .integration import get_remote_memory_path, run_astrasim_analytical, run_cache_astrasim +from .layout_utils import derive_axes_filter from simulate_LLM import visualize_graph @@ -165,25 +166,6 @@ def _build_axis_groups( return axis_groups -def _derive_axes_filter( - axis_order: Sequence[str], - axis_sizes: Mapping[str, int], - dp_count: int, -) -> Optional[List[str]]: - axes: List[str] = [] - for axis in axis_order: - try: - size = max(1, int(axis_sizes.get(axis, 1))) - except Exception: - size = 1 - if size > 1 and axis not in axes: - axes.append(axis) - if dp_count > 1: - axes = [axis for axis in axes if axis != "dp"] - axes.append("dp") - return axes or None - - def _assign_collective_labels( tp_collective_groups: Mapping[str, List[Any]], collective_info: Mapping[Any, Dict[str, Any]], @@ -1485,6 +1467,7 @@ def run_astra_simulation_only_onepath( output_dir: str = "astra_comparison_output", dp_override: Optional[int] = None, persist_artifacts: Optional[bool] = None, + faulty_links_override: Optional[Sequence[Tuple[int, int, float]]] = None, ): """ Run AstraSim simulation on DeepFlow graph and print results. @@ -1493,6 +1476,7 @@ def run_astra_simulation_only_onepath( fwdbwd_root: Forward and backward graph root node time_calc_obj: TimeCalculationLLM object with hw_config and dp attributes output_dir: Directory for temporary files and results + faulty_links_override: Optional remapped faulty link list for this run """ print("\n" + "="*60) print("ASTRASIM SIMULATION RESULTS") @@ -1559,7 +1543,9 @@ def run_astra_simulation_only_onepath( rank_layout = getattr(fwdbwd_root, "_astrasim_rank_layout", None) axis_order, axis_sizes = _extract_axis_layout(rank_layout) - axes_filter = _derive_axes_filter(axis_order, axis_sizes, dp_count) + axes_filter = derive_axes_filter(axis_order, axis_sizes, dp_count) + if not axes_filter: + axes_filter = None if synthetic_pair: axes_filter = ["synthetic2"] astra_configs = generate_astrasim_configs_from_hw( @@ -1567,6 +1553,8 @@ def run_astra_simulation_only_onepath( work_dir, rank_count, axes_filter=axes_filter, + faulty_links_override=faulty_links_override, + reuse_cached_network=not persist, ) # Run AstraSim simulation on forward graph (cached via manifest) diff --git a/astrasim_lib/fault_projection.py b/astrasim_lib/fault_projection.py new file mode 100644 index 00000000..54a2e12b --- /dev/null +++ b/astrasim_lib/fault_projection.py @@ -0,0 +1,159 @@ +"""Utilities for projecting faulty link endpoints onto axis subsets.""" + +from __future__ import annotations + +from dataclasses import dataclass +from typing import Dict, Iterable, Mapping, Optional, Sequence, Tuple, List + +from .layout_utils import AxisLayout, decode_axis_coordinates, build_subset_strides + + +def _coords_dict(coords: Tuple[Tuple[str, int], ...]) -> Dict[str, int]: + return {axis: value for axis, value in coords} + + +def _normalize_axes(axes: Sequence[str]) -> Tuple[str, ...]: + normalized: List[str] = [] + for entry in axes: + axis = str(entry).strip().lower() + if not axis: + continue + if axis not in normalized: + normalized.append(axis) + return tuple(normalized) + + +@dataclass(frozen=True) +class FaultSpaceEntry: + original: Tuple[int, int, float] + src_coords: Tuple[Tuple[str, int], ...] + dst_coords: Tuple[Tuple[str, int], ...] + affected_axes: Tuple[str, ...] + affected_dimensions: Tuple[int, ...] + + +@dataclass(frozen=True) +class FaultProjectionDetail: + original: Tuple[int, int, float] + remapped: Tuple[int, int, float] + subset_axes: Tuple[str, ...] + affected_axes: Tuple[str, ...] + affected_dimensions: Tuple[int, ...] + src_coords: Tuple[Tuple[str, int], ...] + dst_coords: Tuple[Tuple[str, int], ...] + + +@dataclass(frozen=True) +class FaultProjectionResult: + subset_axes: Tuple[str, ...] + remapped_links: Tuple[Tuple[int, int, float], ...] + entries: Tuple[FaultProjectionDetail, ...] + covered_originals: Tuple[Tuple[int, int, float], ...] + skipped_originals: Tuple[Tuple[int, int, float], ...] + + +class FaultSpace: + """Decoded view of global faulty links for downstream projections.""" + + def __init__( + self, + layout: AxisLayout, + faulty_links: Iterable[Tuple[int, int, float]], + axis_to_dimension: Optional[Mapping[str, int]] = None, + ) -> None: + entries: List[FaultSpaceEntry] = [] + layout_order = layout.axis_order + for idx, link in enumerate(faulty_links): + if len(link) != 3: + raise ValueError(f"faulty_links[{idx}] must be a (src, dst, weight) tuple.") + src, dst, weight = link + src_coords = decode_axis_coordinates(int(src), layout) + dst_coords = decode_axis_coordinates(int(dst), layout) + src_dict = _coords_dict(src_coords) + dst_dict = _coords_dict(dst_coords) + affected_axes = tuple( + axis for axis in layout_order if src_dict.get(axis) != dst_dict.get(axis) + ) + if axis_to_dimension is not None: + dims = { + axis_to_dimension[axis] + for axis in affected_axes + if axis in axis_to_dimension + } + affected_dimensions = tuple(sorted(dims)) + else: + affected_dimensions = tuple() + entries.append( + FaultSpaceEntry( + original=(int(src), int(dst), float(weight)), + src_coords=src_coords, + dst_coords=dst_coords, + affected_axes=affected_axes, + affected_dimensions=affected_dimensions, + ) + ) + self._entries: Tuple[FaultSpaceEntry, ...] = tuple(entries) + self._layout = layout + + @property + def layout(self) -> AxisLayout: + return self._layout + + @property + def entries(self) -> Tuple[FaultSpaceEntry, ...]: + return self._entries + + def project(self, subset_axes: Sequence[str]) -> FaultProjectionResult: + normalized = _normalize_axes(subset_axes) + if not normalized: + raise ValueError("subset_axes must not be empty when projecting faults.") + subset_axes_set = set(normalized) + subset_strides = build_subset_strides(normalized, self._layout.axis_sizes) + details: List[FaultProjectionDetail] = [] + remapped_links: List[Tuple[int, int, float]] = [] + covered: List[Tuple[int, int, float]] = [] + unmapped: List[Tuple[int, int, float]] = [] + + for entry in self._entries: + if not any(axis in subset_axes_set for axis in entry.affected_axes): + unmapped.append(entry.original) + continue + src_dict = _coords_dict(entry.src_coords) + dst_dict = _coords_dict(entry.dst_coords) + src_linear = _linearize_subset(src_dict, normalized, subset_strides) + dst_linear = _linearize_subset(dst_dict, normalized, subset_strides) + remapped = (src_linear, dst_linear, entry.original[2]) + details.append( + FaultProjectionDetail( + original=entry.original, + remapped=remapped, + subset_axes=normalized, + affected_axes=entry.affected_axes, + affected_dimensions=entry.affected_dimensions, + src_coords=entry.src_coords, + dst_coords=entry.dst_coords, + ) + ) + remapped_links.append(remapped) + covered.append(entry.original) + + return FaultProjectionResult( + subset_axes=normalized, + remapped_links=tuple(remapped_links), + entries=tuple(details), + covered_originals=tuple(covered), + skipped_originals=tuple(unmapped), + ) + + +def _linearize_subset( + coords: Mapping[str, int], + subset_axes: Sequence[str], + subset_strides: Mapping[str, int], +) -> int: + linear = 0 + for axis in subset_axes: + coord = coords.get(axis, 0) + stride = subset_strides[axis] + linear += coord * stride + return linear diff --git a/astrasim_lib/integration.py b/astrasim_lib/integration.py index 6b942d79..8108dfc2 100644 --- a/astrasim_lib/integration.py +++ b/astrasim_lib/integration.py @@ -318,6 +318,7 @@ def run_cache_astrasim( workload_prefix: Optional[str] = None, comm_group_json: Optional[str] = None, axes_filter: Optional[Sequence[str]] = None, + transform_2d_to_1d: bool = False, files = None, ) -> Tuple[List[float], float]: """Run AstraSim with caching to avoid recomputation when inputs match.""" @@ -376,6 +377,7 @@ def run_cache_astrasim( out_dir=astra_config_dir, npus_count=npus_count, axes_filter=axes_filter_tuple, + transform_2d_to_1d=transform_2d_to_1d, ) topo_list = files.get("topology_list") npus_per_dim = files.get("npus_per_dim") diff --git a/astrasim_lib/layout_utils.py b/astrasim_lib/layout_utils.py new file mode 100644 index 00000000..7cea167a --- /dev/null +++ b/astrasim_lib/layout_utils.py @@ -0,0 +1,154 @@ +"""Shared helpers for reasoning about DeepFlow axis layouts.""" + +from __future__ import annotations + +from dataclasses import dataclass +from typing import Any, Dict, Mapping, Sequence, Tuple, List + + +def _normalize_axis_name(name: str) -> str: + if not isinstance(name, str): + raise TypeError(f"Axis name must be a string, got {type(name)!r}") + normalized = name.strip().lower() + if not normalized: + raise ValueError("Axis name cannot be empty or whitespace.") + return normalized + + +@dataclass(frozen=True) +class AxisLayout: + """Immutable view of the flattened axis layout that defines hw_id numbering.""" + + axis_order: Tuple[str, ...] + axis_sizes: Dict[str, int] + axis_strides: Dict[str, int] + total_ranks: int + + def __post_init__(self) -> None: + if any(size < 1 for size in self.axis_sizes.values()): + raise ValueError("Axis sizes must all be >= 1.") + for axis in self.axis_order: + if axis not in self.axis_sizes: + raise ValueError(f"Axis '{axis}' is missing a declared size.") + if axis not in self.axis_strides: + raise ValueError(f"Axis '{axis}' is missing a declared stride.") + + +def axis_layout_from_descriptor(descriptor: Mapping[str, Any]) -> AxisLayout: + """Create an :class:`AxisLayout` from `_build_rank_layout_descriptor` output.""" + if descriptor is None: + raise ValueError("Axis layout descriptor must not be None.") + + axis_order_raw = descriptor.get("axis_order", []) + if not isinstance(axis_order_raw, Sequence): + raise TypeError("Axis layout descriptor axis_order must be a sequence.") + axis_order: Tuple[str, ...] = tuple(_normalize_axis_name(entry) for entry in axis_order_raw) + if not axis_order: + raise ValueError("Axis layout descriptor axis_order is empty.") + + sizes_raw = descriptor.get("axis_sizes", {}) + if not isinstance(sizes_raw, Mapping): + raise TypeError("Axis layout descriptor axis_sizes must be a mapping.") + axis_sizes: Dict[str, int] = {} + for key, value in sizes_raw.items(): + axis = _normalize_axis_name(str(key)) + try: + size = int(value) + except (TypeError, ValueError) as exc: + raise ValueError(f"Axis '{axis}' size must be an integer.") from exc + if size < 1: + raise ValueError(f"Axis '{axis}' size must be >= 1.") + axis_sizes[axis] = size + + strides_raw = descriptor.get("axis_strides", {}) + if strides_raw and not isinstance(strides_raw, Mapping): + raise TypeError("Axis layout descriptor axis_strides must be a mapping when provided.") + + axis_strides: Dict[str, int] = {} + if strides_raw: + for key, value in strides_raw.items(): + axis = _normalize_axis_name(str(key)) + try: + stride = int(value) + except (TypeError, ValueError) as exc: + raise ValueError(f"Axis '{axis}' stride must be an integer.") from exc + if stride < 0: + raise ValueError(f"Axis '{axis}' stride must be >= 0.") + axis_strides[axis] = stride + + if not axis_strides: + span = 1 + for axis in axis_order: + axis_strides[axis] = span + span *= axis_sizes.get(axis, 1) + + total = 1 + for axis in axis_order: + total *= axis_sizes[axis] + + return AxisLayout( + axis_order=axis_order, + axis_sizes=axis_sizes, + axis_strides=axis_strides, + total_ranks=total, + ) + + +def decode_axis_coordinates(hw_id: int, layout: AxisLayout) -> Tuple[Tuple[str, int], ...]: + """Return ordered axis coordinates for ``hw_id``.""" + if hw_id < 0: + raise ValueError(f"Hardware id must be >= 0, got {hw_id}.") + if hw_id >= layout.total_ranks: + raise ValueError( + f"Hardware id {hw_id} exceeds layout span ({layout.total_ranks} ranks)." + ) + coords: List[Tuple[str, int]] = [] + for axis in layout.axis_order: + stride = layout.axis_strides[axis] + size = layout.axis_sizes[axis] + if size == 1: + coords.append((axis, 0)) + continue + coord = (hw_id // stride) % size + coords.append((axis, coord)) + return tuple(coords) + + +def build_subset_strides( + subset_axes: Sequence[str], + axis_sizes: Mapping[str, int], +) -> Dict[str, int]: + """Compute strides for a subset axis ordering.""" + strides: Dict[str, int] = {} + span = 1 + for raw_axis in subset_axes: + axis = _normalize_axis_name(raw_axis) + if axis not in axis_sizes: + raise ValueError(f"Axis '{axis}' is not defined in the layout axis_sizes.") + size = axis_sizes[axis] + if size < 1: + raise ValueError(f"Axis '{axis}' size must be >= 1.") + strides[axis] = span + span *= size + return strides + + +def derive_axes_filter( + axis_order: Sequence[str], + axis_sizes: Mapping[str, int], + dp_count: int, +) -> List[str]: + """Return the axes filter used to select network dimensions for AstraSim.""" + axes: List[str] = [] + for entry in axis_order: + axis = _normalize_axis_name(entry) + try: + size = max(1, int(axis_sizes.get(axis, 1))) + except Exception: + size = 1 + if size > 1 and axis not in axes: + axes.append(axis) + if dp_count > 1: + axes = [axis for axis in axes if axis != "dp"] + axes.append("dp") + return axes diff --git a/config.py b/config.py index 85560466..7115fd64 100644 --- a/config.py +++ b/config.py @@ -354,7 +354,6 @@ class NetworkDimensionLayout: collective_override: Dict[str, str] = field(default_factory=dict) parallelisms: Tuple[str, ...] = field(default_factory=tuple) energy_per_bit: float = 0.0 - faulty_links: Tuple[Tuple[int, int, float], ...] = field(default_factory=tuple) @classmethod def from_raw( @@ -476,51 +475,6 @@ def from_raw( expected_product=computed_product, ) - faulty_links_raw = raw.get("faulty_links", []) - faulty_links: List[Tuple[int, int, float]] = [] - if faulty_links_raw: - if not isinstance(faulty_links_raw, Sequence) or isinstance(faulty_links_raw, (str, bytes)): - raise ValueError( - f"network dimension '{label}' faulty_links must be a sequence of [src, dst, weight] entries" - ) - for idx, entry in enumerate(faulty_links_raw): - if ( - not isinstance(entry, Sequence) - or isinstance(entry, (str, bytes)) - or len(entry) != 3 - ): - raise ValueError( - f"network dimension '{label}' faulty_links[{idx}] must be a three-item sequence [src, dst, weight]" - ) - src_raw, dst_raw, weight_raw = entry - try: - src = int(src_raw) - except (TypeError, ValueError) as exc: - raise ValueError( - f"network dimension '{label}' faulty_links[{idx}][0] must be an integer endpoint" - ) from exc - try: - dst = int(dst_raw) - except (TypeError, ValueError) as exc: - raise ValueError( - f"network dimension '{label}' faulty_links[{idx}][1] must be an integer endpoint" - ) from exc - if src < 0 or dst < 0: - raise ValueError( - f"network dimension '{label}' faulty_links[{idx}] endpoints must be >= 0" - ) - try: - weight = float(weight_raw) - except (TypeError, ValueError) as exc: - raise ValueError( - f"network dimension '{label}' faulty_links[{idx}][2] must be a numeric reliability weight" - ) from exc - if weight < 0.0 or weight > 1.0: - raise ValueError( - f"network dimension '{label}' faulty_links[{idx}] weight must be between 0.0 and 1.0" - ) - faulty_links.append((src, dst, weight)) - return cls( id=dim_id, label=label, @@ -532,7 +486,6 @@ def from_raw( collective_override=collective_override, parallelisms=tuple(normalized_parallelisms), energy_per_bit=energy_per_bit, - faulty_links=tuple(faulty_links), ) @property @@ -611,15 +564,23 @@ def _compute_dimension_parallelism_product( def _parse_network_layout( network_spec, parallelism_params: Dict[str, object], -) -> Tuple[NetworkDimensionLayout, ...]: +) -> Tuple[Tuple[NetworkDimensionLayout, ...], Tuple[Tuple[int, int, float], ...]]: if network_spec is None: - return tuple() + return tuple(), tuple() + + faulty_links: Tuple[Tuple[int, int, float], ...] = tuple() + dimensions_spec = network_spec + if isinstance(network_spec, dict): + faulty_links = _parse_faulty_links("network", network_spec.get("faulty_links", [])) + dimensions_spec = network_spec.get("dimensions") + if dimensions_spec is None: + raise ValueError("network.dimensions must be specified when network is a mapping") - if not isinstance(network_spec, Sequence) or isinstance(network_spec, (str, bytes)): - raise ValueError("network section must be a sequence of dimension mappings") + if not isinstance(dimensions_spec, Sequence) or isinstance(dimensions_spec, (str, bytes)): + raise ValueError("network.dimensions must be a sequence of dimension mappings") dimensions: List[NetworkDimensionLayout] = [] - for index, entry in enumerate(network_spec): + for index, entry in enumerate(dimensions_spec): dimensions.append( NetworkDimensionLayout.from_raw( entry, @@ -627,12 +588,13 @@ def _parse_network_layout( index=index, ) ) - return tuple(dimensions) + return tuple(dimensions), faulty_links @dataclass(frozen=True) class NetworkLayoutConfig: dimensions: Tuple[NetworkDimensionLayout, ...] + faulty_links: Tuple[Tuple[int, int, float], ...] = field(default_factory=tuple) parallelism_map: Dict[str, NetworkDimensionLayout] = field(default_factory=dict) def primary_dimension(self) -> Optional[NetworkDimensionLayout]: @@ -651,8 +613,53 @@ def link_for_parallelism(self, name: str) -> Tuple[float, float]: return dim.effective_bandwidth, dim.latency +def _parse_faulty_links(owner_label: str, faulty_links_raw) -> Tuple[Tuple[int, int, float], ...]: + entries: List[Tuple[int, int, float]] = [] + if not faulty_links_raw: + return tuple() + if not isinstance(faulty_links_raw, Sequence) or isinstance(faulty_links_raw, (str, bytes)): + raise ValueError( + f"{owner_label} faulty_links must be a sequence of [src, dst, weight] entries" + ) + for idx, entry in enumerate(faulty_links_raw): + if not isinstance(entry, Sequence) or isinstance(entry, (str, bytes)) or len(entry) != 3: + raise ValueError( + f"{owner_label} faulty_links[{idx}] must be a three-item sequence [src, dst, weight]" + ) + src_raw, dst_raw, weight_raw = entry + try: + src = int(src_raw) + except (TypeError, ValueError) as exc: + raise ValueError( + f"{owner_label} faulty_links[{idx}][0] must be an integer endpoint" + ) from exc + try: + dst = int(dst_raw) + except (TypeError, ValueError) as exc: + raise ValueError( + f"{owner_label} faulty_links[{idx}][1] must be an integer endpoint" + ) from exc + if src < 0 or dst < 0: + raise ValueError( + f"{owner_label} faulty_links[{idx}] endpoints must be >= 0" + ) + try: + weight = float(weight_raw) + except (TypeError, ValueError) as exc: + raise ValueError( + f"{owner_label} faulty_links[{idx}][2] must be a numeric reliability weight" + ) from exc + if weight < 0.0 or weight > 1.0: + raise ValueError( + f"{owner_label} faulty_links[{idx}] weight must be between 0.0 and 1.0" + ) + entries.append((src, dst, weight)) + return tuple(entries) + + def _build_network_layout_config( dimensions: Sequence[NetworkDimensionLayout], + faulty_links: Sequence[Tuple[int, int, float]] = (), ) -> NetworkLayoutConfig: parallelism_map: Dict[str, NetworkDimensionLayout] = {} for dim in dimensions: @@ -662,7 +669,11 @@ def _build_network_layout_config( f"parallelism '{pname}' assigned to multiple network dimensions" ) parallelism_map[pname] = dim - return NetworkLayoutConfig(dimensions=tuple(dimensions), parallelism_map=parallelism_map) + return NetworkLayoutConfig( + dimensions=tuple(dimensions), + faulty_links=tuple(faulty_links), + parallelism_map=parallelism_map, + ) _PARALLELISM_DEFAULTS: Dict[str, object] = { @@ -1066,10 +1077,10 @@ def parse_config(filename, config_type): scheduling_for_network = {str(k).lower(): v for k, v in parallelism_params.items()} network_spec = config_dict.get("network") - network_dimensions = _parse_network_layout(network_spec, scheduling_for_network) + network_dimensions, network_faults = _parse_network_layout(network_spec, scheduling_for_network) if not network_dimensions: raise ValueError("network section must define at least one dimension") - network_layout_config = _build_network_layout_config(network_dimensions) + network_layout_config = _build_network_layout_config(network_dimensions, network_faults) tech_config = TechConfig.from_dict(config_dict["tech_param"]) diff --git a/time_calculation.py b/time_calculation.py index b108b5ab..3d9deb1b 100644 --- a/time_calculation.py +++ b/time_calculation.py @@ -33,12 +33,15 @@ def _astra_collective(self, kind: str, participants: int, size_bytes: int, axis: return 0.0 byte_count = int(math.ceil(size_bytes)) axes_filter = [str(axis).lower()] if axis else None + # for collectives ONLY, we cannot use 2D topologies (Mesh2D, Torus2D, KingMesh2D) + # transform them to their 1D equivalents (Ring, Mesh, HyperCube (?)) _, max_sec = run_cache_astrasim( self.hw_config, comm=kind, npus_count=part, size_bytes=byte_count, axes_filter=axes_filter, + transform_2d_to_1d=True, ) return float(max_sec) diff --git a/tools/exec_with_fault.py b/tools/exec_with_fault.py new file mode 100644 index 00000000..8ee47b2a --- /dev/null +++ b/tools/exec_with_fault.py @@ -0,0 +1,212 @@ +#!/usr/bin/env python3 +""" +Execute the LLM AstraSim debug harness under multiple faulty-link scenarios. + +For each predefined configuration this script patches the +``configs/hardware-config/a100_80GB_faraz.yaml`` file with the desired +``network.faulty_links`` list, invokes ``./examples/llm_astra_inference_debug_graphviz.sh``, +parses the reported training time, and records the results in a summary table. + +The original YAML contents are restored once the sweep finishes (even if a run fails). +""" + +from __future__ import annotations + +import argparse +import contextlib +import copy +import os +import re +import subprocess +import sys +from dataclasses import dataclass +from pathlib import Path +from typing import Any, Iterable, List, Sequence + +import yaml + + +PROJECT_ROOT = Path(__file__).resolve().parent.parent +CONFIG_PATH = PROJECT_ROOT / "configs" / "hardware-config" / "a100_80GB_faraz.yaml" +RUN_SCRIPT = "./examples/llm_astra_inference_debug_graphviz.sh" + +# Faulty-link sweep configurations. +# Each entry provides a human-readable label and the list of faulty links to inject. +# Faulty links follow the AstraSim convention: [src, dst, reliability_weight] +FAULT_SWEEP_CASES: Sequence[dict[str, Any]] = [ + {"label": "baseline", "faulty_links": []}, + {"label": "tp_top_left_harsh", "faulty_links": [[0, 1, 0.3]]}, + {"label": "tp_top_left", "faulty_links": [[0, 1, 0.5]]}, + {"label": "tp_sequential", "faulty_links": [[0, 1, 0.5], [1, 2, 0.5]]}, + {"label": "tp_top_left_vert", "faulty_links": [[0, 1, 0.5], [0, 4, 0.5]]}, + {"label": "tp_sequential_harsh", "faulty_links": [[0, 1, 0.3], [1, 2, 0.3]]}, + {"label": "tp_top_left_vert_harsh", "faulty_links": [[0, 1, 0.3], [0, 4, 0.3]]}, +] + + +TRAINING_TIME_PATTERN = re.compile(r"Training time for batch:\s*([0-9.]+)s") + + +@dataclass +class SweepResult: + label: str + links: Sequence[Sequence[float]] + returncode: int + training_time: float | None + stdout: str + stderr: str + + +def _load_yaml(path: Path) -> Any: + with path.open("r", encoding="utf-8") as handle: + return yaml.safe_load(handle) + + +def _dump_yaml(path: Path, data: Any) -> None: + with path.open("w", encoding="utf-8") as handle: + yaml.safe_dump(data, handle, sort_keys=False) + + +def _set_faulty_links(doc: Any, links: Sequence[Sequence[float]]) -> None: + network_block = doc.get("network") + if isinstance(network_block, dict): + network_block["faulty_links"] = copy.deepcopy(list(links)) + return + raise ValueError("Hardware config network section must be a mapping with 'faulty_links'.") + + +def _run_script(cwd: Path) -> subprocess.CompletedProcess[str]: + return subprocess.run( + [RUN_SCRIPT], + cwd=str(cwd), + check=False, + text=True, + capture_output=True, + shell=False, + ) + + +def _parse_training_time(output: str) -> float | None: + match = TRAINING_TIME_PATTERN.search(output) + if not match: + return None + try: + return float(match.group(1)) + except ValueError: + return None + + +def _format_summary(results: Iterable[SweepResult]) -> str: + lines: List[str] = [] + header = f"{'Label':<24} {'Training Time (s)':>18} {'Return Code':>12}" + lines.append(header) + lines.append("-" * len(header)) + for result in results: + time_str = f"{result.training_time:.2f}" if result.training_time is not None else "n/a" + lines.append(f"{result.label:<24} {time_str:>18} {result.returncode:>12}") + return "\n".join(lines) + + +def run_sweep(cases: Sequence[dict[str, Any]], *, dry_run: bool = False) -> List[SweepResult]: + original_yaml = _load_yaml(CONFIG_PATH) + original_links = copy.deepcopy(original_yaml.get("network", {}).get("faulty_links")) + results: List[SweepResult] = [] + + try: + for case in cases: + label = case.get("label", "unnamed") + links = case.get("faulty_links", []) + print(f"[exec_with_fault] Running case '{label}' with links={links}") + + next_yaml = copy.deepcopy(original_yaml) + _set_faulty_links(next_yaml, links) + if not dry_run: + _dump_yaml(CONFIG_PATH, next_yaml) + + completed: subprocess.CompletedProcess[str] + if dry_run: + completed = subprocess.CompletedProcess( + args=[RUN_SCRIPT], returncode=0, stdout="", stderr="" + ) + else: + completed = _run_script(PROJECT_ROOT) + + combined_output = f"{completed.stdout}\n{completed.stderr}" + training_time = _parse_training_time(combined_output) + if training_time is None: + print( + f"[exec_with_fault] Warning: training time not found in output for case '{label}'.", + file=sys.stderr, + ) + + results.append( + SweepResult( + label=label, + links=links, + returncode=completed.returncode, + training_time=training_time, + stdout=completed.stdout, + stderr=completed.stderr, + ) + ) + if completed.returncode != 0: + print( + f"[exec_with_fault] Case '{label}' exited with return code {completed.returncode}.", + file=sys.stderr, + ) + finally: + # Restore original faulty links (or remove the key if it was missing) + restored_yaml = _load_yaml(CONFIG_PATH) + network_block = restored_yaml.get("network") + if isinstance(network_block, dict): + if original_links is None: + network_block.pop("faulty_links", None) + else: + network_block["faulty_links"] = original_links + if not dry_run: + _dump_yaml(CONFIG_PATH, restored_yaml) + + return results + + +def main(argv: Sequence[str] | None = None) -> int: + parser = argparse.ArgumentParser(description="Execute AstraSim runs with predefined faulty links.") + parser.add_argument( + "--dry-run", + action="store_true", + help="Do not modify files or run the benchmark; useful for quick validation.", + ) + parser.add_argument( + "--cases", + nargs="*", + help="Optional subset of case labels to run. Defaults to the entire FAULT_SWEEP_CASES list.", + ) + args = parser.parse_args(argv) + + if args.dry_run: + print("[exec_with_fault] Dry run enabled: configurations will not be written or executed.") + + selected_cases = [] + if args.cases: + case_set = set(args.cases) + for case in FAULT_SWEEP_CASES: + if case.get("label") in case_set: + selected_cases.append(case) + missing = case_set - {case.get("label") for case in selected_cases} + if missing: + parser.error(f"Unknown case label(s): {', '.join(sorted(missing))}") + else: + selected_cases = list(FAULT_SWEEP_CASES) + + if not selected_cases: + parser.error("No cases selected for execution.") + + results = run_sweep(selected_cases, dry_run=args.dry_run) + print("\nSweep summary:\n") + print(_format_summary(results)) + + return 0 + + +if __name__ == "__main__": + raise SystemExit(main()) diff --git a/tools/fault_meta_sweep.py b/tools/fault_meta_sweep.py new file mode 100644 index 00000000..3425643e --- /dev/null +++ b/tools/fault_meta_sweep.py @@ -0,0 +1,189 @@ +#!/usr/bin/env python3 +""" +Meta sweep orchestrator for DeepFlow fault sensitivity experiments. + +Runs a sequence of configurations by invoking ``fault_sweep.py`` with tailored +arguments, collects artefacts under ``tools/fault_sweep_results/``, and skips +any configuration whose plot already exists. +""" + +import subprocess +import sys +from pathlib import Path +from typing import List + + +RESULTS_DIR = Path("tools", "fault_sweep_results") +RESULTS_DIR.mkdir(parents=True, exist_ok=True) + + +def run_config(name: str, extra_args: List[str]) -> None: + plot_path = RESULTS_DIR / f"{name}.png" + if plot_path.exists(): + print(f"[meta] Skipping '{name}' (plot already exists).") + return + + report_path = RESULTS_DIR / f"{name}.tsv" + results_json = RESULTS_DIR / f"{name}_results.json" + + cmd = [ + sys.executable, + "tools/fault_sweep.py", + "--plot-output", + str(plot_path), + "--report-output", + str(report_path), + "--results-json", + str(results_json), + ] + cmd.extend(extra_args) + + print(f"[meta] Running configuration '{name}'") + completed = subprocess.run(cmd, text=True) + if completed.returncode != 0: + raise RuntimeError(f"Configuration '{name}' failed with exit code {completed.returncode}") + + +def main() -> None: + common_args = [ + "--target-num-gpus", + "256", + "--fault-iter", + "200", + "--fault-workers", + "86", + "--fault-mag", + "0.5,0.1", + "--num-faults", + "1", + "--min-tp", + "1", + "--max-tp", + "16", + "--min-cp", + "1", + "--max-cp", + "256", + "--min-dp", + "1", + "--max-dp", + "256", + "--min-lp", + "1", + "--max-lp", + "128", + "--seed", + "1337", + ] + + configs = [ + ( + "dense_1fault_alldim_256", + common_args + + [ + "--sample-count", + "25", + "--fault-iter", + "50", + "--fault-workers", + "96", + ], + ), + ( + "dense_1fault_dim0_256", + common_args + + [ + "--sample-count", + "48", + "--allowed-fault-dims", + "0", + ], + ), + ( + "sparse_1fault_alldim_tpcponly_256", + [ + "--target-num-gpus", + "256", + "--fault-iter", + "100", + "--fault-workers", + "86", + "--fault-mag", + "0.5,0.1", + "--num-faults", + "1", + "--sample-count", + "10", + "--min-tp", + "1", + "--max-tp", + "64", + "--min-cp", + "1", + "--max-cp", + "128", + "--min-dp", + "1", + "--max-dp", + "1", + "--min-lp", + "1", + "--max-lp", + "1", + "--seed", + "1337", + ], + ), + ( + "sparse_1fault_alldim_dplponly_256", + [ + "--target-num-gpus", + "256", + "--fault-iter", + "200", + "--fault-workers", + "86", + "--fault-mag", + "0.5,0.1", + "--num-faults", + "1", + "--sample-count", + "20", + "--min-tp", + "1", + "--max-tp", + "1", + "--min-cp", + "1", + "--max-cp", + "1", + "--min-dp", + "1", + "--max-dp", + "256", + "--min-lp", + "1", + "--max-lp", + "128", + "--seed", + "1337", + ], + ), + ( + "dense_1fault_alldim_256_harsh", + common_args + + [ + "--sample-count", + "48", + "--fault-mag", + "0.3,0.2", + ], + ), + ] + + for name, args in configs: + run_config(name, args) + + +if __name__ == "__main__": + main() diff --git a/tools/fault_sweep.py b/tools/fault_sweep.py new file mode 100644 index 00000000..c85d3e57 --- /dev/null +++ b/tools/fault_sweep.py @@ -0,0 +1,1209 @@ +#!/usr/bin/env python3 +""" +Fault sensitivity sweep for DeepFlow parallelism configurations. + +Generates a representative set of parallelism tuples that map to a fixed GPU +budget, evaluates baseline runtimes, and then injects random soft faults across +network dimensions to gauge runtime variability. +""" + +import argparse +import atexit +import datetime +import copy +import json +import math +import os +import random +import signal +import shlex +import sys +import threading +from concurrent.futures import ProcessPoolExecutor, as_completed +from pathlib import Path +from typing import Dict, Iterable, List, Optional, Sequence, Tuple + +import matplotlib.pyplot as plt +import numpy as np +import yaml +from tqdm import tqdm + +try: # Optional styling dependency + import seaborn as sns # type: ignore +except ImportError: # pragma: no cover + sns = None + +sys.path.append(os.path.dirname(os.path.dirname(os.path.abspath(__file__)))) + +import config # noqa: E402 + +from parallelism_sweep import ( # noqa: E402 + HARDWARE_CONFIG_PATH, + MODEL_CONFIG_PATH, + ASTRA_CACHE_MODE, + determine_model_mode, + evaluate_parallelism, + read_yaml, + set_astrasim_cache_mode, + tp_cp_product_is_power_of_two_square, +) + + +class FlowSeq(list): + """List subclass that forces YAML flow-style serialization.""" + + +def _represent_flow_seq(dumper, data): + return dumper.represent_sequence("tag:yaml.org,2002:seq", data, flow_style=True) + + +yaml.SafeDumper.add_representer(FlowSeq, _represent_flow_seq) + + +def _quantize_weight(value: float, decimals: int = 2) -> float: + return round(float(value), decimals) + +# ----------------------------------------------------------------------------- +# Fault sweep configuration +# ----------------------------------------------------------------------------- + +TARGET_NUM_GPUS = 64 +SAMPLE_COUNT = 10 +FAULT_ITER = 25 +FAULT_WORKERS = 100 +FAULT_MAG = [0.5, 0.1] # May also be a (mean, std) tuple +NUM_FAULTS = [1] + +MIN_ALLOWED_TP = 1 +MAX_ALLOWED_TP = 128 +MIN_ALLOWED_CP = 1 +MAX_ALLOWED_CP = 128 +MIN_ALLOWED_DP = 1 +MAX_ALLOWED_DP = 128 +MIN_ALLOWED_LP = 1 +MAX_ALLOWED_LP = 128 +ALLOWED_FAULT_DIMS: Optional[List[int]] = None + +PLOT_OUTPUT_PATH = "tools/fault_sweep.png" +REPORT_OUTPUT_PATH = "tools/fault_sweep.tsv" +FAULT_OUTPUT_DIR = os.path.join(os.getcwd(), "astra_cache_faults") +RESULTS_JSON_PATH = Path("tools/results.json") + +RANDOM_SEED = 1337 +RESULTS_WRITE_LOCK = threading.Lock() +_ACTIVE_EXECUTOR: Optional[ProcessPoolExecutor] = None +_OLD_SIGINT_HANDLER = None +_OLD_SIGTERM_HANDLER = None + +# When True, discard configurations whose tp*cp product is not a square power of two. +ENFORCE_SQUARE_TP_CP = True + +# Debug dumping controls +DEBUG_MODE = False +DEBUG_OUTPUT_BASE = Path("dbg_hw_conf") +DEBUG_RUN_DIR: Optional[Path] = None +DEBUG_SAVED_PATHS: List[Path] = [] +_DEBUG_FILE_INDEX = 0 + +# Failure dump controls +FAILURE_OUTPUT_BASE = DEBUG_OUTPUT_BASE / "failures" +FAILURE_DUMP_DIR: Optional[Path] = None +FAILURE_SAVED_PATHS: List[Path] = [] +_FAILURE_FILE_INDEX = 0 + + +def _parse_int_list(value: Optional[str]) -> List[int]: + if not value: + return [] + parts = [item.strip() for item in value.split(',')] + result: List[int] = [] + for part in parts: + if not part: + continue + result.append(int(part)) + return result + + +def _parse_float_tuple(value: Optional[str]) -> Tuple[float, float]: + if not value: + return (0.0, 0.0) + parts = [item.strip() for item in value.split(',') if item.strip()] + if len(parts) == 1: + val = float(parts[0]) + return (val, 0.0) + if len(parts) >= 2: + return (float(parts[0]), float(parts[1])) + return (0.0, 0.0) + +# Network topologies in AstraSim that understand per-link fault annotations. +# The neighbor builders below encode each topology's link structure so that +# fault injection only targets physically valid edges. +TOPOLOGIES_WITH_FAULTY_LINK_SUPPORT = { + "ring", + "torus2d", + "mesh", + "mesh2d", + "kingmesh2d", +} + +def _build_ring_neighbors(node_count: int) -> List[List[int]]: + if node_count <= 0: + return [] + if node_count == 1: + return [[]] + neighbors: List[List[int]] = [] + for idx in range(node_count): + left = (idx - 1) % node_count + right = (idx + 1) % node_count + entries: List[int] = [] + for candidate in (left, right): + if candidate != idx and candidate not in entries: + entries.append(candidate) + neighbors.append(entries) + return neighbors + + +def _build_mesh_neighbors(node_count: int) -> List[List[int]]: + if node_count <= 0: + return [] + neighbors: List[List[int]] = [] + for idx in range(node_count): + entries: List[int] = [] + if idx - 1 >= 0: + entries.append(idx - 1) + if idx + 1 < node_count: + entries.append(idx + 1) + neighbors.append(entries) + return neighbors + + +def _build_mesh2d_neighbors(node_count: int) -> List[List[int]]: + if node_count <= 0: + return [] + dim = int(math.isqrt(node_count)) + if dim * dim != node_count: + raise ValueError(f"Mesh2D topology requires perfect square node count, got {node_count}.") + neighbors: List[List[int]] = [[] for _ in range(node_count)] + for idx in range(node_count): + row, col = divmod(idx, dim) + entries: List[int] = [] + if col - 1 >= 0: + entries.append(row * dim + (col - 1)) + if col + 1 < dim: + entries.append(row * dim + (col + 1)) + if row - 1 >= 0: + entries.append((row - 1) * dim + col) + if row + 1 < dim: + entries.append((row + 1) * dim + col) + neighbors[idx] = entries + return neighbors + + +def _build_torus2d_neighbors(node_count: int) -> List[List[int]]: + if node_count <= 0: + return [] + dim = int(math.isqrt(node_count)) + if dim * dim != node_count: + raise ValueError(f"Torus2D topology requires perfect square node count, got {node_count}.") + neighbors: List[List[int]] = [[] for _ in range(node_count)] + if dim == 1: + return neighbors + for idx in range(node_count): + row, col = divmod(idx, dim) + left = row * dim + ((col - 1) % dim) + right = row * dim + ((col + 1) % dim) + up = ((row - 1) % dim) * dim + col + down = ((row + 1) % dim) * dim + col + entries = [] + for candidate in (left, right, up, down): + if candidate != idx and candidate not in entries: + entries.append(candidate) + neighbors[idx] = entries + return neighbors + + +def _build_kingmesh2d_neighbors(node_count: int) -> List[List[int]]: + if node_count <= 0: + return [] + dim = int(math.isqrt(node_count)) + if dim * dim != node_count: + raise ValueError(f"KingMesh2D topology requires perfect square node count, got {node_count}.") + neighbors: List[List[int]] = [[] for _ in range(node_count)] + for idx in range(node_count): + row, col = divmod(idx, dim) + entries: List[int] = [] + for dy in (-1, 0, 1): + for dx in (-1, 0, 1): + if dx == 0 and dy == 0: + continue + nr = row + dy + nc = col + dx + if 0 <= nr < dim and 0 <= nc < dim: + candidate = nr * dim + nc + if candidate != idx: + entries.append(candidate) + neighbors[idx] = entries + return neighbors + + +_NEIGHBOR_BUILDERS = { + "ring": _build_ring_neighbors, + "mesh": _build_mesh_neighbors, + "mesh2d": _build_mesh2d_neighbors, + "torus2d": _build_torus2d_neighbors, + "kingmesh2d": _build_kingmesh2d_neighbors, +} + + +_NEIGHBOR_CACHE: Dict[Tuple[str, int], Tuple[Tuple[int, ...], ...]] = {} + + +def _neighbors_for_topology(topo_type: str, node_count: int) -> Tuple[Tuple[int, ...], ...]: + key = (topo_type, int(node_count)) + cached = _NEIGHBOR_CACHE.get(key) + if cached is not None: + return cached + builder = _NEIGHBOR_BUILDERS.get(topo_type) + if builder is None: + raise ValueError(f"No neighbor builder registered for topology '{topo_type}'.") + neighbor_lists = builder(int(node_count)) + normalized = tuple(tuple(neigh) for neigh in neighbor_lists) + _NEIGHBOR_CACHE[key] = normalized + return normalized + + +def _candidate_fault_dimensions(settings: Dict[str, int]) -> List[Tuple[int, int]]: + candidates: List[Tuple[int, int]] = [] + dim0_nodes = int(settings.get("tp", 1)) * int(settings.get("cp", 1)) + if dim0_nodes >= 2 and (ALLOWED_FAULT_DIMS is None or 0 in ALLOWED_FAULT_DIMS): + candidates.append((0, dim0_nodes)) + dim1_nodes = int(settings.get("lp", 1)) * int(settings.get("dp", 1)) + if dim1_nodes >= 2 and (ALLOWED_FAULT_DIMS is None or 1 in ALLOWED_FAULT_DIMS): + candidates.append((1, dim1_nodes)) + return candidates + + +def parse_args(argv: Optional[Sequence[str]] = None) -> argparse.Namespace: + parser = argparse.ArgumentParser(description="DeepFlow fault sensitivity sweep") + parser.add_argument("--target-num-gpus", type=int, default=TARGET_NUM_GPUS) + parser.add_argument("--sample-count", type=int, default=SAMPLE_COUNT) + parser.add_argument("--fault-iter", type=int, default=FAULT_ITER) + parser.add_argument("--fault-workers", type=int, default=FAULT_WORKERS) + parser.add_argument("--fault-mag", type=str, default=f"{FAULT_MAG[0]},{FAULT_MAG[1] if isinstance(FAULT_MAG, (tuple, list)) else 0.0}", help="Mean[,std] for fault magnitude") + parser.add_argument("--num-faults", type=str, default=",".join(str(n) for n in NUM_FAULTS), help="Comma-separated list of number of faults per evaluation") + parser.add_argument("--min-tp", type=int, default=MIN_ALLOWED_TP) + parser.add_argument("--max-tp", type=int, default=MAX_ALLOWED_TP) + parser.add_argument("--min-cp", type=int, default=MIN_ALLOWED_CP) + parser.add_argument("--max-cp", type=int, default=MAX_ALLOWED_CP) + parser.add_argument("--min-dp", type=int, default=MIN_ALLOWED_DP) + parser.add_argument("--max-dp", type=int, default=MAX_ALLOWED_DP) + parser.add_argument("--min-lp", type=int, default=MIN_ALLOWED_LP) + parser.add_argument("--max-lp", type=int, default=MAX_ALLOWED_LP) + parser.add_argument("--allowed-fault-dims", type=str, default=None, help="Comma-separated list of network dimension indices eligible for faults") + parser.add_argument("--plot-output", type=str, default=PLOT_OUTPUT_PATH) + parser.add_argument("--report-output", type=str, default=REPORT_OUTPUT_PATH) + parser.add_argument("--results-json", type=str, default=str(RESULTS_JSON_PATH)) + parser.add_argument("--seed", type=int, default=RANDOM_SEED) + parser.add_argument( + "--enforce-square-tp-cp", + action="store_true", + help="Require tp*cp to be a square power of two when sampling parallelism tuples.", + ) + parser.add_argument( + "--debug-mode", + action="store_true", + help="Enable debug dumps of hardware configs and exit after generation.", + ) + return parser.parse_args(argv) + + +# ----------------------------------------------------------------------------- +# Parallelism sampling helpers +# ----------------------------------------------------------------------------- + +def _divisors(n: int) -> List[int]: + divs = set() + for i in range(1, int(math.sqrt(n)) + 1): + if n % i == 0: + divs.add(i) + divs.add(n // i) + return sorted(divs) + + +def _enumerate_parallelism_vectors(num_gpus: int) -> List[Tuple[int, int, int, int]]: + vectors: List[Tuple[int, int, int, int]] = [] + for tp in _divisors(num_gpus): + rem_tp = num_gpus // tp + for cp in _divisors(rem_tp): + rem_cp = rem_tp // cp + for dp in _divisors(rem_cp): + lp = rem_cp // dp + vectors.append((tp, cp, dp, lp)) + return vectors + + +def _log_distance(a: Sequence[int], b: Sequence[int]) -> float: + total = 0.0 + for x, y in zip(a, b): + total += abs(math.log2(max(x, 1)) - math.log2(max(y, 1))) + return total + + +def _select_representative_vectors(candidates: List[Tuple[int, int, int, int]], count: int) -> List[Tuple[int, int, int, int]]: + if not candidates: + return [] + selected: List[Tuple[int, int, int, int]] = [] + + # Seed with extreme configurations + extremes = [ + max(candidates, key=lambda v: v[0]), + max(candidates, key=lambda v: v[1]), + max(candidates, key=lambda v: v[2]), + max(candidates, key=lambda v: v[3]), + ] + for cand in extremes: + if cand not in selected: + selected.append(cand) + if len(selected) >= count: + return selected[:count] + + # Farthest-point sampling to promote coverage + remaining = [v for v in candidates if v not in selected] + while remaining and len(selected) < count: + best_vec = None + best_score = -1.0 + for vec in remaining: + min_dist = min(_log_distance(vec, sel) for sel in selected) + if min_dist > best_score: + best_score = min_dist + best_vec = vec + if best_vec is None: + break + selected.append(best_vec) + remaining.remove(best_vec) + + # If still short, append random leftovers + while remaining and len(selected) < count: + candidate = remaining.pop(random.randrange(len(remaining))) + selected.append(candidate) + + return selected[:count] + + +def generate_parallelism_samples(num_gpus: int, sample_count: int) -> List[Dict[str, int]]: + tuples = _enumerate_parallelism_vectors(num_gpus) + filtered: List[Tuple[int, int, int, int]] = [] + for tp, cp, dp, lp in tuples: + if not (MIN_ALLOWED_TP <= tp <= MAX_ALLOWED_TP): + continue + if not (MIN_ALLOWED_CP <= cp <= MAX_ALLOWED_CP): + continue + if not (MIN_ALLOWED_DP <= dp <= MAX_ALLOWED_DP): + continue + if not (MIN_ALLOWED_LP <= lp <= MAX_ALLOWED_LP): + continue + if ENFORCE_SQUARE_TP_CP and not tp_cp_product_is_power_of_two_square(tp, cp): + continue + filtered.append((tp, cp, dp, lp)) + tuples = filtered + chosen = _select_representative_vectors(tuples, sample_count) + samples: List[Dict[str, int]] = [] + for tp, cp, dp, lp in chosen: + settings = { + "tp": tp, + "cp": cp, + "dp": dp, + "lp": lp, + "mb": max(lp, 1), + "tp_sp": True, + } + samples.append(settings) + return samples + + +# ----------------------------------------------------------------------------- +# Fault injection helpers +# ----------------------------------------------------------------------------- + +def sample_fault_magnitude() -> float: + if isinstance(FAULT_MAG, (tuple, list)) and len(FAULT_MAG) == 2: + mean, std = FAULT_MAG + value = random.gauss(mean, std) + else: + value = float(FAULT_MAG) + value = max(0.0, min(0.95, value)) + value = _quantize_weight(value) + return max(0.1, value) # min of 10% bw, max of 95% bw + + +def make_fault_mutator(fault_specs: Sequence[Tuple[int, float, int]]): + fault_specs = [ + (int(dim_idx), float(weight), int(node_count)) + for dim_idx, weight, node_count in fault_specs + ] + + def _mutator(hw_dict: Dict[str, object]) -> None: + network = hw_dict.get("network") + if not isinstance(network, dict): + return + dimensions = network.get("dimensions") + if not isinstance(dimensions, list): + return + fault_entries: List[FlowSeq] = [] + for dimension_index, weight, node_count in fault_specs: + if not (0 <= dimension_index < len(dimensions)): + continue + dim = dimensions[dimension_index] + if not isinstance(dim, dict): + continue + topology = dim.get("topology") + topo_type = "" + if isinstance(topology, dict): + topo_raw = topology.get("type", "") + if isinstance(topo_raw, str): + topo_type = topo_raw.strip().lower() + supports_faulty_links = topo_type in TOPOLOGIES_WITH_FAULTY_LINK_SUPPORT + if not supports_faulty_links: + label = dim.get("label", f"dim{dimension_index}") + raise ValueError( + f"Topology '{topo_type or 'unknown'}' on network dimension '{label}' " + "does not support faulty_links injection." + ) + if node_count < 2: + continue + try: + neighbor_map = _neighbors_for_topology(topo_type, node_count) + except ValueError as exc: + raise ValueError( + f"Unable to determine neighbors for topology '{topo_type or 'unknown'}' " + f"with node_count={node_count}: {exc}" + ) from exc + valid_sources = [idx for idx, neigh in enumerate(neighbor_map) if neigh] + if not valid_sources: + continue + src = random.choice(valid_sources) + dst_candidates = neighbor_map[src] + dst = random.choice(dst_candidates) + weight_clamped = float(max(0.0, min(1.0, weight))) + weight_clamped = _quantize_weight(weight_clamped) + dim.pop("faulty_links", None) + fault_entries.append(FlowSeq([int(src), int(dst), weight_clamped])) + + if fault_entries: + network["faulty_links"] = FlowSeq(fault_entries) + else: + network.pop("faulty_links", None) + + return _mutator + + +# ----------------------------------------------------------------------------- +# Multiprocessing helpers +# ----------------------------------------------------------------------------- + +_WORKER_HW_DICT = None +_WORKER_MODEL_CONFIG = None +_WORKER_MODE = None + + +def _parallelism_worker_init(base_hw_dict, model_config_path, mode): + global _WORKER_HW_DICT, _WORKER_MODEL_CONFIG, _WORKER_MODE + set_astrasim_cache_mode(ASTRA_CACHE_MODE) + _WORKER_HW_DICT = base_hw_dict + _WORKER_MODEL_CONFIG = config.parse_config(model_config_path, config_type=mode) + _WORKER_MODE = mode + + +def _execute_parallelism_task(base_hw_dict, model_config_obj, mode, task: Dict[str, object]) -> Dict[str, object]: + kind = task.get("kind") + settings = dict(task.get("settings", {})) + key = task.get("key") + settings_key = tuple(task.get("settings_key", ())) + task_id = task.get("task_id") + try: + mutator = None + if kind == "fault": + raw_faults = task.get("faults", []) + faults = [tuple(spec) for spec in raw_faults] + mutator = make_fault_mutator(faults) + else: + faults = [] + metrics = evaluate_parallelism( + base_hw_dict, + model_config_obj, + mode, + settings, + hw_mutator=mutator, + ) + result: Dict[str, object] = { + "status": "ok", + "kind": kind, + "key": key, + "settings": settings, + "settings_key": settings_key, + "runtime": float(metrics["runtime"]), + "hw_yaml": metrics.get("hw_yaml"), + "task_id": task_id, + } + if kind == "fault": + result.update( + { + "faults": faults, + "num_faults": int(task.get("num_faults", len(faults))), + } + ) + else: + result.update({"faults": [], "num_faults": 0}) + return result + except Exception as exc: + return { + "status": "error", + "kind": kind, + "key": key, + "settings": settings, + "settings_key": settings_key, + "task_id": task_id, + "error": str(exc), + } + + +def _parallelism_worker_task(task: Dict[str, object]) -> Dict[str, object]: + if _WORKER_HW_DICT is None or _WORKER_MODEL_CONFIG is None or _WORKER_MODE is None: + raise RuntimeError("Worker initialisation missing before task execution.") + return _execute_parallelism_task(_WORKER_HW_DICT, _WORKER_MODEL_CONFIG, _WORKER_MODE, task) + + +def _parallelism_key(settings: Dict[str, object]) -> Tuple[Tuple[str, object], ...]: + return tuple(sorted(settings.items())) + + +def _serialize_settings_key(key: Tuple[Tuple[str, object], ...]) -> List[List[object]]: + return [[k, v] for k, v in key] + + +def _deserialize_settings_key(raw: Iterable[Iterable[object]]) -> Tuple[Tuple[str, object], ...]: + return tuple((str(k), v) for k, v in raw) + + +def _normalise_faults_raw(faults: Iterable[Iterable[object]]) -> List[List[object]]: + normalised: List[List[object]] = [] + for dim, weight, nodes in faults: + normalised.append([int(dim), _quantize_weight(float(weight)), int(nodes)]) + return normalised + + +def _task_identifier( + kind: str, + settings_key: Tuple[Tuple[str, object], ...], + *, + faults: Sequence[Tuple[int, float, int]] | None = None, + num_faults: int | None = None, +) -> str: + payload: Dict[str, object] = { + "kind": kind, + "settings_key": _serialize_settings_key(settings_key), + } + if kind == "fault": + effective_faults = faults or () + payload["num_faults"] = int(num_faults if num_faults is not None else len(effective_faults)) + payload["faults"] = _normalise_faults_raw(effective_faults) + return json.dumps(payload, sort_keys=True) + + +def _load_existing_results(path: Path) -> Dict[str, Dict[str, object]]: + if not path.exists(): + return {} + try: + with path.open("r", encoding="utf-8") as handle: + data = json.load(handle) + if isinstance(data, dict): + return data + print(f"Warning: results file {path} did not contain a JSON object. Ignoring existing data.", file=sys.stderr) + except Exception as exc: + print(f"Warning: failed to load results file {path}: {exc}", file=sys.stderr) + return {} + + +def _record_partial_result(results_store: Dict[str, Dict[str, object]], entry: Dict[str, object]) -> None: + task_id = entry.get("task_id") + if not task_id: + return + results_store[task_id] = { + "kind": entry.get("kind"), + "settings_key": _serialize_settings_key(tuple(entry.get("settings_key", ()))), + "runtime": entry.get("runtime"), + } + if entry.get("kind") == "fault": + results_store[task_id]["num_faults"] = int(entry.get("num_faults", 0)) + results_store[task_id]["faults"] = _normalise_faults_raw(entry.get("faults", [])) + tmp_path = RESULTS_JSON_PATH.with_suffix(".tmp") + tmp_path.parent.mkdir(parents=True, exist_ok=True) + with RESULTS_WRITE_LOCK: + with tmp_path.open("w", encoding="utf-8") as handle: + json.dump(results_store, handle, indent=2) + tmp_path.replace(RESULTS_JSON_PATH) + + +def _dump_debug_hw_config(result: Dict[str, object]) -> Optional[Path]: + if not DEBUG_MODE: + return None + yaml_text = result.get("hw_yaml") + if not isinstance(yaml_text, str) or not yaml_text.strip(): + return None + if DEBUG_RUN_DIR is None: + return None + global _DEBUG_FILE_INDEX + _DEBUG_FILE_INDEX += 1 + settings = result.get("settings") or {} + kind = str(result.get("kind") or "task") + tp = settings.get("tp") + cp = settings.get("cp") + dp = settings.get("dp") + lp = settings.get("lp") + label = f"{kind}_{_DEBUG_FILE_INDEX:03d}" + if all(isinstance(val, (int, float)) for val in (tp, cp, dp, lp)): + label += f"_tp{int(tp)}_cp{int(cp)}_dp{int(dp)}_lp{int(lp)}" + num_faults = result.get("num_faults") + if kind == "fault" and isinstance(num_faults, int): + label += f"_{num_faults}faults" + file_path = DEBUG_RUN_DIR / f"{label}.yaml" + file_path.parent.mkdir(parents=True, exist_ok=True) + with file_path.open("w", encoding="utf-8") as handle: + handle.write(yaml_text) + DEBUG_SAVED_PATHS.append(file_path) + return file_path + + +def _generate_hw_yaml_from_task(base_hw_dict: Dict[str, object], task: Dict[str, object]) -> Optional[str]: + if base_hw_dict is None: + return None + try: + hw_copy: Dict[str, object] = copy.deepcopy(base_hw_dict) + except Exception as exc: + print(f"Warning: failed to clone base hardware config: {exc}", file=sys.stderr) + return None + + settings = dict(task.get("settings", {}) or {}) + parallel_block = hw_copy.setdefault("parallelism", {}) + if isinstance(parallel_block, dict): + for key, value in settings.items(): + parallel_block[key] = value + + mutator = None + faults: Sequence[Tuple[int, float, int]] = () + if task.get("kind") == "fault": + raw_faults = task.get("faults") or [] + try: + faults = [tuple(spec) for spec in raw_faults] # type: ignore[assignment] + except Exception: + faults = () + mutator = make_fault_mutator(faults) + + try: + if mutator is not None: + mutator(hw_copy) + except Exception as exc: + print(f"Warning: failed to apply fault mutator for debug dump: {exc}", file=sys.stderr) + + try: + return yaml.safe_dump(hw_copy, default_flow_style=False, sort_keys=False) + except Exception as exc: + print(f"Warning: failed to serialize debug hardware config: {exc}", file=sys.stderr) + return None + + +def _ensure_failure_dump_dir() -> Path: + global FAILURE_DUMP_DIR + if FAILURE_DUMP_DIR is None: + timestamp = datetime.datetime.now().strftime("%Y%m%d_%H%M%S") + FAILURE_DUMP_DIR = FAILURE_OUTPUT_BASE / timestamp + FAILURE_DUMP_DIR.mkdir(parents=True, exist_ok=True) + print(f"Saving failing hardware configs to {FAILURE_DUMP_DIR}") + return FAILURE_DUMP_DIR + + +def _failure_label_from_settings(kind: str, settings: Dict[str, object], num_faults: Optional[int]) -> str: + label = kind or "task" + tp = settings.get("tp") + cp = settings.get("cp") + dp = settings.get("dp") + lp = settings.get("lp") + if all(isinstance(val, (int, float)) for val in (tp, cp, dp, lp)): + label += f"_tp{int(tp)}_cp{int(cp)}_dp{int(dp)}_lp{int(lp)}" + if kind == "fault" and isinstance(num_faults, int): + label += f"_{num_faults}faults" + return label + + +def _dump_failure_hw_config( + base_hw_dict: Dict[str, object], + result: Dict[str, object], + error_message: Optional[str] = None, +) -> Optional[Path]: + task_like = { + "kind": result.get("kind"), + "settings": result.get("settings"), + "faults": result.get("faults"), + "num_faults": result.get("num_faults"), + } + yaml_text = _generate_hw_yaml_from_task(base_hw_dict, task_like) + failure_dir = _ensure_failure_dump_dir() + global _FAILURE_FILE_INDEX + _FAILURE_FILE_INDEX += 1 + settings = dict(result.get("settings", {}) or {}) + kind = str(result.get("kind") or "task") + num_faults = result.get("num_faults") + label = _failure_label_from_settings(kind, settings, num_faults if isinstance(num_faults, int) else None) + base_filename = f"failure_{_FAILURE_FILE_INDEX:03d}_{label}" + yaml_path = failure_dir / f"{base_filename}.yaml" + command_line = None + if isinstance(yaml_text, str) and yaml_text.strip(): + with yaml_path.open("w", encoding="utf-8") as handle: + handle.write(yaml_text) + FAILURE_SAVED_PATHS.append(yaml_path) + print(f"Stored failing hardware config at {yaml_path}") + cli_hw_path = os.path.relpath(yaml_path, os.getcwd()) + command_line = ( + f"uv run run_perf.py --hardware_config {shlex.quote(cli_hw_path)} " + f"--model_config {shlex.quote(MODEL_CONFIG_PATH)}" + ) + else: + yaml_path = None + if error_message or command_line: + log_path = failure_dir / f"{base_filename}_error.txt" + with log_path.open("w", encoding="utf-8") as handle: + if error_message: + handle.write(str(error_message).strip() + "\n") + if command_line: + handle.write("\n# Reproduce failing configuration\n") + handle.write(command_line + "\n") + FAILURE_SAVED_PATHS.append(log_path) + return yaml_path + + +def _terminate_active_executor() -> None: + global _ACTIVE_EXECUTOR + executor = _ACTIVE_EXECUTOR + if executor is None: + return + _ACTIVE_EXECUTOR = None + try: + executor.shutdown(wait=False, cancel_futures=True) + except Exception: + pass + processes = getattr(executor, "_processes", None) + if isinstance(processes, dict): + for proc in processes.values(): + try: + if proc is not None and proc.is_alive(): + proc.kill() + except Exception: + continue + + +def _signal_handler(signum, frame): # type: ignore[override] + _terminate_active_executor() + if signum == signal.SIGINT: + raise KeyboardInterrupt() + sys.exit(128 + signum) + + +def _install_signal_handlers() -> None: + global _OLD_SIGINT_HANDLER, _OLD_SIGTERM_HANDLER + _OLD_SIGINT_HANDLER = signal.getsignal(signal.SIGINT) + _OLD_SIGTERM_HANDLER = signal.getsignal(signal.SIGTERM) + signal.signal(signal.SIGINT, _signal_handler) + signal.signal(signal.SIGTERM, _signal_handler) + atexit.register(_terminate_active_executor) + + +# ----------------------------------------------------------------------------- +# Reporting and plotting +# ----------------------------------------------------------------------------- + +def summarise_fault_runs(runtimes: Iterable[float]) -> Tuple[float, float, float]: + values = [float(v) for v in runtimes if math.isfinite(v)] + if not values: + return float("nan"), float("nan"), float("nan") + return min(values), max(values), float(sum(values) / len(values)) + + +def write_fault_report(records: List[Dict[str, object]], path: str) -> None: + os.makedirs(os.path.dirname(path), exist_ok=True) + header = [ + "tp", + "cp", + "dp", + "lp", + "baseline_runtime", + "fault_min", + "fault_max", + "fault_mean", + ] + with open(path, "w") as handle: + handle.write("\t".join(header) + "\n") + for record in records: + row = [ + str(record["parallelism"]["tp"]), + str(record["parallelism"]["cp"]), + str(record["parallelism"]["dp"]), + str(record["parallelism"]["lp"]), + f"{record['baseline_runtime']:.6f}", + f"{record['fault_min']:.6f}" if math.isfinite(record["fault_min"]) else "nan", + f"{record['fault_max']:.6f}" if math.isfinite(record["fault_max"]) else "nan", + f"{record['fault_mean']:.6f}" if math.isfinite(record["fault_mean"]) else "nan", + ] + handle.write("\t".join(row) + "\n") + + +def plot_fault_sensitivity(records: List[Dict[str, object]], output_path: str, num_gpus: int) -> None: + if not records: + print("No records to plot.", file=sys.stderr) + return + + if sns is not None: # pragma: no branch + sns.set_theme(style="whitegrid") + + xs = np.arange(len(records)) + baseline = np.array([rec["baseline_runtime"] for rec in records], dtype=float) + fault_mean = np.array([rec["fault_mean"] for rec in records], dtype=float) + fault_min = np.array([rec["fault_min"] for rec in records], dtype=float) + fault_max = np.array([rec["fault_max"] for rec in records], dtype=float) + + config_labels = [rec["label"] for rec in records] + + plt.figure(figsize=(12, 6)) + plt.scatter(xs, baseline, color="#1f78b4", marker="o", s=80, label="Baseline") + + fault_lower = [] + fault_upper = [] + for mean_val, min_val, max_val in zip(fault_mean, fault_min, fault_max): + if math.isfinite(mean_val) and math.isfinite(min_val): + fault_lower.append(max(mean_val - min_val, 0.0)) + else: + fault_lower.append(0.0) + if math.isfinite(mean_val) and math.isfinite(max_val): + fault_upper.append(max(max_val - mean_val, 0.0)) + else: + fault_upper.append(0.0) + + plt.errorbar( + xs, + fault_mean, + yerr=[fault_lower, fault_upper], + fmt="s", + color="#d7301f", + ecolor="#fb6a4a", + elinewidth=2, + capsize=6, + capthick=1.6, + markersize=7, + label="Fault mean ± range", + ) + + # Candlestick-style shading for min/max bounds + for x, fmin, fmax in zip(xs, fault_min, fault_max): + if math.isfinite(fmin) and math.isfinite(fmax): + plt.fill_between([x - 0.22, x + 0.22], [fmin, fmin], [fmax, fmax], color="#fdd0a2", alpha=0.4) + + plt.xticks(xs, config_labels, rotation=45, ha="right") + plt.ylabel("Runtime (s)") + plt.title(f"Fault Sensitivity Across Parallelism Configurations (Num GPUs = {num_gpus})") + plt.grid(alpha=0.3, axis="y") + plt.legend() + plt.tight_layout() + + os.makedirs(os.path.dirname(output_path), exist_ok=True) + plt.savefig(output_path, dpi=200) + plt.close() + print(f"Saved fault sensitivity plot to {output_path}") + + +# ----------------------------------------------------------------------------- +# Main execution +# ----------------------------------------------------------------------------- + +def main(cli_args: Optional[Sequence[str]] = None) -> int: + global TARGET_NUM_GPUS, SAMPLE_COUNT, FAULT_ITER, FAULT_WORKERS, FAULT_MAG + global NUM_FAULTS, MIN_ALLOWED_TP, MAX_ALLOWED_TP, MIN_ALLOWED_CP, MAX_ALLOWED_CP + global MIN_ALLOWED_DP, MAX_ALLOWED_DP, MIN_ALLOWED_LP, MAX_ALLOWED_LP + global ALLOWED_FAULT_DIMS, PLOT_OUTPUT_PATH, REPORT_OUTPUT_PATH, RESULTS_JSON_PATH + global RANDOM_SEED, ENFORCE_SQUARE_TP_CP, DEBUG_MODE, DEBUG_RUN_DIR, DEBUG_SAVED_PATHS + global _DEBUG_FILE_INDEX, FAILURE_DUMP_DIR, FAILURE_SAVED_PATHS, _FAILURE_FILE_INDEX + + args = parse_args(cli_args) + + TARGET_NUM_GPUS = args.target_num_gpus + SAMPLE_COUNT = args.sample_count + FAULT_ITER = args.fault_iter + FAULT_WORKERS = args.fault_workers + + mag_mean, mag_std = _parse_float_tuple(args.fault_mag) + FAULT_MAG = (mag_mean, mag_std) + + parsed_num_faults = _parse_int_list(args.num_faults) + NUM_FAULTS = parsed_num_faults or [1] + + MIN_ALLOWED_TP = max(1, args.min_tp) + MAX_ALLOWED_TP = max(MIN_ALLOWED_TP, args.max_tp) + MIN_ALLOWED_CP = max(1, args.min_cp) + MAX_ALLOWED_CP = max(MIN_ALLOWED_CP, args.max_cp) + MIN_ALLOWED_DP = max(1, args.min_dp) + MAX_ALLOWED_DP = max(MIN_ALLOWED_DP, args.max_dp) + MIN_ALLOWED_LP = max(1, args.min_lp) + MAX_ALLOWED_LP = max(MIN_ALLOWED_LP, args.max_lp) + + allowed_dims_list = _parse_int_list(args.allowed_fault_dims) if args.allowed_fault_dims else [] + ALLOWED_FAULT_DIMS = allowed_dims_list if allowed_dims_list else None + + PLOT_OUTPUT_PATH = args.plot_output + REPORT_OUTPUT_PATH = args.report_output + RESULTS_JSON_PATH = Path(args.results_json) + + RANDOM_SEED = args.seed + random.seed(RANDOM_SEED) + ENFORCE_SQUARE_TP_CP = ENFORCE_SQUARE_TP_CP or bool(args.enforce_square_tp_cp) + DEBUG_MODE = bool(args.debug_mode) + _DEBUG_FILE_INDEX = 0 + _FAILURE_FILE_INDEX = 0 + FAILURE_SAVED_PATHS.clear() + FAILURE_DUMP_DIR = None + if DEBUG_MODE: + FAULT_WORKERS = 1 + DEBUG_SAVED_PATHS.clear() + timestamp = datetime.datetime.now().strftime("%Y%m%d_%H%M%S") + DEBUG_RUN_DIR = DEBUG_OUTPUT_BASE / timestamp + DEBUG_RUN_DIR.mkdir(parents=True, exist_ok=True) + print(f"Debug mode enabled. Hardware configs will be written to {DEBUG_RUN_DIR}") + else: + DEBUG_RUN_DIR = None + DEBUG_SAVED_PATHS.clear() + + set_astrasim_cache_mode(ASTRA_CACHE_MODE) + _install_signal_handlers() + + base_hw_dict = read_yaml(HARDWARE_CONFIG_PATH) + mode = determine_model_mode(MODEL_CONFIG_PATH) + model_config_obj = config.parse_config(MODEL_CONFIG_PATH, config_type=mode) + + existing_results = _load_existing_results(RESULTS_JSON_PATH) + if existing_results: + print(f"Loaded {len(existing_results)} cached result entries from {RESULTS_JSON_PATH}") + results_store: Dict[str, Dict[str, object]] = dict(existing_results) + + parallelism_samples = generate_parallelism_samples(TARGET_NUM_GPUS, SAMPLE_COUNT) + if not parallelism_samples: + print("Unable to generate parallelism samples.", file=sys.stderr) + return + + records: List[Dict[str, object]] = [] + records_by_key: Dict[Tuple[Tuple[str, object], ...], Dict[str, object]] = {} + task_queue: List[Dict[str, object]] = [] + total_fault_tasks = 0 + + for settings in parallelism_samples: + settings_copy = dict(settings) + key = _parallelism_key(settings_copy) + settings_key = key + + candidates_for_config = _candidate_fault_dimensions(settings_copy) + if not candidates_for_config: + continue + + records_by_key[key] = { + "parallelism": settings_copy, + "baseline_runtime": None, + "fault_runtimes": [], + "fault_details": [], + } + + baseline_task_id = _task_identifier("baseline", settings_key) + baseline_cached = existing_results.get(baseline_task_id) + if baseline_cached is not None: + runtime = baseline_cached.get("runtime") + if runtime is not None: + records_by_key[key]["baseline_runtime"] = float(runtime) + else: + task_queue.append( + { + "kind": "baseline", + "settings": settings_copy, + "key": key, + "settings_key": settings_key, + "task_id": baseline_task_id, + } + ) + + for num_faults in NUM_FAULTS: + repeats = max(1, int(num_faults)) + for _ in range(FAULT_ITER): + fault_specs: List[Tuple[int, float, int]] = [] + for _ in range(repeats): + fault_value = sample_fault_magnitude() + dim_choice, node_count = random.choice(candidates_for_config) + fault_specs.append((dim_choice, fault_value, node_count)) + + fault_task_id = _task_identifier( + "fault", + settings_key, + faults=fault_specs, + num_faults=repeats, + ) + cached_fault = existing_results.get(fault_task_id) + if cached_fault is not None: + runtime = cached_fault.get("runtime") + if runtime is not None: + records_by_key[key]["fault_runtimes"].append(float(runtime)) + cached_faults = _normalise_faults_raw(cached_fault.get("faults", [])) + records_by_key[key]["fault_details"].append( + { + "num_faults": cached_fault.get("num_faults"), + "faults": cached_faults, + } + ) + continue + total_fault_tasks += 1 + task_queue.append( + { + "kind": "fault", + "settings": settings_copy, + "key": key, + "settings_key": settings_key, + "task_id": fault_task_id, + "faults": fault_specs, + "num_faults": repeats, + } + ) + + if not task_queue: + print("No new evaluation tasks; reusing cached results.") + else: + print( + f"Prepared {len(task_queue)} evaluation task(s) across " + f"{len(parallelism_samples)} configuration(s) " + f"(fault tasks: {total_fault_tasks})." + ) + + if DEBUG_MODE: + if not task_queue: + print("Debug mode enabled but no evaluation tasks were generated.") + else: + for task in task_queue: + yaml_text = _generate_hw_yaml_from_task(base_hw_dict, task) + debug_entry = { + "kind": task.get("kind"), + "settings": task.get("settings"), + "num_faults": task.get("num_faults"), + "faults": task.get("faults"), + "hw_yaml": yaml_text, + } + _dump_debug_hw_config(debug_entry) + if DEBUG_RUN_DIR: + if DEBUG_SAVED_PATHS: + print(f"Debug mode: saved {len(DEBUG_SAVED_PATHS)} hardware config(s) to {DEBUG_RUN_DIR}") + else: + print(f"Debug mode enabled but no hardware configs were generated. Directory: {DEBUG_RUN_DIR}") + else: + print("Debug mode enabled but no output directory was initialised.") + return 0 + + def handle_task_result(result: Dict[str, object]) -> None: + if not isinstance(result, dict): + return + key = result.get("key") + rec = records_by_key.get(key) + if rec is None: + return + status = result.get("status") + if status != "ok": + error_msg = result.get("error") + print( + f"Warning: task failed for settings {result.get('settings')}: {error_msg}", + file=sys.stderr, + ) + _dump_failure_hw_config(base_hw_dict, result, error_msg) + return + kind = result.get("kind") + runtime = float(result.get("runtime", float("nan"))) + if kind == "baseline": + rec["baseline_runtime"] = runtime + else: + rec["fault_runtimes"].append(runtime) + rec["fault_details"].append( + { + "num_faults": result.get("num_faults"), + "faults": _normalise_faults_raw(result.get("faults", [])), + } + ) + _record_partial_result(results_store, result) + _dump_debug_hw_config(result) + + if task_queue: + if FAULT_WORKERS and FAULT_WORKERS > 1: + available_cpus = max(1, os.cpu_count() or 1) + worker_count = min(max(1, FAULT_WORKERS), len(task_queue), available_cpus) + print(f"Launching ProcessPoolExecutor with {worker_count} worker(s).") + executor = ProcessPoolExecutor( + max_workers=worker_count, + initializer=_parallelism_worker_init, + initargs=(base_hw_dict, MODEL_CONFIG_PATH, mode), + ) + try: + global _ACTIVE_EXECUTOR + _ACTIVE_EXECUTOR = executor + futures = [executor.submit(_parallelism_worker_task, task) for task in task_queue] + for fut in tqdm(as_completed(futures), total=len(futures), desc="Evaluations", unit="task"): + try: + result = fut.result() + except Exception as exc: # pragma: no cover - defensive + print(f"Warning: parallel task raised an exception: {exc}", file=sys.stderr) + continue + handle_task_result(result) + finally: + _terminate_active_executor() + else: + print("Running tasks sequentially (FAULT_WORKERS <= 1).") + for task in tqdm(task_queue, desc="Evaluations", unit="task"): + try: + result = _execute_parallelism_task(base_hw_dict, model_config_obj, mode, task) + except Exception as exc: # pragma: no cover - defensive + print(f"Warning: sequential task raised an exception: {exc}", file=sys.stderr) + continue + handle_task_result(result) + + for settings in parallelism_samples: + key = _parallelism_key(settings) + state = records_by_key.get(key) + if not state: + continue + baseline_runtime = state.get("baseline_runtime") + if baseline_runtime is None: + print(f"Warning: missing baseline runtime for settings {settings}", file=sys.stderr) + continue + fault_min, fault_max, fault_mean = summarise_fault_runs(state.get("fault_runtimes", [])) + record = { + "parallelism": state["parallelism"], + "baseline_runtime": float(baseline_runtime), + "fault_min": fault_min, + "fault_max": fault_max, + "fault_mean": fault_mean, + "label": f"tp{settings['tp']}-cp{settings['cp']}-dp{settings['dp']}-lp{settings['lp']}", + } + records.append(record) + + if not records: + print("No successful sweep results collected.", file=sys.stderr) + return 0 + + write_fault_report(records, REPORT_OUTPUT_PATH) + print(f"Wrote fault sweep report to {REPORT_OUTPUT_PATH}") + plot_fault_sensitivity(records, PLOT_OUTPUT_PATH, TARGET_NUM_GPUS) + if FAILURE_SAVED_PATHS and FAILURE_DUMP_DIR: + print( + f"Captured {len(FAILURE_SAVED_PATHS)} failing hardware config artefact(s) in {FAILURE_DUMP_DIR}" + ) + return 0 + + +if __name__ == "__main__": + sys.exit(main()) diff --git a/tools/parallelism_sweep.py b/tools/parallelism_sweep.py index 29afeb84..0637f6a1 100644 --- a/tools/parallelism_sweep.py +++ b/tools/parallelism_sweep.py @@ -47,16 +47,16 @@ # ----------------------------------------------------------------------------- # Paths to the baseline configuration files -HARDWARE_CONFIG_PATH = "configs/hardware-config/a100_80GB.yaml" +HARDWARE_CONFIG_PATH = "configs/hardware-config/a100_80GB_faraz.yaml" MODEL_CONFIG_PATH = "configs/model-config/Llama3.1-405B.yaml" # Parallelism values to sweep (dense grid). Edit to suit your search space. # Keys should match the entries under the YAML "parallelism" section. PARALLELISM_SWEEP = { - "tp": [2**i for i in range(0, 3)], - "cp": [2**i for i in range(0, 3)], - "dp": [2**i for i in range(0, 3)], - "lp": [2**i for i in range(0, 3)], + "tp": [2**i for i in range(0, 9)], + "cp": [2**i for i in range(0, 9)], + "dp": [2**i for i in range(0, 9)], + "lp": [2**i for i in range(0, 9)], } # Optional knobs that still live inside the parallelism section but do not @@ -67,8 +67,11 @@ # GPU count filter: only evaluate combinations whose TP*CP*DP*LP fall inside # this inclusive range. -GPU_COUNT_MIN = 16 -GPU_COUNT_MAX = 64 +GPU_COUNT_MIN = 64 +GPU_COUNT_MAX = 1024 + +# When True, discard configurations whose tp*cp product is not a square power of two. +ENFORCE_SQUARE_TP_CP = False # Select the y-axis metric for the scatter plot and report. Accepted values: # "runtime" -> plot raw runtime in seconds (lower is better) @@ -89,10 +92,10 @@ ASTRA_CACHE_MODE = "NO_CACHE" # Options: NO_CACHE, CACHE_READONLY, CACHE_READWRITE # Plotting behaviour toggles -MEM_AWARE_FILTER = True # When True, skip memory-violating configurations in plots. +MEM_AWARE_FILTER = False # When True, skip memory-violating configurations in plots. # Maximum number of parallel worker processes (set <= available CPUs - 1). Set to 1 to disable multiprocessing. -MAX_WORKERS = 16 +MAX_WORKERS = 100 # ----------------------------------------------------------------------------- @@ -194,19 +197,45 @@ def total_gpu_count(parallel_cfg): return total -def make_temp_hw_config(base_hw_dict, parallel_settings): - """Return parsed HW config for the given parallelism override.""" +def tp_cp_product_is_power_of_two_square(tp_value, cp_value): + """Return True when tp*cp is a square whose root is also a power of two.""" + if tp_value == 1 and cp_value == 1: + return False # special case + try: + tp_int = int(tp_value) + cp_int = int(cp_value) + except (TypeError, ValueError): + return False + if tp_int <= 0 or cp_int <= 0: + return False + product = tp_int * cp_int + root = math.isqrt(product) + if root * root != product: + return False + return (root & (root - 1)) == 0 + + +def make_temp_hw_config(base_hw_dict, parallel_settings, hw_mutator=None): + """Return (parsed HW config, YAML string) for the given override.""" updated = copy.deepcopy(base_hw_dict) parallel_block = updated.setdefault("parallelism", {}) for key, value in parallel_settings.items(): parallel_block[key] = value + if hw_mutator: + hw_mutator(updated) + try: + debug_yaml = yaml.safe_dump(updated, default_flow_style=False) + except Exception: + debug_yaml = None + tmp_file = tempfile.NamedTemporaryFile("w", suffix=".yaml", delete=False) try: yaml.safe_dump(updated, tmp_file, default_flow_style=False, sort_keys=False) tmp_file.flush() tmp_file.close() - return config.parse_config(tmp_file.name, config_type="hardware") + hw_config = config.parse_config(tmp_file.name, config_type="hardware") + return hw_config, debug_yaml finally: try: tmp_file.close() @@ -329,8 +358,16 @@ def _forward_backward(shape) -> float: return float(total_flops) -def evaluate_parallelism(hw_dict, model_config_obj, mode, parallel_settings): - hw_config = make_temp_hw_config(hw_dict, parallel_settings) +def evaluate_parallelism(hw_dict, model_config_obj, mode, parallel_settings, hw_mutator=None): + hw_config, debug_yaml = make_temp_hw_config(hw_dict, parallel_settings, hw_mutator=hw_mutator) + # if debug_yaml: + # print("=== DEBUG HW CONFIG ===") + # print(debug_yaml) + # try: + # with open("debug.yaml", "w") as debug_handle: + # debug_handle.write(debug_yaml) + # except Exception as exc: + # print(f"Warning: failed to write debug.yaml: {exc}", file=sys.stderr) temp_dir = tempfile.mkdtemp(prefix="parallelism_sweep_") try: calculator = TimeCalculationLLM(hw_config, model_config_obj, mode, output_dir=temp_dir) @@ -352,6 +389,7 @@ def evaluate_parallelism(hw_dict, model_config_obj, mode, parallel_settings): "achieved_flops": achieved_flops, "memory_exceeded": getattr(calculator, "memory_capacity_exceeded", False), "memory_violation_gb": getattr(calculator, "memory_capacity_violation_gb", 0.0), + "hw_yaml": debug_yaml, } finally: shutil.rmtree(temp_dir, ignore_errors=True) @@ -454,6 +492,11 @@ def parse_args(): action="store_true", help="Skip evaluation and only regenerate plots from the existing TSV report.", ) + parser.add_argument( + "--enforce-square-tp-cp", + action="store_true", + help="Require tp*cp to be a square power of two when enumerating configurations.", + ) return parser.parse_args() @@ -795,6 +838,8 @@ def main(): args = parse_args() set_astrasim_cache_mode(ASTRA_CACHE_MODE) + enforce_square = ENFORCE_SQUARE_TP_CP or args.enforce_square_tp_cp + results: List[Dict[str, object]] = [] if args.plot_only: @@ -836,9 +881,13 @@ def main(): evaluated = 0 filtered_tasks: List[Tuple[Tuple[str, object], ...]] = [] + skipped_square_constraint = 0 for items in task_items: settings = dict(items) num_gpus = total_gpu_count(settings) + if enforce_square and not tp_cp_product_is_power_of_two_square(settings.get("tp"), settings.get("cp")): + skipped_square_constraint += 1 + continue if GPU_COUNT_MIN <= num_gpus <= GPU_COUNT_MAX: filtered_tasks.append(items) else: @@ -847,6 +896,8 @@ def main(): if not filtered_tasks: print("No configurations within GPU count bounds.") return + if enforce_square and skipped_square_constraint: + print(f"Skipped {skipped_square_constraint} configuration(s) due to tp*cp square constraint.") available_cpus = max(1, os.cpu_count() or 1) if MAX_WORKERS is None or MAX_WORKERS <= 0: From bf3c0731d34a6ebe6c40adf2af454df19a6bd180 Mon Sep 17 00:00:00 2001 From: gkarfakis19 Date: Thu, 6 Nov 2025 18:05:11 -0800 Subject: [PATCH 25/77] hierarchical mode faults enabled after large reorganization of transformer graphs (now we make many graphs). DP Still doesnt work fully (WIP). Also added print logging system --- LLM_execution.py | 267 +++++++++++++++++++++++++++++++-------- astrasim_lib/executor.py | 28 +++- run_perf.py | 49 ++++--- simulate_LLM.py | 10 +- util.py | 101 ++++++++++++++- 5 files changed, 371 insertions(+), 84 deletions(-) diff --git a/LLM_execution.py b/LLM_execution.py index 5f29b94e..2171f92c 100644 --- a/LLM_execution.py +++ b/LLM_execution.py @@ -4,13 +4,16 @@ import os from dataclasses import dataclass from enum import Enum -from typing import Any, Dict, Tuple, Optional, List, Set, Sequence +from typing import Any, Dict, Tuple, Optional, List, Set, Sequence, Mapping +from collections import defaultdict +from collections import defaultdict import simulate_LLM from astrasim_lib import run_astra_simulation_only_onepath from astrasim_lib.fault_projection import FaultProjectionResult, FaultSpace from astrasim_lib.layout_utils import axis_layout_from_descriptor from simulate_LLM import Graph +from util import log_message def _env_flag(name: str) -> bool: @@ -638,8 +641,13 @@ def __init__( self.transformer_forward_root = transformer_forward_root self.transformer_backward_root = transformer_backward_root self.flattened_root: Optional[Any] = None - self._transformer_rank_layout: Optional[Dict[str, Any]] = None - self._pipeline_rank_layout: Optional[Dict[str, Any]] = None + self._transformer_rank_layout: Dict[str, Any] = {} + self._pipeline_rank_layout: Dict[str, Any] = {} + self._network_dimensions: Tuple[Any, ...] = tuple() + self._axis_dimension_map: Dict[str, int] = {} + self._transformer_stage_faults: Dict[int, Tuple[Tuple[int, int, float], ...]] = {} + self._transformer_stage_timings: Dict[int, TransformerTimings] = {} + self._transformer_baseline_timings: Optional[TransformerTimings] = None self._rank_layout = self._build_rank_layout_descriptor() self._fault_space: Optional[FaultSpace] = None self._fault_projections: Dict[str, FaultProjectionResult] = {} @@ -651,6 +659,7 @@ def _build_rank_layout_descriptor(self) -> Dict[str, Any]: dimensions = getattr(layout, "dimensions", None) if layout is not None else None if not dimensions: return {} + self._network_dimensions = tuple(dimensions) def _safe_int(value: Any, default: int = 1) -> int: try: @@ -731,13 +740,7 @@ def _safe_int(value: Any, default: int = 1) -> int: "axis_strides": axis_strides, "stage_span": span, } - pretty_sizes = ", ".join(f"{axis}={axis_sizes[axis]}" for axis in axis_order) - dp_msg = f", dp={dp_size}" if dp_size > 1 else "" - print( - "[DeepFlow][hw-id] Flattening axes for hw_id assignment: " - f"{axis_order or ['']} ({pretty_sizes or 'single device'}{dp_msg}) | " - "dp ranks are appended separately." - ) + def _subset_layout(allowed: Sequence[str]) -> Optional[Dict[str, Any]]: subset = [axis for axis in axis_order if axis in allowed and axis_sizes.get(axis, 1) > 1] @@ -763,20 +766,78 @@ def _subset_layout(allowed: Sequence[str]) -> Optional[Dict[str, Any]]: return descriptor + def _log_fault_summary(self, axis_order: Sequence[str], axis_sizes: Mapping[str, int]) -> None: + if not axis_order: + return + space = self._fault_space + if space is None or not space.entries: + return + network_dims = getattr(self, "_network_dimensions", tuple()) + if not network_dims: + log_message("[DeepFlow][faults] Hardware dimensions unavailable; skipping fault summary.") + return + + def coords_to_dict(coords: Tuple[Tuple[str, int], ...]) -> Dict[str, int]: + return {name: value for name, value in coords} + + entries = space.entries + for dim_index, dim in enumerate(network_dims): + axes = [str(axis).strip().lower() for axis in getattr(dim, "parallelisms", ())] + if not axes: + continue + replication_axes: List[str] = [] + for future_dim in network_dims[dim_index + 1 :]: + replication_axes.extend( + str(axis).strip().lower() for axis in getattr(future_dim, "parallelisms", ()) if axis + ) + total_clusters = 1 + for axis in replication_axes: + total_clusters *= max(1, axis_sizes.get(axis, 1)) + axes_label = ", ".join(axes) or "" + table_rows: List[Tuple[str, str, float]] = [] + for entry in entries: + if not any(axis in axes for axis in entry.affected_axes): + continue + src_dict = coords_to_dict(entry.src_coords) + dst_dict = coords_to_dict(entry.dst_coords) + src_label = f"{entry.original[0]} (" + ", ".join(f"{axis}={src_dict.get(axis, 0)}" for axis in axes) + ")" + dst_label = f"{entry.original[1]} (" + ", ".join(f"{axis}={dst_dict.get(axis, 0)}" for axis in axes) + ")" + higher_coords = " ".join(f"{axis}={src_dict.get(axis, 0)}" for axis in replication_axes) + table_rows.append((higher_coords, f"{src_label} <-> {dst_label}", float(entry.original[2]))) + + if not table_rows: + log_message(f" • dim{dim_index} ({axes_label}) → Affected HW clusters: 0 / {total_clusters}", category="faults") + continue + + log_message( + f" • dim{dim_index} ({axes_label}) → Affected HW clusters: {len(table_rows)} / {total_clusters or 1}", + category="faults", + ) + header_axes = " ".join(replication_axes) if replication_axes else "" + log_message(f" {header_axes:<12} | {'source ↔ dest':<40} | derate", category="faults") + log_message(f" {'-' * max(len(header_axes),12)}-+-{'-' * 40}-+-------", category="faults") + + for higher_str, pair, derate in table_rows: + log_message( + f" {higher_str:<12} | {pair:<40} | {derate:>6.2f}", + category="faults", + ) + log_message(f" {'-' * max(len(header_axes),12)}-+-{'-' * 40}-+-------", category="faults") + + + + def _initialize_fault_mappings(self) -> None: hw_config = getattr(self.time_calc, "hw_config", None) network_layout = getattr(hw_config, "network_layout", None) faulty_links: Tuple[Tuple[int, int, float], ...] = tuple( getattr(network_layout, "faulty_links", ()) or () ) - if not faulty_links: - self._fault_space = None - self._fault_projections = {} - return - if not self._rank_layout: - raise ValueError("Faulty links require a populated rank layout descriptor.") + if faulty_links and self.time_calc.execution_mode in {ExecutionMode.ANALYTICAL, ExecutionMode.HYBRID}: + raise ValueError("Faulty links require full AstraSim execution; analytical/hybrid modes are not supported.") axis_layout = axis_layout_from_descriptor(self._rank_layout) axis_dim_map = self._axis_to_dimension_map(network_layout) + self._axis_dimension_map = dict(axis_dim_map) self._fault_space = FaultSpace( axis_layout, faulty_links, @@ -784,6 +845,10 @@ def _initialize_fault_mappings(self) -> None: ) self._fault_projections = self._build_fault_projections_from_space(self._fault_space) self._validate_fault_coverage() + self._transformer_stage_faults = self._build_transformer_stage_fault_map() + axis_order = self._rank_layout.get("axis_order", []) if isinstance(self._rank_layout, dict) else [] + axis_sizes = self._rank_layout.get("axis_sizes", {}) if isinstance(self._rank_layout, dict) else {} + self._log_fault_summary(axis_order, axis_sizes) def _axis_to_dimension_map(self, network_layout) -> Dict[str, int]: mapping: Dict[str, int] = {} @@ -845,6 +910,32 @@ def _validate_fault_coverage(self) -> None: f"{formatted}" ) + def _axis_value_from_coords(self, coords: Tuple[Tuple[str, int], ...], axis: str) -> int: + for name, value in coords: + if name == axis: + return int(value) + layout = self._rank_layout or {} + axis_sizes = layout.get("axis_sizes", {}) if isinstance(layout, dict) else {} + if axis_sizes.get(axis, 1) <= 1: + return 0 + raise ValueError(f"Axis '{axis}' not present in coordinate tuple for faulty link.") + + def _build_transformer_stage_fault_map(self) -> Dict[int, Tuple[Tuple[int, int, float], ...]]: + projection = self._fault_projection_for("transformer") + if not projection: + return {} + stage_faults: Dict[int, List[Tuple[int, int, float]]] = {} + for detail in projection.entries: + src_stage = self._axis_value_from_coords(detail.src_coords, "lp") + dst_stage = self._axis_value_from_coords(detail.dst_coords, "lp") + if src_stage != dst_stage: + raise ValueError( + "Transformer faulty link spans multiple pipeline stages; " + "hierarchical execution requires faults limited to a single stage." + ) + stage_faults.setdefault(src_stage, []).append(detail.remapped) + return {stage: tuple(links) for stage, links in stage_faults.items()} + def _fault_projection_for(self, label: str) -> Optional[FaultProjectionResult]: return self._fault_projections.get(label) @@ -1096,72 +1187,117 @@ def enqueue_children(obj: Any) -> None: def _run_transformer_astrasim(self, mode: ExecutionMode) -> Optional[TransformerTimings]: del mode # mode currently unused but kept for signature consistency - # Use hierarchical artifact directory when persisting artifacts for transformer simulation - artifact_dir = self.time_calc.output_dir - artifact_dir_fwd = artifact_dir - artifact_dir_bwd = artifact_dir - os.makedirs(artifact_dir, exist_ok=True) - if self.time_calc.persist_astrasim_artifacts: - artifact_dir = os.path.join(self.time_calc.output_dir, "astra_hier") - artifact_dir_fwd = os.path.join(artifact_dir, "fwd") - artifact_dir_bwd = os.path.join(artifact_dir, "bwd") - os.makedirs(artifact_dir_fwd, exist_ok=True) - os.makedirs(artifact_dir_bwd, exist_ok=True) + if not self.transformer_forward_root and not self.transformer_backward_root: + if getattr(self, "_transformer_stage_faults", {}): + raise ValueError("Transformer faults require transformer graph metadata, but none is available.") + return None + + layout = getattr(self, "_transformer_rank_layout", {}) + if self.transformer_forward_root: + setattr(self.transformer_forward_root, "_astrasim_rank_layout", layout) + if self.transformer_backward_root: + setattr(self.transformer_backward_root, "_astrasim_rank_layout", layout) + + persist = self.time_calc.persist_astrasim_artifacts + os.makedirs(self.time_calc.output_dir, exist_ok=True) + self._transformer_stage_timings = {} + self._transformer_baseline_timings = None + + # Baseline run (no transformer faults) + baseline_fwd_dir, baseline_bwd_dir = self._transformer_artifact_dirs(label=None, persist=persist) + baseline_timings, baseline_fwd_per_rank, baseline_bwd_per_rank = self._execute_transformer_run( + baseline_fwd_dir, + baseline_bwd_dir, + faulty_links_override=(), + ) + self._transformer_baseline_timings = baseline_timings + self.time_calc.transformer_astrasim_per_rank_forward = baseline_fwd_per_rank + self.time_calc.transformer_astrasim_per_rank_backward = baseline_bwd_per_rank + self.time_calc.transformer_astrasim_time_forward = baseline_timings.forward + self.time_calc.transformer_astrasim_time_backward = baseline_timings.backward + + # Per-stage fault runs + stage_faults = getattr(self, "_transformer_stage_faults", {}) + for fault_index, (stage_id, fault_links) in enumerate(sorted(stage_faults.items())): + label = f"fault{fault_index}" + stage_fwd_dir, stage_bwd_dir = self._transformer_artifact_dirs(label=label, persist=persist) + stage_timings, _, _ = self._execute_transformer_run( + stage_fwd_dir, + stage_bwd_dir, + faulty_links_override=fault_links, + ) + self._transformer_stage_timings[stage_id] = stage_timings + + return baseline_timings + + def _transformer_artifact_dirs(self, label: Optional[str], persist: bool) -> Tuple[str, str]: + if not persist: + base_dir = self.time_calc.output_dir + os.makedirs(base_dir, exist_ok=True) + return base_dir, base_dir + base_dir = os.path.join(self.time_calc.output_dir, "astra_hier") + os.makedirs(base_dir, exist_ok=True) + if label is None: + fwd_dir = os.path.join(base_dir, "fwd") + bwd_dir = os.path.join(base_dir, "bwd") + else: + fwd_dir = os.path.join(base_dir, f"{label}_fwd") + bwd_dir = os.path.join(base_dir, f"{label}_bwd") + os.makedirs(fwd_dir, exist_ok=True) + os.makedirs(bwd_dir, exist_ok=True) + return fwd_dir, bwd_dir + def _execute_transformer_run( + self, + artifact_dir_fwd: str, + artifact_dir_bwd: str, + *, + faulty_links_override: Optional[Tuple[Tuple[int, int, float], ...]], + ) -> Tuple[TransformerTimings, Optional[List[float]], Optional[List[float]]]: fwd_per_rank = None bwd_per_rank = None fwd_max = 0 bwd_max = 0 - transformer_fault_override = self._fault_override("transformer") + if self.transformer_forward_root: - layout = getattr(self, "_transformer_rank_layout", None) - if layout: - setattr(self.transformer_forward_root, "_astrasim_rank_layout", layout) - elif hasattr(self.transformer_forward_root, "_astrasim_rank_layout"): - delattr(self.transformer_forward_root, "_astrasim_rank_layout") fwd_per_rank, fwd_max = run_astra_simulation_only_onepath( self.transformer_forward_root, self.time_calc, artifact_dir_fwd, dp_override=1, persist_artifacts=self.time_calc.persist_astrasim_artifacts, - faulty_links_override=transformer_fault_override, + faulty_links_override=faulty_links_override, ) + if fwd_max <= 0: + raise RuntimeError("AstraSim transformer forward execution returned non-positive duration") + if self.transformer_backward_root: - layout = getattr(self, "_transformer_rank_layout", None) - if layout: - setattr(self.transformer_backward_root, "_astrasim_rank_layout", layout) - elif hasattr(self.transformer_backward_root, "_astrasim_rank_layout"): - delattr(self.transformer_backward_root, "_astrasim_rank_layout") bwd_per_rank, bwd_max = run_astra_simulation_only_onepath( self.transformer_backward_root, self.time_calc, artifact_dir_bwd, dp_override=1, persist_artifacts=self.time_calc.persist_astrasim_artifacts, - faulty_links_override=transformer_fault_override, + faulty_links_override=faulty_links_override, ) + if bwd_max < 0: + raise RuntimeError("AstraSim transformer backward execution returned non-positive duration") - self.time_calc.transformer_astrasim_per_rank_forward = fwd_per_rank - self.time_calc.transformer_astrasim_per_rank_backward = bwd_per_rank - self.time_calc.transformer_astrasim_time_forward = fwd_max - self.time_calc.transformer_astrasim_time_backward = bwd_max - - if fwd_max < 0 or bwd_max < 0: - raise RuntimeError("AstraSim transformer execution returned non-positive duration") - - return TransformerTimings(forward=fwd_max, backward=bwd_max) + return TransformerTimings(forward=fwd_max, backward=bwd_max), fwd_per_rank, bwd_per_rank def _apply_transformer_time(self, timings: TransformerTimings) -> None: if timings.forward < 0 or timings.backward < 0: raise ValueError("AstraSim transformer times must be positive") + baseline_timings = self._transformer_baseline_timings or timings + stage_timings = getattr(self, "_transformer_stage_timings", {}) + comp_times = getattr(self.pipeline_graph, "comp_times", None) if isinstance(comp_times, dict): if "transformer_f" in comp_times: - comp_times["transformer_f"] = timings.forward + comp_times["transformer_f"] = baseline_timings.forward if "transformer_b" in comp_times: - comp_times["transformer_b"] = timings.backward + comp_times["transformer_b"] = baseline_timings.backward visited: Set[int] = set() roots: List[Any] @@ -1171,9 +1307,22 @@ def _apply_transformer_time(self, timings: TransformerTimings) -> None: roots = [self.pipeline_root] for root in roots: - self._assign_transformer_durations(root, visited, timings.forward, timings.backward) + self._assign_transformer_durations( + root, + visited, + stage_timings, + baseline_timings.forward, + baseline_timings.backward, + ) - def _assign_transformer_durations(self, node: Any, visited: Set[int], forward_value: float, backward_value: float) -> None: + def _assign_transformer_durations( + self, + node: Any, + visited: Set[int], + stage_timings: Dict[int, TransformerTimings], + forward_default: float, + backward_default: float, + ) -> None: if node is None: return node_id = id(node) @@ -1184,10 +1333,16 @@ def _assign_transformer_durations(self, node: Any, visited: Set[int], forward_va if isinstance(node, simulate_LLM.Node): base_name = str(getattr(node, "name", "") or "") if base_name.startswith("transformer_layer"): + hw_stage = None + try: + hw_stage = int(getattr(node, "hw_id", None)) + except (TypeError, ValueError): + hw_stage = None + timing_override = stage_timings.get(hw_stage) if hw_stage is not None else None if getattr(node, "fwd", True): - node.duration = forward_value + node.duration = timing_override.forward if timing_override else forward_default else: - node.duration = backward_value + node.duration = timing_override.backward if timing_override else backward_default for child in getattr(node, "children", []): - self._assign_transformer_durations(child, visited, forward_value, backward_value) + self._assign_transformer_durations(child, visited, stage_timings, forward_default, backward_default) diff --git a/astrasim_lib/executor.py b/astrasim_lib/executor.py index c03fb2bc..42a107e6 100644 --- a/astrasim_lib/executor.py +++ b/astrasim_lib/executor.py @@ -42,6 +42,7 @@ from .integration import get_remote_memory_path, run_astrasim_analytical, run_cache_astrasim from .layout_utils import derive_axes_filter from simulate_LLM import visualize_graph +from util import relpath_display def _env_truthy(name: str) -> bool: @@ -428,14 +429,37 @@ def _render_et_file(et_path: str) -> None: except Exception as exc: print(f"[WARN] Failed to render Graphviz graph for {et_path}: {exc}") - for et_path in et_paths: - message = f" | ET Graph saved to {et_path}.svg" + if not et_paths: + return + + if len(et_paths) == 1: + et_path = et_paths[0] + display_path = relpath_display(f"{et_path}.svg") + message = f" | ET Graph saved to {display_path}" graphviz_async.submit( f"et:{os.path.basename(et_path)}", _render_et_file, et_path, print_message=message, ) + return + + first_path = et_paths[0] + display_first = relpath_display(f"{first_path}.svg") + summary_message = f" | {len(et_paths)} ET graphs saved to {display_first} (..)" + graphviz_async.submit( + f"et:{os.path.basename(first_path)}", + _render_et_file, + first_path, + print_message=summary_message, + ) + for et_path in et_paths[1:]: + graphviz_async.submit( + f"et:{os.path.basename(et_path)}", + _render_et_file, + et_path, + print_message=None, + ) diff --git a/run_perf.py b/run_perf.py index 265fac13..246279e7 100755 --- a/run_perf.py +++ b/run_perf.py @@ -11,6 +11,7 @@ import yaml import shutil import util +from util import log_message, flush_log_queue, extend_log import graphviz_async from tile import TiledGEMM, formatBytes @@ -49,6 +50,7 @@ def _report_total_wall_time() -> None: # Best-effort only pass +atexit.register(flush_log_queue) atexit.register(_report_total_wall_time) def parse_arguments(): @@ -252,13 +254,12 @@ def _run_llm_training(exp_hw_config, exp_model_config, exp_dir, mode): handle.write("\n".join(topology_lines)) handle.write("\n") - print("Training time for batch: {:.2f}s".format(tc_llm.get_time())) - for line in topology_lines: - print(line) - print("LLM training results written to {}".format(output_file)) + log_message("Training time for batch: {:.2f}s".format(tc_llm.get_time()), category="results") + extend_log(topology_lines, category="network") + log_message("LLM training results written to {}".format(output_file), category="results") warning_message = tc_llm.memory_capacity_warning() if warning_message: - print(warning_message) + log_message(warning_message) def _run_llm_inference(exp_hw_config, exp_model_config, exp_dir, mode): @@ -270,20 +271,22 @@ def _run_llm_inference(exp_hw_config, exp_model_config, exp_dir, mode): total_time = inference_timing["total_inference_time"] decode_rates = inference_timing.get("decode_tokens_per_s") or {} - print( + log_message( "LLM inference time: {:.2f}s (mode={})".format( total_time, tc_inf.execution_mode.value - ) + ), + category="results", ) - print( + log_message( "LLM time to first token: {:.2f}s".format( inference_timing["time_to_first_token"], - ) + ), + category="results", ) dp_replicas = max(1, getattr(tc_inf, "dp", 1)) batch_size = getattr(tc_inf, "batch_size", 1) if dp_replicas > 1: - print(f"Data parallel replicas: {dp_replicas}") + log_message(f"Data parallel replicas: {dp_replicas}", category="results") if decode_rates: # decode_rates are per-generation rates (tokens per second per generation) start_gen_rate = decode_rates.get("start", 0.0) @@ -292,17 +295,18 @@ def _run_llm_inference(exp_hw_config, exp_model_config, exp_dir, mode): mid_step = int(decode_rates.get("midpoint_step", 0.0)) # Print per-generation rates - print( + log_message( "Decode sequences/s: start={:.2f}, mid(token {})={:.2f}, end={:.2f}".format( start_gen_rate, mid_step, mid_gen_rate, end_gen_rate, - ) + ), + category="results", ) # Print aggregate decode throughput (with batch_size and dp multipliers) - print( + log_message( "Aggregate decode throughput tok/s (batch={}, dp={}): start={:.2f}, mid(token {})={:.2f}, end={:.2f}".format( batch_size, dp_replicas, @@ -310,7 +314,8 @@ def _run_llm_inference(exp_hw_config, exp_model_config, exp_dir, mode): mid_step, mid_gen_rate * batch_size * dp_replicas, end_gen_rate * batch_size * dp_replicas, - ) + ), + category="results", ) topology_lines = util.network_topology_summary_inference(exp_hw_config) @@ -341,12 +346,11 @@ def _run_llm_inference(exp_hw_config, exp_model_config, exp_dir, mode): handle.write("\n".join(topology_lines)) handle.write("\n") - for line in topology_lines: - print(line) - print("LLM inference results written to {}".format(output_path)) + extend_log(topology_lines, category="network") + log_message("LLM inference results written to {}".format(output_path), category="results") warning_message = tc_inf.memory_capacity_warning() if warning_message: - print(warning_message) + log_message(warning_message) if __name__ == "__main__": args = parse_arguments() @@ -364,7 +368,6 @@ def _run_llm_inference(exp_hw_config, exp_model_config, exp_dir, mode): os.makedirs(exp_dir, exist_ok=True) if mode == "LLM": - print("Using LLM parameters for computation...") run_LLM( exp_hw_config_path=config_hardware_path, exp_model_config_path=config_model_path, @@ -373,7 +376,7 @@ def _run_llm_inference(exp_hw_config, exp_model_config, exp_dir, mode): ) elif mode == "LSTM": - print("Using LSTM parameters for computation...") + log_message("Using LSTM parameters for computation...") run_LSTM( exp_hw_config_path=config_hardware_path, exp_model_config_path=config_model_path, @@ -382,7 +385,7 @@ def _run_llm_inference(exp_hw_config, exp_model_config, exp_dir, mode): ) elif mode == "GEMM": - print("Using GEMM parameters for computation...") + log_message("Using GEMM parameters for computation...") run_GEMM( exp_hw_config_path=config_hardware_path, exp_model_config_path=config_model_path, @@ -391,3 +394,7 @@ def _run_llm_inference(exp_hw_config, exp_model_config, exp_dir, mode): ) else: print("Invalid mode selected. Please choose 'LLM', 'LSTM', or 'GEMM'.") + flush_log_queue() + sys.exit(1) + + flush_log_queue() diff --git a/simulate_LLM.py b/simulate_LLM.py index 0afd6a84..fbc2f620 100644 --- a/simulate_LLM.py +++ b/simulate_LLM.py @@ -7,6 +7,7 @@ import os import graphviz_async +import util debug = False BYTES_PER_GIB = 1024 ** 3 @@ -1498,10 +1499,13 @@ def _count_transformer_layers_per_device(root_obj: Any, num_devices: int) -> Lis def save_graph(self, roots, output_folder = "output/LLM/", filename="graph"): os.makedirs(output_folder, exist_ok=True) - printstr = " | Graph saved to %s%s.svg" % (output_folder, filename) + base_path = os.path.normpath(f"{output_folder}{filename}") + svg_path = f"{base_path}.svg" + display_path = util.relpath_display(svg_path) + printstr = f" | Graph saved to {display_path}" def _render_graph() -> None: - dot_fw = visualize_graph(roots, filename=output_folder + filename) - dot_fw.render(output_folder + filename, format="svg", cleanup=True) + dot_fw = visualize_graph(roots, filename=base_path) + dot_fw.render(base_path, format="svg", cleanup=True) graphviz_async.submit(f"{filename}.svg", _render_graph, print_message=printstr) diff --git a/util.py b/util.py index 772b28dd..b3eef075 100644 --- a/util.py +++ b/util.py @@ -2,6 +2,8 @@ import sys import os import config +import threading +from typing import Iterable, List, Optional, Tuple core=1 DRAM=1 @@ -12,6 +14,90 @@ DF_CACHE_DIR = "./cache" +_LOG_MESSAGES: List[Tuple[Optional[str], str]] = [] +_LOG_LOCK = threading.Lock() + +_SECTION_ORDER = [ + ("network", "TOPOLOGY/NETWORK"), + ("faults", "FAULTY LINKS"), + ("results", "RESULTS"), +] + +_REPO_ROOT = os.path.abspath(os.environ.get("DEEPFLOW_REPO_ROOT", os.getcwd())) + + +def log_message(message: str, category: Optional[str] = None) -> None: + """Append ``message`` to the shared log queue.""" + if message is None: + return + text = str(message) + if not text: + return + cat_normalized = str(category).strip().lower() if category else None + with _LOG_LOCK: + _LOG_MESSAGES.append((cat_normalized, text)) + + +def extend_log(lines: Iterable[str], category: Optional[str] = None) -> None: + """Append multiple lines to the log queue.""" + for line in lines: + log_message(line, category=category) + + +def drain_log_messages() -> List[Tuple[Optional[str], str]]: + """Return and clear all queued log messages (category, message).""" + with _LOG_LOCK: + if not _LOG_MESSAGES: + return [] + drained = list(_LOG_MESSAGES) + _LOG_MESSAGES.clear() + return drained + + +def flush_log_queue() -> None: + """Print and clear the queued log messages grouped by category.""" + entries = drain_log_messages() + print("\n") + if not entries: + return + categorized = {cat: [] for cat, _ in _SECTION_ORDER} + uncategorized: List[str] = [] + for category, message in entries: + if category in categorized: + categorized[category].append(message) + else: + uncategorized.append(message) + + sections_printed = False + section_border = "=" * 60 + for cat, title in _SECTION_ORDER: + lines = categorized.get(cat) or [] + if not lines: + continue + sections_printed = True + print(f"{section_border}\n{title}\n{section_border}") + for line in lines: + print(line) + if sections_printed: + print(section_border) + + for line in uncategorized: + print(line) + + +def relpath_display(path: str) -> str: + """Return ``path`` relative to the repo root when possible.""" + if not path: + return "" + abs_path = os.path.abspath(path) + try: + rel = os.path.relpath(abs_path, start=_REPO_ROOT) + except Exception: + return abs_path + if rel.startswith(".."): + return abs_path + return rel + def _collect_parallelism_values(hw_config): sch_config = getattr(hw_config, "sch_config", None) @@ -41,7 +127,18 @@ def _format_parallelism_terms(dim, parallelism_values): def network_topology_summary_training(hw_config): parallelism_values = _collect_parallelism_values(hw_config) dimensions = list(getattr(hw_config.network_layout, "dimensions", ())) - lines = [f"Network Topology [dims={len(dimensions)}]"] + ordered_axes = ["tp", "cp", "lp", "dp"] + formatted_terms = [] + for axis in ordered_axes: + value = parallelism_values.get(axis) + if value is None: + continue + formatted_terms.append(f"{axis}:{value}") + formatted_parallelisms = ", ".join(formatted_terms) if formatted_terms else "none" + lines = [ + f"Parallelisms: {formatted_parallelisms}", + f"Network Topology [dims={len(dimensions)}]", + ] aggregate = 1 for dim in dimensions: @@ -60,7 +157,7 @@ def network_topology_summary_training(hw_config): f" • {dim.id} {dim.label} : {axis_repr} ⇒ size {size_display}" ) - lines.append(f" => aggregate = {aggregate} GPUs") + lines.append(f" ⇒ total {aggregate} devices") return lines From 99681a6d2a06b60761393d6984e8784c6782fbc4 Mon Sep 17 00:00:00 2001 From: NanoMembers Date: Thu, 6 Nov 2025 23:53:47 -0800 Subject: [PATCH 26/77] Support no-DP graph generation for gradient-accumulation timing --- LLM_execution.py | 32 +++++++++++-------- simulate_inf.py | 2 ++ time_calculation_LLM.py | 71 +++++++++++++++++++++++++++++++++++++---- time_calculation_inf.py | 2 ++ 4 files changed, 88 insertions(+), 19 deletions(-) diff --git a/LLM_execution.py b/LLM_execution.py index 2171f92c..579c3e1d 100644 --- a/LLM_execution.py +++ b/LLM_execution.py @@ -632,6 +632,7 @@ def __init__( transformer_graph: Optional[Graph] = None, transformer_forward_root: Optional[Any] = None, transformer_backward_root: Optional[Any] = None, + no_data_parallel: bool = False, ) -> None: self.time_calc = time_calc self.pipeline_graph = pipeline_graph @@ -648,6 +649,7 @@ def __init__( self._transformer_stage_faults: Dict[int, Tuple[Tuple[int, int, float], ...]] = {} self._transformer_stage_timings: Dict[int, TransformerTimings] = {} self._transformer_baseline_timings: Optional[TransformerTimings] = None + self.no_data_parallel = bool(no_data_parallel) self._rank_layout = self._build_rank_layout_descriptor() self._fault_space: Optional[FaultSpace] = None self._fault_projections: Dict[str, FaultProjectionResult] = {} @@ -962,10 +964,16 @@ def run(self, mode: ExecutionMode) -> ExecutionResult: def _run_pipeline_with_analytical_comm(self, declared_mode: ExecutionMode) -> ExecutionResult: if declared_mode == ExecutionMode.HYBRID: - filename = "/hybrid_graph" + if self.no_data_parallel: + filename = "/hybrid_graph_no_dp" + else: + filename = "/hybrid_graph" timed_root = self.pipeline_root else: # must be "ANALYTICAL" - filename = "/analytical_graph" + if self.no_data_parallel: + filename = "/analytical_graph_no_dp" + else: + filename = "/analytical_graph" timed_root = self.pipeline_graph.convert_comm_sizes_to_times( self.pipeline_root, self.time_calc.network_model, @@ -987,12 +995,7 @@ def _run_pipeline_with_analytical_comm(self, declared_mode: ExecutionMode) -> Ex def _run_hybrid(self) -> ExecutionResult: generate_graphs = _env_flag("DEEPFLOW_VISUALIZE_GRAPHS") - if generate_graphs: - self.transformer_graph.save_graph( - self.transformer_forward_root, - self.time_calc.output_dir, - "/hybrid_graph_transformer", - ) + transformer_time = self._run_transformer_astrasim(ExecutionMode.HYBRID) if generate_graphs: self.transformer_graph.save_graph( @@ -1030,10 +1033,11 @@ def _run_full_astrasim_hierarchical(self) -> ExecutionResult: run_kwargs["dp_override"] = 1 if _env_flag("DEEPFLOW_VISUALIZE_GRAPHS") and self.pipeline_root is not None: + filename = "/pipeline_graph_hierarchical_no_dp" if self.no_data_parallel else "/pipeline_graph_hierarchical" self.pipeline_graph.save_graph( self.pipeline_root, self.time_calc.output_dir, - "/pipeline_graph_hierarchical", + filename, ) pipeline_layout = getattr(self, "_pipeline_rank_layout", None) @@ -1065,11 +1069,12 @@ def _run_full_astrasim_flattened(self) -> ExecutionResult: rank_layout=self._rank_layout, ) - if _env_flag("DEEPFLOW_VISUALIZE_GRAPHS") and self.pipeline_root is not None: + if _env_flag("DEEPFLOW_VISUALIZE_GRAPHS") and self.pipeline_root is not None : + filename = "/pipeline_graph_pre_flatten_no_dp" if self.no_data_parallel else "/pipeline_graph_pre_flatten" self.pipeline_graph.save_graph( self.pipeline_root, self.time_calc.output_dir, - "/pipeline_graph_pre_flatten", + filename, ) flattened_root = flattener.build(self.pipeline_root) if flattened_root is None: @@ -1079,10 +1084,11 @@ def _run_full_astrasim_flattened(self) -> ExecutionResult: setattr(flattened_root, "_astrasim_rank_layout", self._rank_layout) self.time_calc.flattened_pipeline_root = flattened_root if _env_flag("DEEPFLOW_VISUALIZE_GRAPHS") and self.pipeline_root is not None: + filename = "/pipeline_graph_post_flatten_no_dp" if self.no_data_parallel else "/pipeline_graph_post_flatten" self.pipeline_graph.save_graph( flattened_root, self.time_calc.output_dir, - "/pipeline_graph_post_flatten", + filename, ) self.pipeline_root = flattened_root # output_dir = "./astra_flattened_graph" @@ -1122,6 +1128,7 @@ def _run_full_astrasim_flattened(self) -> ExecutionResult: raise RuntimeError("AstraSim flattened execution returned no per-rank timings") expected_rank_count = effective_dp * len(unique_hw_ids) + # Special case: If expected rank count is 1, then 2 is fine, but we prune the extra result # this is done, since astrasim backend only supports >1 ranks, so we generate extra fake result for that case. if expected_rank_count == 1: @@ -1131,7 +1138,6 @@ def _run_full_astrasim_flattened(self) -> ExecutionResult: f"expected {expected_rank_count}, got {len(per_rank_sec)}" ) per_rank_sec = per_rank_sec[:1] - if len(per_rank_sec) != expected_rank_count: raise RuntimeError( "AstraSim rank count mismatch for flattened execution: " diff --git a/simulate_inf.py b/simulate_inf.py index e97afad3..744bcb3b 100644 --- a/simulate_inf.py +++ b/simulate_inf.py @@ -196,6 +196,8 @@ def _execute_decode_step( ( pipeline_graph, pipeline_root, + _, + _, transformer_graph, transformer_forward_root, transformer_backward_root, diff --git a/time_calculation_LLM.py b/time_calculation_LLM.py index f3b58f3a..20c7218a 100644 --- a/time_calculation_LLM.py +++ b/time_calculation_LLM.py @@ -165,6 +165,8 @@ def __init__(self, hw_config, model_config, mode, output_dir: Optional[str] = No self.transformer_astrasim_per_rank_backward: Optional[List[float]] = None self.pipeline_astrasim_time: Optional[float] = None self.pipeline_astrasim_per_rank: Optional[List[float]] = None + self.pipeline_graph_no_dp: Optional[Graph] = None + self.pipeline_root_no_dp: Optional[Any] = None def _sequence_parallel_degree(self) -> int: """Return tensor-parallel degree used for sequence-parallel collectives. @@ -2026,7 +2028,16 @@ def _prepare_execution_graphs( zero3_embedding_gather_bytes: float = 0.0, zero3_transformer_gather_bytes: float = 0.0, zero3_softmax_gather_bytes: float = 0.0, - ) -> Tuple[Graph, Any, Optional[Graph], Optional[Any], Optional[Any], Dict[str, Tuple[float, float]]]: + ) -> Tuple[ + Graph, + Any, + Optional[Graph], + Optional[Any], + Optional[Any], + Dict[str, Tuple[float, float]], + Optional[Graph], + Optional[Any], + ]: """Build pipeline/transformer graphs shared across training and inference.""" if not include_pipeline_backward and not include_transformer_backward: @@ -2260,12 +2271,33 @@ def _resolve_operation(name: str) -> OperationTiming: comm_metadata=comm_metadata, misc_metadata=misc_metadata, ) + graph_root = pipeline_graph_obj.construct_fwd_bwd_graph(include_backward=include_pipeline_backward) + pipeline_graph_obj_no_dp = None + graph_root_no_dp = None + need_no_dp_variant = getattr(self, "gradient_accumulation_steps", 1) > 1 + if need_no_dp_variant: + pipeline_graph_obj_no_dp = simulate_LLM.Graph( + mode="pipeline", + dp=1, + lp=self.lp, + tp=self.tp, + cp=self.cp, + comp_times=comp_times, + comm_metadata=comm_metadata, + misc_metadata=misc_metadata, + ) + + graph_root_no_dp = pipeline_graph_obj_no_dp.construct_fwd_bwd_graph(include_backward=include_pipeline_backward) + + interconnect_params = self._build_interconnect_params() return ( pipeline_graph_obj, graph_root, + pipeline_graph_obj_no_dp, + graph_root_no_dp, transformer_graph, transformer_forward_root, transformer_backward_root, @@ -2384,6 +2416,8 @@ def calc_time_llm(self): ( pipeline_graph_obj, graph_root, + pipeline_graph_obj_no_dp, + graph_root_no_dp, transformer_graph, transformer_forward_root, transformer_backward_root, @@ -2416,9 +2450,30 @@ def calc_time_llm(self): self.pipeline_graph = pipeline_graph_obj self.pipeline_root = graph_root self.pipeline_interconnect = interconnect_params + self.pipeline_graph_no_dp = pipeline_graph_obj_no_dp + self.pipeline_root_no_dp = graph_root_no_dp mode = self.execution_mode - + time_fw_bw_no_dp: Optional[float] = None + if self.gradient_accumulation_steps > 1: + if not (self.pipeline_graph_no_dp and self.pipeline_root_no_dp): + raise RuntimeError("Gradient accumulation steps > 1 requires a no-DP pipeline graph") + dispatcher_no_dp = LLMExecutionDispatcher( + time_calc=self, + pipeline_graph=self.pipeline_graph_no_dp, + pipeline_root=self.pipeline_root_no_dp, + interconnect_params=self.pipeline_interconnect, + transformer_graph=self.transformer_graph, + transformer_forward_root=self.transformer_forward_root, + transformer_backward_root=self.transformer_backward_root, + no_data_parallel=True, + ) + try: + result_no_dp = dispatcher_no_dp.run(mode) + except NotImplementedError as exc: + raise NotImplementedError(f"{exc}. Selected execution mode '{mode.value}'.") from exc + time_fw_bw_no_dp = result_no_dp.total_time + dispatcher = LLMExecutionDispatcher( time_calc=self, pipeline_graph=self.pipeline_graph, @@ -2427,15 +2482,15 @@ def calc_time_llm(self): transformer_graph=self.transformer_graph, transformer_forward_root=self.transformer_forward_root, transformer_backward_root=self.transformer_backward_root, + no_data_parallel=False, ) - mode = self.execution_mode try: result = dispatcher.run(mode) except NotImplementedError as exc: raise NotImplementedError(f"{exc}. Selected execution mode '{mode.value}'.") from exc - time_fw_bw = result.total_time + pipeline_root = result.graph_root self.pipeline_graph = dispatcher.pipeline_graph self.pipeline_root = pipeline_root @@ -2452,8 +2507,12 @@ def calc_time_llm(self): print(f"Actual transformer forward time: {self.transformer_astrasim_time_forward:.4f}s") print(f"Actual transformer backward time: {self.transformer_astrasim_time_backward:.4f}s") - self.tot_time = time_fw_bw * self.gradient_accumulation_steps - + if self.gradient_accumulation_steps > 1: #if gradient accumulation is enabled, non-last steps have no dp overhead, only last step has dp overhead + if time_fw_bw_no_dp is None: + raise RuntimeError("Missing no-DP timing for gradient accumulation computation") + self.tot_time = time_fw_bw_no_dp * (self.gradient_accumulation_steps - 1) + time_fw_bw + else: + self.tot_time = time_fw_bw return self.tot_time def _simulate_with_memory( diff --git a/time_calculation_inf.py b/time_calculation_inf.py index 5f84d759..c71a56f8 100644 --- a/time_calculation_inf.py +++ b/time_calculation_inf.py @@ -473,6 +473,8 @@ def calc_time(self) -> float: ( pipeline_graph, pipeline_root, + _, + _, transformer_graph, transformer_forward_root, _, From a661557a692fa319fb6b0dcd62edf5a96754fb3b Mon Sep 17 00:00:00 2001 From: gkarfakis19 Date: Fri, 7 Nov 2025 20:40:28 -0800 Subject: [PATCH 27/77] implemented DP faults (harder than I thought) --- HIER_BUG.md | 25 ------------- LLM_execution.py | 74 ++++++++++++++++++++++++------------ astrasim_lib/executor.py | 15 +++++++- simulate_LLM.py | 81 ++++++++++++++++++++++++++++++++++------ 4 files changed, 135 insertions(+), 60 deletions(-) delete mode 100644 HIER_BUG.md diff --git a/HIER_BUG.md b/HIER_BUG.md deleted file mode 100644 index 013ac907..00000000 --- a/HIER_BUG.md +++ /dev/null @@ -1,25 +0,0 @@ -# Hierarchical Mode Fault Bugs - -## 1. Fault Projection Leak -Hierarchical execution runs AstraSim separately for the transformer (TP/CP axes) and pipeline (LP axis) graphs. Each run filters the hardware layout down to the axes needed for that graph, but `generate_astrasim_configs_from_hw()` still injects the original `network.faulty_links` list without remapping the endpoints into the reduced coordinate space. Any fault whose flattened IDs fall inside the smaller rank range now appears to live on *every* simulated axis subset. This reinterprets a single physical faulty link multiple times (e.g., once in the transformer run and again in the LP run), so hierarchical mode applies the same degradation repeatedly and slows far more than flattened execution. - -## 2. Single-Replica Transformer Timing -The hierarchical pipeline measures AstraSim on exactly one transformer replica and copies that timing to every transformer node in the pipeline graph. That assumption holds only when all stages experience identical network conditions. Faulty links are physical, rank-specific degradations—different pipeline stages map to different hardware coordinates, so only the stages whose TP/CP cluster actually touches the faulty link should be slowed. Because hierarchical mode reuses the single measurement for every stage, *all* stages inherit the worst-case transformer time, even if the fault affects just one cluster. Consequently hierarchical mode always over-penalizes the transformer, regardless of where the fault truly sits. - -## Remediation Plan (hierarchical & hybrid) -1. **Decode Fault Coordinates** - - Reconstruct the global axis layout (`tp`, `cp`, `lp`, `dp`) exactly as flattened mode assigns hw_ids. - - For every `faulty_links` entry, convert `src/dst` into `(tp,cp,lp,dp)` tuples, then re-linearize per axis subset. This projection will emit two remapped lists: one for TP/CP-only graphs (transformer) and one for LP/DP graphs (pipeline). Hybrid mode follows the same flow because it still invokes AstraSim for transformer+pipeline pieces separately. - -2. **Route Faults to the Correct Graph** - - Use the projected coordinates to classify each link by the highest dimension it touches (dim0 → transformer cluster, dim≥1 → pipeline/DP axes). Only pass transformer-dimension faults into the transformer AstraSim jobs; only pass higher-dimension faults into the pipeline job. - -3. **Per-Cluster Transformer Simulations** - - Identify the set of TP/CP clusters impacted by faults. Run AstraSim `(N + 1)` times per direction: one baseline (no transformer faults) plus one run per faulty cluster with its specific remapped fault list attached. Hybrid mode needs the same forward/backward pairs. - - Cache the timing per faulty cluster (forward/backward) alongside the baseline. - -4. **Annotate Pipeline Graph Nodes** - - When applying transformer timings in hierarchical/hybrid modes, assign the baseline duration to transformer nodes whose cluster is fault-free, and assign the corresponding per-cluster duration to nodes mapped to faulty hardware. This ensures only the affected stages slow down. - -5. **Pipeline Fault Injection** - - Feed the LP/DP-projected fault list into the pipeline AstraSim execution. Because the endpoints are already remapped to the LP/DP-only axis ordering, the existing hierarchical pipeline simulation will see the correct degraded links. diff --git a/LLM_execution.py b/LLM_execution.py index 579c3e1d..b2c1bea3 100644 --- a/LLM_execution.py +++ b/LLM_execution.py @@ -646,8 +646,8 @@ def __init__( self._pipeline_rank_layout: Dict[str, Any] = {} self._network_dimensions: Tuple[Any, ...] = tuple() self._axis_dimension_map: Dict[str, int] = {} - self._transformer_stage_faults: Dict[int, Tuple[Tuple[int, int, float], ...]] = {} - self._transformer_stage_timings: Dict[int, TransformerTimings] = {} + self._transformer_stage_dp_faults: Dict[Tuple[int, int], Tuple[Tuple[int, int, float], ...]] = {} + self._transformer_stage_timings: Dict[Tuple[int, int], TransformerTimings] = {} self._transformer_baseline_timings: Optional[TransformerTimings] = None self.no_data_parallel = bool(no_data_parallel) self._rank_layout = self._build_rank_layout_descriptor() @@ -693,8 +693,9 @@ def _safe_int(value: Any, default: int = 1) -> int: dim_axes = [str(axis).strip().lower() for axis in getattr(dim, "parallelisms", ())] declared = int(getattr(dim, "size", 1)) - if enforce_layout and declared > 1 and not first_active_checked: - canon = sorted(dim_axes) + axes_without_dp = [axis for axis in dim_axes if axis != "dp"] + if enforce_layout and declared > 1 and not first_active_checked and axes_without_dp: + canon = sorted(axes_without_dp) if canon != ["cp", "tp"] and canon != ["tp", "cp"] and (axis_sizes["tp"] > 1 or axis_sizes["cp"] > 1): raise ValueError( "For hierarchical/hybrid AstraSim modes, the first active network " @@ -703,8 +704,6 @@ def _safe_int(value: Any, default: int = 1) -> int: first_active_checked = True for name in dim_axes: - if name == "dp": - continue if name not in axis_sizes: raise ValueError( f"Unsupported parallelism axis '{name}' in network layout. " @@ -764,7 +763,7 @@ def _subset_layout(allowed: Sequence[str]) -> Optional[Dict[str, Any]]: # (DP replicas are handled externally). Store these subsets so callers # can attach the appropriate layout before invoking AstraSim. self._transformer_rank_layout = _subset_layout(["tp", "cp"]) - self._pipeline_rank_layout = _subset_layout(["lp"]) + self._pipeline_rank_layout = _subset_layout(["lp", "dp"]) return descriptor @@ -847,7 +846,7 @@ def _initialize_fault_mappings(self) -> None: ) self._fault_projections = self._build_fault_projections_from_space(self._fault_space) self._validate_fault_coverage() - self._transformer_stage_faults = self._build_transformer_stage_fault_map() + self._transformer_stage_dp_faults = self._build_transformer_stage_dp_fault_map() axis_order = self._rank_layout.get("axis_order", []) if isinstance(self._rank_layout, dict) else [] axis_sizes = self._rank_layout.get("axis_sizes", {}) if isinstance(self._rank_layout, dict) else {} self._log_fault_summary(axis_order, axis_sizes) @@ -917,16 +916,22 @@ def _axis_value_from_coords(self, coords: Tuple[Tuple[str, int], ...], axis: str if name == axis: return int(value) layout = self._rank_layout or {} - axis_sizes = layout.get("axis_sizes", {}) if isinstance(layout, dict) else {} + if isinstance(layout, dict): + layout_axes = layout.get("axis_order", []) + if axis not in layout_axes: + return 0 + axis_sizes = layout.get("axis_sizes", {}) + else: + axis_sizes = {} if axis_sizes.get(axis, 1) <= 1: return 0 raise ValueError(f"Axis '{axis}' not present in coordinate tuple for faulty link.") - def _build_transformer_stage_fault_map(self) -> Dict[int, Tuple[Tuple[int, int, float], ...]]: + def _build_transformer_stage_dp_fault_map(self) -> Dict[Tuple[int, int], Tuple[Tuple[int, int, float], ...]]: projection = self._fault_projection_for("transformer") if not projection: return {} - stage_faults: Dict[int, List[Tuple[int, int, float]]] = {} + stage_dp_faults: Dict[Tuple[int, int], List[Tuple[int, int, float]]] = {} for detail in projection.entries: src_stage = self._axis_value_from_coords(detail.src_coords, "lp") dst_stage = self._axis_value_from_coords(detail.dst_coords, "lp") @@ -935,8 +940,12 @@ def _build_transformer_stage_fault_map(self) -> Dict[int, Tuple[Tuple[int, int, "Transformer faulty link spans multiple pipeline stages; " "hierarchical execution requires faults limited to a single stage." ) - stage_faults.setdefault(src_stage, []).append(detail.remapped) - return {stage: tuple(links) for stage, links in stage_faults.items()} + src_dp = self._axis_value_from_coords(detail.src_coords, "dp") + dst_dp = self._axis_value_from_coords(detail.dst_coords, "dp") + affected_dp = {src_dp, dst_dp} + for dp_idx in affected_dp: + stage_dp_faults.setdefault((dp_idx, src_stage), []).append(detail.remapped) + return {key: tuple(links) for key, links in stage_dp_faults.items()} def _fault_projection_for(self, label: str) -> Optional[FaultProjectionResult]: return self._fault_projections.get(label) @@ -1194,7 +1203,7 @@ def _run_transformer_astrasim(self, mode: ExecutionMode) -> Optional[Transformer del mode # mode currently unused but kept for signature consistency if not self.transformer_forward_root and not self.transformer_backward_root: - if getattr(self, "_transformer_stage_faults", {}): + if getattr(self, "_transformer_stage_dp_faults", {}): raise ValueError("Transformer faults require transformer graph metadata, but none is available.") return None @@ -1223,16 +1232,16 @@ def _run_transformer_astrasim(self, mode: ExecutionMode) -> Optional[Transformer self.time_calc.transformer_astrasim_time_backward = baseline_timings.backward # Per-stage fault runs - stage_faults = getattr(self, "_transformer_stage_faults", {}) - for fault_index, (stage_id, fault_links) in enumerate(sorted(stage_faults.items())): - label = f"fault{fault_index}" + stage_dp_faults = getattr(self, "_transformer_stage_dp_faults", {}) + for fault_index, ((dp_idx, stage_id), fault_links) in enumerate(sorted(stage_dp_faults.items())): + label = f"fault{fault_index}_dp{dp_idx}_stage{stage_id}" stage_fwd_dir, stage_bwd_dir = self._transformer_artifact_dirs(label=label, persist=persist) stage_timings, _, _ = self._execute_transformer_run( stage_fwd_dir, stage_bwd_dir, faulty_links_override=fault_links, ) - self._transformer_stage_timings[stage_id] = stage_timings + self._transformer_stage_timings[(dp_idx, stage_id)] = stage_timings return baseline_timings @@ -1312,6 +1321,8 @@ def _apply_transformer_time(self, timings: TransformerTimings) -> None: else: roots = [self.pipeline_root] + dp_count = max(1, getattr(self.time_calc, "dp", 1)) + for root in roots: self._assign_transformer_durations( root, @@ -1319,15 +1330,17 @@ def _apply_transformer_time(self, timings: TransformerTimings) -> None: stage_timings, baseline_timings.forward, baseline_timings.backward, + dp_count, ) def _assign_transformer_durations( self, node: Any, visited: Set[int], - stage_timings: Dict[int, TransformerTimings], + stage_timings: Dict[Tuple[int, int], TransformerTimings], forward_default: float, backward_default: float, + dp_count: int, ) -> None: if node is None: return @@ -1344,11 +1357,26 @@ def _assign_transformer_durations( hw_stage = int(getattr(node, "hw_id", None)) except (TypeError, ValueError): hw_stage = None - timing_override = stage_timings.get(hw_stage) if hw_stage is not None else None + def _per_dp_durations(is_forward: bool) -> List[float]: + values: List[float] = [] + for dp_idx in range(dp_count): + timing_override = stage_timings.get((dp_idx, hw_stage)) if hw_stage is not None else None + if timing_override: + values.append(timing_override.forward if is_forward else timing_override.backward) + else: + default = forward_default if is_forward else backward_default + values.append(default) + return values + if getattr(node, "fwd", True): - node.duration = timing_override.forward if timing_override else forward_default + per_dp_values = _per_dp_durations(is_forward=True) + else: + per_dp_values = _per_dp_durations(is_forward=False) + + if dp_count > 1: + node.duration = tuple(per_dp_values) else: - node.duration = timing_override.backward if timing_override else backward_default + node.duration = per_dp_values[0] for child in getattr(node, "children", []): - self._assign_transformer_durations(child, visited, stage_timings, forward_default, backward_default) + self._assign_transformer_durations(child, visited, stage_timings, forward_default, backward_default, dp_count) diff --git a/astrasim_lib/executor.py b/astrasim_lib/executor.py index 42a107e6..c1dbf97c 100644 --- a/astrasim_lib/executor.py +++ b/astrasim_lib/executor.py @@ -674,6 +674,19 @@ def dfs(obj: Any) -> None: def rank_for(stage: int, dp_idx: int) -> int: return stage_to_ranks[stage][dp_idx] + def _compute_duration_seconds(task: Any, dp_idx: int) -> float: + """Return the per-DP duration (in seconds) for ``task``.""" + profile = getattr(task, "duration_profile", None) + if profile: + if len(profile) != dp_count: + raise ValueError( + f"Duration profile for node '{getattr(task, 'name', '')}' " + f"has length {len(profile)} but dp_count={dp_count}." + ) + return float(profile[dp_idx]) + duration_sec = getattr(task, "duration", 0.0) or 0.0 + return float(duration_sec) + # Step 3: Recover the multi-dimensional axis layout (tp/cp/lp, etc.) so that # communicator membership can be reconstructed deterministically later on. rank_layout = getattr(graph_root, "_astrasim_rank_layout", None) @@ -1339,7 +1352,7 @@ def ensure_pipeline(parent: Any, target: Any, dp_idx: int) -> int: if dep not in unique_deps: unique_deps.append(dep) - duration_sec = getattr(task, "duration", 0.0) or 0.0 + duration_sec = _compute_duration_seconds(task, dp_idx) duration_micros = int(round(duration_sec * 1e6)) if duration_sec else 0 node_id = trace.next_id comp_node = new_comp_node( diff --git a/simulate_LLM.py b/simulate_LLM.py index fbc2f620..353acc70 100644 --- a/simulate_LLM.py +++ b/simulate_LLM.py @@ -16,7 +16,7 @@ def __init__(self, name, op_id, hw_id, duration, fwd=True): self.name = name self.op_id = op_id self.hw_id = hw_id - self.duration = duration + self._duration_data = self._normalize_duration(duration) self.done = False self.finish_time = -1 self.parents = [] @@ -31,6 +31,39 @@ def add_child(self, obj): def __repr__(self): return f"Node({self.name},op={self.op_id},hw={self.hw_id},{round(self.duration, 2)})" + + @staticmethod + def _normalize_duration(value): + if isinstance(value, (list, tuple)): + entries = tuple(float(v) for v in value) + if not entries: + raise ValueError("Duration tuple must contain at least one entry.") + return entries + return float(value) + + @property + def duration(self): + data = self._duration_data + if isinstance(data, tuple): + return data[0] + return data + + @duration.setter + def duration(self, value): + self._duration_data = self._normalize_duration(value) + + @property + def duration_profile(self) -> Optional[Tuple[float, ...]]: + data = self._duration_data + if isinstance(data, tuple): + return data + return None + + def duration_storage(self): + data = self._duration_data + if isinstance(data, tuple): + return tuple(data) + return data class Data_batch: def __init__(self, name, batch_id, duration): @@ -204,7 +237,7 @@ def ensure_clone(obj: Any, cls: str) -> Any: return cached if isinstance(obj, Node): - clone = Node(obj.name, obj.op_id, obj.hw_id, obj.duration, fwd=obj.fwd) + clone = Node(obj.name, obj.op_id, obj.hw_id, obj.duration_storage(), fwd=obj.fwd) clone.memory = getattr(obj, "memory", 0) elif isinstance(obj, Edge): clone = Edge( @@ -380,7 +413,7 @@ def ensure_clone(obj: Any) -> Any: return cached if isinstance(obj, Node): - clone = Node(obj.name, obj.op_id, obj.hw_id, obj.duration, fwd=obj.fwd) + clone = Node(obj.name, obj.op_id, obj.hw_id, obj.duration_storage(), fwd=obj.fwd) clone.memory = getattr(obj, "memory", 0) elif isinstance(obj, Edge): clone = Edge( @@ -1516,11 +1549,35 @@ def visualize_graph(roots, filename="graph"): dot = Digraph(comment="Computation Graph", format="svg") visited = set() - def _format_duration(value: float) -> str: - ms = value * 1e3 - if abs(ms) > 1000: - return f"{value:.2f}s" - return f"{ms:.2f}ms" + def _format_duration(value: float, profile: Optional[Tuple[float, ...]] = None) -> str: + def _format_single(entry: float) -> str: + ms = entry * 1e3 + if abs(ms) > 1000: + return f"{entry:.2f}s" + return f"{ms:.2f}ms" + + if not profile: + return _format_single(value) + + groups: List[Dict[str, Any]] = [] + for idx, entry in enumerate(profile): + matched = False + for group in groups: + if math.isclose(entry, group["value"], rel_tol=1e-9, abs_tol=1e-12): + group["indices"].append(idx) + matched = True + break + if not matched: + groups.append({"value": entry, "indices": [idx]}) + + if len(groups) == 1: + return _format_single(groups[0]["value"]) + + parts = [] + for group in groups: + indices = ",".join(str(i) for i in group["indices"]) + parts.append(f"{{{indices}}}: {_format_single(group['value'])}") + return ", ".join(parts) def _node_color(node) -> str: if isinstance(node, Node): @@ -1541,17 +1598,19 @@ def _node_label(node) -> str: if isinstance(node, Data_batch): return f"{node.name}\n(batch_id={node.batch_id}, dur={_format_duration(node.duration)})" if isinstance(node, Node): + duration_display = _format_duration(node.duration, node.duration_profile) return ( f"{node.name}\n(op_id={node.op_id}, hw_id={node.hw_id}, " - f"dur={_format_duration(node.duration)})" + f"dur={duration_display})" ) if isinstance(node, Edge): + duration_display = _format_duration(node.duration) if getattr(node, "local_hw_id", None) is not None: return ( f"{node.name}\n(op_id={node.op_id}, local_hw_id={node.local_hw_id}, " - f"dur={_format_duration(node.duration)})" + f"dur={duration_display})" ) - return f"{node.name}\n(op_id={node.op_id}, dur={_format_duration(node.duration)})" + return f"{node.name}\n(op_id={node.op_id}, dur={duration_display})" return str(node) def _visit(node): From 10d4ac1394ce5d3489df6519a86ba2cc27943559 Mon Sep 17 00:00:00 2001 From: gkarfakis19 Date: Sat, 8 Nov 2025 03:32:42 -0800 Subject: [PATCH 28/77] gmap optimize 2D topologies implemented (not tested) --- .gitignore | 3 +- LLM_execution.py | 43 ++ astrasim_lib/executor.py | 163 ++++- astrasim_lib/gmap.py | 494 +++++++++++++++ config.py | 16 + configs/hardware-config/gmap_mesh2d_demo.yaml | 167 +++++ configs/model-config/gmap_mesh2d_demo.yaml | 20 + pyproject.toml | 1 + run_perf.py | 5 - time_calculation.py | 2 + tools/gmap_sweep.py | 582 ++++++++++++++++++ uv.lock | 14 + 12 files changed, 1496 insertions(+), 14 deletions(-) create mode 100644 astrasim_lib/gmap.py create mode 100644 configs/hardware-config/gmap_mesh2d_demo.yaml create mode 100644 configs/model-config/gmap_mesh2d_demo.yaml create mode 100644 tools/gmap_sweep.py diff --git a/.gitignore b/.gitignore index d4467d02..62f6ee97 100644 --- a/.gitignore +++ b/.gitignore @@ -11,4 +11,5 @@ dbg_hw_conf/ tools/fault_sweep_results* tools/*.png tools/*.tsv -tools/*.json \ No newline at end of file +tools/*.json +tools/gmap_artifacts/ \ No newline at end of file diff --git a/LLM_execution.py b/LLM_execution.py index b2c1bea3..963c00d7 100644 --- a/LLM_execution.py +++ b/LLM_execution.py @@ -651,6 +651,7 @@ def __init__( self._transformer_baseline_timings: Optional[TransformerTimings] = None self.no_data_parallel = bool(no_data_parallel) self._rank_layout = self._build_rank_layout_descriptor() + self._first_dim_optimize_cfg: Optional[Dict[str, Any]] = getattr(self, "_first_dim_optimize_cfg", None) self._fault_space: Optional[FaultSpace] = None self._fault_projections: Dict[str, FaultProjectionResult] = {} self._initialize_fault_mappings() @@ -662,6 +663,28 @@ def _build_rank_layout_descriptor(self) -> Dict[str, Any]: if not dimensions: return {} self._network_dimensions = tuple(dimensions) + optimize_cfg: Optional[Dict[str, Any]] = None + for idx, dim in enumerate(dimensions): + if getattr(dim, "optimize_2dmap", False): + if idx != 0: + raise ValueError("optimize_2dmap is only supported on the first network dimension.") + if optimize_cfg is not None: + raise ValueError("Multiple network dimensions requested optimize_2dmap; only one is supported.") + topo_type = getattr(dim, "topology_type", None) + if not topo_type: + raise ValueError("optimize_2dmap requires a topology type on the target dimension.") + if str(topo_type).strip().lower() == "kingmesh2d": + raise NotImplementedError("optimize_2dmap is not implemented for KingMesh2D.") + size_value = getattr(dim, "size", None) + if size_value is None: + raise ValueError("optimize_2dmap requires an explicit dimension size.") + optimize_cfg = { + "dimension_index": idx, + "topology": str(topo_type), + "size": int(size_value), + "parallelisms": tuple(getattr(dim, "parallelisms", ()) or ()), + } + self._first_dim_optimize_cfg = optimize_cfg def _safe_int(value: Any, default: int = 1) -> int: try: @@ -742,6 +765,16 @@ def _safe_int(value: Any, default: int = 1) -> int: "stage_span": span, } + return descriptor + + def _attach_optimize_hint(self, root: Any) -> None: + if root is None: + return + if self._first_dim_optimize_cfg: + setattr(root, "_optimize_2dmap", dict(self._first_dim_optimize_cfg)) + elif hasattr(root, "_optimize_2dmap"): + delattr(root, "_optimize_2dmap") + def _subset_layout(allowed: Sequence[str]) -> Optional[Dict[str, Any]]: subset = [axis for axis in axis_order if axis in allowed and axis_sizes.get(axis, 1) > 1] @@ -767,6 +800,14 @@ def _subset_layout(allowed: Sequence[str]) -> Optional[Dict[str, Any]]: return descriptor + def _attach_optimize_hint(self, root: Any) -> None: + if root is None: + return + if self._first_dim_optimize_cfg: + setattr(root, "_optimize_2dmap", dict(self._first_dim_optimize_cfg)) + elif hasattr(root, "_optimize_2dmap"): + delattr(root, "_optimize_2dmap") + def _log_fault_summary(self, axis_order: Sequence[str], axis_sizes: Mapping[str, int]) -> None: if not axis_order: return @@ -1054,6 +1095,7 @@ def _run_full_astrasim_hierarchical(self) -> ExecutionResult: setattr(self.pipeline_root, "_astrasim_rank_layout", pipeline_layout) elif hasattr(self.pipeline_root, "_astrasim_rank_layout"): delattr(self.pipeline_root, "_astrasim_rank_layout") + self._attach_optimize_hint(self.pipeline_root) per_rank_sec, max_sec = run_astra_simulation_only_onepath( self.pipeline_root, self.time_calc, @@ -1091,6 +1133,7 @@ def _run_full_astrasim_flattened(self) -> ExecutionResult: flattener._propagate_local_hw_ids(flattened_root) setattr(flattened_root, "_astrasim_rank_layout", self._rank_layout) + self._attach_optimize_hint(flattened_root) self.time_calc.flattened_pipeline_root = flattened_root if _env_flag("DEEPFLOW_VISUALIZE_GRAPHS") and self.pipeline_root is not None: filename = "/pipeline_graph_post_flatten_no_dp" if self.no_data_parallel else "/pipeline_graph_post_flatten" diff --git a/astrasim_lib/executor.py b/astrasim_lib/executor.py index c1dbf97c..e9d2acdb 100644 --- a/astrasim_lib/executor.py +++ b/astrasim_lib/executor.py @@ -28,6 +28,7 @@ from graphviz import Digraph from .config_generation import ASTRA_DEBUG, generate_astrasim_configs_from_hw +from . import gmap from .et_utils import ( chakra_decode, chakra_encode, @@ -99,9 +100,9 @@ def _clean_astrasim_artifacts(directory: str) -> None: _TP_MEMBERS_TO_GID: Dict[Tuple[int, ...], str] = {} -def _extract_axis_layout(rank_layout: Optional[Dict[str, Any]]) -> Tuple[List[str], Dict[str, int]]: +def _extract_axis_layout(rank_layout: Optional[Dict[str, Any]]) -> Tuple[List[str], Dict[str, int], Dict[str, int]]: if not isinstance(rank_layout, dict): - return [], {} + return [], {}, {} axis_order = list(rank_layout.get("axis_order", [])) raw_sizes = rank_layout.get("axis_sizes", {}) if isinstance(axis_order, tuple): @@ -113,7 +114,85 @@ def _extract_axis_layout(rank_layout: Optional[Dict[str, Any]]) -> Tuple[List[st axis_sizes[str(key)] = max(1, int(value)) except Exception: axis_sizes[str(key)] = 1 - return axis_order, axis_sizes + raw_strides = rank_layout.get("axis_strides", {}) + axis_strides: Dict[str, int] = {} + if isinstance(raw_strides, dict): + for key, value in raw_strides.items(): + try: + axis_strides[str(key)] = int(value) + except Exception: + axis_strides[str(key)] = 0 + else: + span = 1 + for axis in axis_order: + axis_strides[axis] = span + span *= axis_sizes.get(axis, 1) + return axis_order, axis_sizes, axis_strides + + +def _remap_stages_for_mapping( + *, + permutation: Sequence[int], + collector: Optional[gmap.GMapCollector], + stage_ids: List[int], + stage_to_ranks: Dict[int, List[int]], + stage_index: Dict[int, int], + rank_traces: Dict[int, "_RankTrace"], + rank_meta: Dict[int, Dict[str, int]], + dp_count: int, + num_stages: int, + output_dir: str, +) -> Tuple[List[int], Dict[int, int]]: + """Reorder stages so AstraSim ranks match the SCOTCH permutation. + + We apply the permutation *before* ET emission so the rest of the converter + remains unaware of the remap. Only stage ordering / rank bookkeeping is + touched here, which keeps the change localized and avoids renaming files + after the fact. + """ + + if collector is None or not permutation: + return stage_ids, stage_index + + if len(permutation) != collector.vertex_count: + raise ValueError("Permutation length does not match first-dimension vertex count.") + + def _sort_key(stage: int) -> Tuple[int, ...]: + coords = collector.stage_axis_coords.get(stage, {}) + higher = tuple(int(coords.get(axis, 0)) for axis in collector.replication_axes) + local_idx = collector.local_index_for_stage(stage) + target_idx = permutation[local_idx] + return higher + (target_idx,) + + ordered_stages = sorted(stage_ids, key=_sort_key) + new_stage_index = {stage: idx for idx, stage in enumerate(ordered_stages)} + + new_stage_to_ranks: Dict[int, List[int]] = {stage: [0] * dp_count for stage in ordered_stages} + new_rank_traces: Dict[int, _RankTrace] = {} + new_rank_meta: Dict[int, Dict[str, int]] = {} + + for stage in stage_ids: + for dp_idx in range(dp_count): + old_rank = stage_to_ranks[stage][dp_idx] + new_idx = new_stage_index[stage] + new_rank = dp_idx * num_stages + new_idx + trace = rank_traces[old_rank] + trace.rank = new_rank + trace.path = os.path.join(output_dir, f"llm_graph.{new_rank}.et") + new_rank_traces[new_rank] = trace + new_stage_to_ranks[stage][dp_idx] = new_rank + new_rank_meta[new_rank] = {"stage": stage, "dp": dp_idx} + + stage_to_ranks.clear() + stage_to_ranks.update(new_stage_to_ranks) + rank_traces.clear() + rank_traces.update(new_rank_traces) + rank_meta.clear() + rank_meta.update(new_rank_meta) + + stage_ids[:] = ordered_stages + stage_index = new_stage_index + return ordered_stages, stage_index def _compute_stage_axis_coords( @@ -690,9 +769,14 @@ def _compute_duration_seconds(task: Any, dp_idx: int) -> float: # Step 3: Recover the multi-dimensional axis layout (tp/cp/lp, etc.) so that # communicator membership can be reconstructed deterministically later on. rank_layout = getattr(graph_root, "_astrasim_rank_layout", None) - axis_order, axis_sizes = _extract_axis_layout(rank_layout) + axis_order, axis_sizes, axis_strides = _extract_axis_layout(rank_layout) stage_axis_coords = _compute_stage_axis_coords(stage_ids, axis_order, axis_sizes) axis_groups = _build_axis_groups(axis_order, axis_sizes, stage_axis_coords, stage_to_ranks) + if hasattr(graph_root, "_optimize_2dmap"): + optimize_cfg = getattr(graph_root, "_optimize_2dmap") + else: + optimize_cfg = None + collector = gmap.begin_collection(optimize_cfg, axis_sizes, stage_axis_coords, output_dir) # Step 4: Attribute each comm edge to the stage that owns it. This lets the # dependency walkers treat collectives the same way they treat compute nodes. @@ -738,6 +822,35 @@ def _compute_duration_seconds(task: Any, dp_idx: int) -> float: pipeline_edge_map: Dict[Tuple[Any, Any], Any] = {} + def _record_pipeline_edges_for_first_dim(collector_obj: gmap.GMapCollector) -> None: + if not collector_obj.include_pipeline: + return + for obj in all_objects: + if getattr(obj, "comm_type", None) != "pipeline": + continue + parents = getattr(obj, "parents", []) + src_stage = None + for parent in parents: + if hasattr(parent, "hw_id") and parent.hw_id is not None and int(parent.hw_id) >= 0: + src_stage = int(parent.hw_id) + break + if src_stage is None: + raise ValueError("Unable to resolve source stage for pipeline edge during gmap collection.") + children = getattr(obj, "children", []) + dst_stage = None + for child in children: + if hasattr(child, "hw_id") and child.hw_id is not None and int(child.hw_id) >= 0: + dst_stage = int(child.hw_id) + break + if dst_stage is None: + raise ValueError("Unable to resolve destination stage for pipeline edge during gmap collection.") + if src_stage == dst_stage: + continue + if not hasattr(obj, "comm_size_bytes"): + raise ValueError("Pipeline edge is missing comm_size_bytes for gmap collection.") + size_bytes = int(getattr(obj, "comm_size_bytes")) + collector_obj.record_pipeline(src_stage=src_stage, dst_stage=dst_stage, size_bytes=size_bytes) + # Step 5a: Walk backwards from a compute node to classify dependencies into # stage-local, cross-stage (pipeline), or collective inputs. def analyze_for_compute(node: Any) -> Tuple[Set[Any], Set[Any], Set[Any]]: @@ -978,7 +1091,7 @@ def analyze_for_collective(edge: Any) -> Tuple[Set[Any], Set[Any], Set[Any]]: collective_info: Dict[Any, Dict[str, Any]] = {} for edge, stage in edge_stage.items(): stage_deps, pipeline_deps, collective_deps = analyze_for_collective(edge) - collective_info[edge] = { + entry = { "stage": stage, "stage_deps": stage_deps, "pipeline_deps": pipeline_deps, @@ -989,6 +1102,20 @@ def analyze_for_collective(edge: Any) -> Tuple[Set[Any], Set[Any], Set[Any]]: "interconnect_type": getattr(edge, "comm_interconnect_type", None), "participants": int(getattr(edge, "participants", 0) or 0), } + collective_info[edge] = entry + if collector: + axis_name = entry["interconnect_type"] + if axis_name: + axis_norm = str(axis_name).strip().lower() + if axis_norm in collector.subset_axes: + should_double = entry["comm_type"] in (pb.ALL_REDUCE, pb.ALL_TO_ALL) + collector.record_collective( + axis=axis_norm, + stage_id=stage, + size_bytes=entry["size"], + participant_count=entry["participants"], + double_weight=should_double, + ) if debug_enabled: print("[ConverterDebug] === COMPUTE INFO (stage deps / pipeline deps / collectives) ===") @@ -1012,6 +1139,8 @@ def analyze_for_collective(edge: Any) -> Tuple[Set[Any], Set[Any], Set[Any]]: f"[ConverterDebug] collective {name}: stage={info['stage']}" f" stage_deps={stage_names} pipeline_deps={pipe_names} collectives={coll_names}" ) + if collector: + _record_pipeline_edges_for_first_dim(collector) # Step 6: Group collectives by label/axis so we can reconstruct communicator # memberships (tp/cp/lp) without relying on the original ordering. tp_collective_groups: Dict[str, List[Any]] = defaultdict(list) @@ -1048,6 +1177,7 @@ def analyze_for_collective(edge: Any) -> Tuple[Set[Any], Set[Any], Set[Any]]: # Ensure mapping present even for stages discovered only via collectives stage_idx = stage_index.setdefault(stage, len(stage_index)) stage_to_ranks[stage] = [dp_idx * num_stages + stage_idx for dp_idx in range(dp_count)] + stage_ids.append(stage) stage_tasks[stage].add(edge) stage_adj: Dict[int, Dict[Any, Set[Any]]] = {stage: defaultdict(set) for stage in stage_tasks} @@ -1280,10 +1410,27 @@ def ensure_pipeline(parent: Any, target: Any, dp_idx: int) -> int: return recv_id + mapping_result = gmap.finalize_collection(collector) + if mapping_result: + stage_ids, stage_index = _remap_stages_for_mapping( + permutation=mapping_result.permutation, + collector=collector, + stage_ids=stage_ids, + stage_to_ranks=stage_to_ranks, + stage_index=stage_index, + rank_traces=rank_traces, + rank_meta=rank_meta, + dp_count=dp_count, + num_stages=num_stages, + output_dir=output_dir, + ) + # Step 10: Emit ET nodes per stage (collectives first, then compute nodes), # iterating in the topological order computed above. - for stage in stage_order: - order = stage_order[stage] + for stage in stage_ids: + order = stage_order.get(stage) + if not order: + continue for task in order: if task in collective_info: info = collective_info[task] @@ -1579,7 +1726,7 @@ def run_astra_simulation_only_onepath( exit() rank_layout = getattr(fwdbwd_root, "_astrasim_rank_layout", None) - axis_order, axis_sizes = _extract_axis_layout(rank_layout) + axis_order, axis_sizes, _ = _extract_axis_layout(rank_layout) axes_filter = derive_axes_filter(axis_order, axis_sizes, dp_count) if not axes_filter: axes_filter = None diff --git a/astrasim_lib/gmap.py b/astrasim_lib/gmap.py new file mode 100644 index 00000000..1e97d494 --- /dev/null +++ b/astrasim_lib/gmap.py @@ -0,0 +1,494 @@ +"""First-dimension traffic aggregation and SCOTCH mapping helpers.""" + +from __future__ import annotations + +import math +import os +import tempfile +from dataclasses import dataclass +from typing import Dict, List, Mapping, Optional, Sequence, Tuple + +import numpy as np + +from .layout_utils import build_subset_strides + +import scotchpy +from scotchpy.common import proper_int as sct_i +from scotchpy.maporder import Mapping as ScotchMapping +from scotchpy import strat as scotch_strat + +_EDGE_WEIGHT_SCALE_BYTES = 1024 * 1024 * 1024 # convert bytes -> gigabytes for SCOTCH stability + + +@dataclass(frozen=True) +class MappingResult: + permutation: Tuple[int, ...] + graph_path: str + map_path: str + metrics_path: str + + +class GMapCollector: + """Accumulates per-pair byte weights for the first network dimension.""" + + def __init__( + self, + *, + topology: str, + subset_axes: Tuple[str, ...], + axis_sizes: Mapping[str, int], + subset_strides: Mapping[str, int], + stage_axis_coords: Mapping[int, Mapping[str, int]], + output_dir: str, + vertex_count: int, + include_pipeline: bool, + ) -> None: + if not subset_axes: + raise ValueError("optimize_2dmap requires at least one axis in the first network dimension") + + self.topology = str(topology).strip().lower() + self.subset_axes = tuple(subset_axes) + self.axis_sizes = dict(axis_sizes) + self.subset_strides = dict(subset_strides) + self.stage_axis_coords = stage_axis_coords + self.output_dir = output_dir + self.vertex_count = int(vertex_count) + self.include_pipeline = include_pipeline + self._edge_weights: Dict[Tuple[int, int], float] = {} + + missing = [axis for axis in self.subset_axes if axis not in self.axis_sizes] + if missing: + raise ValueError(f"Subset axes missing from axis_sizes: {missing}") + + product = 1 + for axis in self.subset_axes: + product *= int(self.axis_sizes[axis]) + if product != self.vertex_count: + raise ValueError( + f"Vertex count {self.vertex_count} does not match product of subset axis sizes {product}." + ) + + self.replication_axes: Tuple[str, ...] = tuple( + axis + for axis in self.axis_sizes + if axis not in self.subset_axes and int(self.axis_sizes[axis]) > 1 + ) + self._axis_groups = self._build_axis_groups() + self._canonical_signature = self._select_canonical_signature() + + # ------------------------------------------------------------------ + # Canonical cluster helpers (only one LP/DP slice contributes weights) + # ------------------------------------------------------------------ + + def _select_canonical_signature(self) -> Optional[Tuple[Tuple[str, int], ...]]: + if not self.replication_axes: + return None + zero_signature = tuple((axis, 0) for axis in self.replication_axes) + for coords in self.stage_axis_coords.values(): + if all(int(coords.get(axis, 0)) == 0 for axis in self.replication_axes): + return zero_signature + signatures = sorted( + { + tuple((axis, int(coords.get(axis, 0))) for axis in self.replication_axes) + for coords in self.stage_axis_coords.values() + } + ) + if not signatures: + raise ValueError("Unable to determine canonical cluster for optimize_2dmap.") + return signatures[0] + + def _is_canonical_stage(self, stage_id: int) -> bool: + if not self.replication_axes: + return True + coords = self.stage_axis_coords.get(stage_id) + if coords is None: + raise ValueError(f"Unknown stage id {stage_id} when checking canonical membership.") + signature = tuple((axis, int(coords.get(axis, 0))) for axis in self.replication_axes) + return signature == self._canonical_signature + + # ------------------------------------------------------------------ + # Axis indexing helpers + # ------------------------------------------------------------------ + + def _coord(self, coords: Mapping[str, int], axis: str) -> int: + if axis not in coords: + raise ValueError(f"Axis '{axis}' missing from coordinate tuple {coords}.") + return int(coords[axis]) + + def _build_axis_groups(self) -> Dict[str, Dict[Tuple[Tuple[str, int], ...], Tuple[int, ...]]]: + axis_groups: Dict[str, Dict[Tuple[Tuple[str, int], ...], List[int]]] = { + axis: {} for axis in self.subset_axes + } + for stage_id, coords in self.stage_axis_coords.items(): + for axis in self.subset_axes: + key = self._axis_key_entries(axis, coords) + axis_groups[axis].setdefault(key, []).append(stage_id) + normalized: Dict[str, Dict[Tuple[Tuple[str, int], ...], Tuple[int, ...]]] = {} + for axis, entries in axis_groups.items(): + normalized[axis] = { + key: tuple(sorted(stages)) + for key, stages in entries.items() + } + return normalized + + def local_index_for_stage(self, stage_id: int) -> int: + coords = self.stage_axis_coords.get(stage_id) + if coords is None: + raise ValueError(f"Unknown stage id {stage_id} when computing local index.") + local = 0 + for axis in self.subset_axes: + coord = self._coord(coords, axis) + stride = self.subset_strides[axis] + local += coord * stride + return local + + def coords_from_local_index(self, local_index: int) -> Dict[str, int]: + coords: Dict[str, int] = {} + for axis in reversed(self.subset_axes): + stride = self.subset_strides[axis] + size = int(self.axis_sizes[axis]) + if stride == 0: + coords[axis] = 0 + continue + coords[axis] = (local_index // stride) % size + return coords + + # ------------------------------------------------------------------ + # Byte aggregation + # ------------------------------------------------------------------ + + def _per_pair_bytes(self, total_bytes: int, participants: int, double_weight: bool) -> float: + if participants <= 1: + raise ValueError("participants must be > 1 to compute per-pair bytes") + if total_bytes < 0: + raise ValueError("total_bytes must be non-negative") + per_pair = float(total_bytes) / (participants * (participants - 1)) + if double_weight: + per_pair *= 2.0 + return per_pair + + def _accumulate_pairs(self, local_indices: Sequence[int], per_pair_bytes: float) -> None: + unique = sorted(set(int(idx) for idx in local_indices)) + if len(unique) <= 1 or per_pair_bytes <= 0: + return + for i, src in enumerate(unique): + for dst in unique[i + 1 :]: + key = (src, dst) + self._edge_weights[key] = self._edge_weights.get(key, 0.0) + per_pair_bytes + + def _axis_key_entries(self, axis: str, coords: Mapping[str, int]) -> Tuple[Tuple[str, int], ...]: + key_entries: List[Tuple[str, int]] = [] + for other_axis in self.subset_axes: + if other_axis == axis: + continue + key_entries.append((other_axis, self._coord(coords, other_axis))) + for repl_axis in self.replication_axes: + key_entries.append((repl_axis, self._coord(coords, repl_axis))) + return tuple(key_entries) + + def _participants_for_axis(self, axis: str, coords: Mapping[str, int]) -> Optional[Tuple[int, ...]]: + axis_map = self._axis_groups.get(axis) + if not axis_map: + return None + key = self._axis_key_entries(axis, coords) + return axis_map.get(key) + + def record_collective( + self, + *, + axis: str, + stage_id: int, + size_bytes: int, + participant_count: int, + double_weight: bool, + ) -> None: + if not self._is_canonical_stage(stage_id): + return + axis_norm = str(axis).strip().lower() + coords = self.stage_axis_coords.get(stage_id) + if coords is None: + raise ValueError(f"Unknown stage id {stage_id} when resolving participants.") + axis_candidates = [axis_norm] + [ax for ax in self.subset_axes if ax != axis_norm] + participants: Optional[Tuple[int, ...]] = None + chosen_axis = axis_norm + for candidate in axis_candidates: + group = self._participants_for_axis(candidate, coords) + if not group: + continue + if participant_count > 1 and len(group) == participant_count: + participants = group + chosen_axis = candidate + break + if participants is None: + participants = group + chosen_axis = candidate + if not participants: + raise ValueError(f"Failed to resolve participants for axis '{axis_norm}' at stage {stage_id}.") + axis_norm = chosen_axis + effective_participants = participant_count if participant_count > 1 else len(participants) + effective_participants = max(effective_participants, len(participants)) + if effective_participants <= 1: + return + local_indices = [self.local_index_for_stage(stage) for stage in participants] + per_pair = self._per_pair_bytes(size_bytes, effective_participants, double_weight) + self._accumulate_pairs(local_indices, per_pair) + + def record_pipeline(self, src_stage: int, dst_stage: int, size_bytes: int) -> None: + if not self.include_pipeline: + return + if not (self._is_canonical_stage(src_stage) and self._is_canonical_stage(dst_stage)): + return + if size_bytes < 0: + raise ValueError("Pipeline edge size must be non-negative") + local_indices = [self.local_index_for_stage(src_stage), self.local_index_for_stage(dst_stage)] + per_pair = self._per_pair_bytes(size_bytes, 2, False) + self._accumulate_pairs(local_indices, per_pair) + + # ------------------------------------------------------------------ + # Graph emission helpers + # ------------------------------------------------------------------ + + def emit_graph(self) -> str: + os.makedirs(self.output_dir, exist_ok=True) + path = os.path.join(self.output_dir, "first_dim_comm.grf") + adjacency: List[List[Tuple[int, int]]] = [[] for _ in range(self.vertex_count)] + for (src, dst), weight in self._edge_weights.items(): + weight /= _EDGE_WEIGHT_SCALE_BYTES + w_int = max(1, int(round(weight))) + adjacency[src].append((dst, w_int)) + adjacency[dst].append((src, w_int)) + num_arcs = sum(len(neigh) for neigh in adjacency) + with open(path, "w", encoding="ascii") as handle: + handle.write("0\n") + handle.write(f"{self.vertex_count} {num_arcs}\n") + handle.write("0 010\n") + for neighbors in adjacency: + neighbors.sort(key=lambda item: item[0]) + line_parts: List[str] = [str(len(neighbors))] + for neighbor, weight in neighbors: + line_parts.extend([str(weight), str(neighbor)]) + handle.write(" ".join(line_parts) + "\n") + return path + + def build_csr(self) -> Tuple[np.ndarray, np.ndarray, np.ndarray]: + nbrs = [[] for _ in range(self.vertex_count)] + weights = [[] for _ in range(self.vertex_count)] + for (src, dst), weight in self._edge_weights.items(): + if src == dst or weight <= 0: + continue + # print(f"src: {src}, dst: {dst}, weight: {weight}") + w_int = _scotch_weight(weight) + # new weight + # print(f"w_int: {w_int}") + nbrs[src].append(dst) + weights[src].append(w_int) + nbrs[dst].append(src) + weights[dst].append(w_int) + dtype = sct_i if sct_i is not None else np.int64 + xadj = np.zeros(self.vertex_count + 1, dtype=dtype) + adjncy: List[int] = [] + adjwgt: List[int] = [] + for idx in range(self.vertex_count): + xadj[idx + 1] = xadj[idx] + len(nbrs[idx]) + adjncy.extend(nbrs[idx]) + adjwgt.extend(weights[idx]) + return xadj, np.array(adjncy, dtype=dtype), np.array(adjwgt, dtype=dtype) + + +def begin_collection( + optimize_cfg: Optional[Mapping[str, object]], + axis_sizes: Mapping[str, int], + stage_axis_coords: Mapping[int, Mapping[str, int]], + output_dir: str, +) -> Optional[GMapCollector]: + if not optimize_cfg: + return None + subset_axes_raw = optimize_cfg.get("parallelisms") + if not subset_axes_raw: + raise ValueError("optimize_2dmap requires at least one parallelism axis in the first dimension.") + subset_axes: Tuple[str, ...] = tuple( + str(axis).strip().lower() + for axis in subset_axes_raw + if str(axis).strip() + ) + subset_strides = build_subset_strides(subset_axes, axis_sizes) + vertex_count = 1 + for axis in subset_axes: + if axis not in axis_sizes: + raise ValueError(f"Axis '{axis}' from optimize_2dmap is missing in axis_sizes.") + vertex_count *= int(axis_sizes[axis]) + declared_size = int(optimize_cfg.get("size", vertex_count)) + if declared_size != vertex_count: + raise ValueError( + f"optimize_2dmap size {declared_size} does not match product of subset axes {vertex_count}." + ) + topology = str(optimize_cfg.get("topology", "")).strip() + if not topology: + raise ValueError("optimize_2dmap requires a topology string.") + include_pipeline = "lp" in subset_axes + return GMapCollector( + topology=topology, + subset_axes=subset_axes, + axis_sizes=axis_sizes, + subset_strides=subset_strides, + stage_axis_coords=stage_axis_coords, + output_dir=output_dir, + vertex_count=vertex_count, + include_pipeline=include_pipeline, + ) + + +def finalize_collection(collector: Optional[GMapCollector]) -> Optional[MappingResult]: + if collector is None: + return None + graph_path = collector.emit_graph() + xadj, adjncy, adjwgt = collector.build_csr() + g = scotchpy.Graph() + g.build( + vertnbr=int(collector.vertex_count), + verttab=xadj, + vendtab=None, + velotab=None, + vlbltab=None, + edgenbr=int(len(adjncy)), + edgetab=adjncy, + edlotab=adjwgt, + ) + dims = _infer_target_dims(collector) + arch = _build_architecture(collector.topology, dims) + baseline_map = np.arange(collector.vertex_count, dtype=np.int64) + before_metrics = _collect_map_metrics(g, arch, baseline_map.copy(), run_map=False) + maptab = np.empty(collector.vertex_count, dtype=np.int64) + after_metrics = _collect_map_metrics(g, arch, maptab, run_map=True) + before_commexp, after_commexp = _commexpan_values(before_metrics, after_metrics) + should_apply_mapping = True + if before_commexp is not None and after_commexp is not None: + should_apply_mapping = after_commexp < before_commexp + map_path = os.path.join(collector.output_dir, "first_dim_comm.out") + final_map = maptab if should_apply_mapping else baseline_map.copy() + _write_mapping_file(map_path, final_map) + metrics_path = os.path.join(collector.output_dir, "first_dim_comm.metrics") + _write_metrics_report(metrics_path, before_metrics, after_metrics) + _maybe_print_commexpan_delta(before_commexp, after_commexp, should_apply_mapping) + if _env_truthy("DEEPFLOW_GMAP_ONLY"): + status = "applied" if should_apply_mapping else "skipped" + print( + "[GMapDebug] DEEPFLOW_GMAP_ONLY=1; " + f"{status} mapping. Artifacts: map={map_path}, metrics={metrics_path}" + ) + raise SystemExit(0) + permutation = tuple(int(val) for val in final_map) + return MappingResult( + permutation=permutation, + graph_path=graph_path, + map_path=map_path, + metrics_path=metrics_path, + ) + + +def _infer_target_dims(collector: GMapCollector) -> Tuple[int, int]: + total = int(collector.vertex_count) + if total <= 0: + raise ValueError("optimize_2dmap requires a positive vertex count.") + side = int(math.isqrt(total)) + if side * side != total: + raise ValueError( + f"optimize_2dmap first-dimension size {total} is not a perfect square; Mesh2D/Torus2D mapping now requires square layouts." + ) + return side, side + + +def _build_architecture(topology: str, dims: Tuple[int, int]): + dim_x, dim_y = dims + topo = topology.lower() + arch = scotchpy.Arch() + if topo == "mesh2d": + arch.mesh2(int(dim_x), int(dim_y)) + return arch + if topo == "torus2d": + arch.torus2(int(dim_x), int(dim_y)) + return arch + if topo == "kingmesh2d": + raise NotImplementedError("optimize_2dmap does not yet support KingMesh2D mappings.") + raise ValueError(f"Unsupported topology '{topology}' for optimize_2dmap.") + + +def _scotch_weight(weight_bytes: float) -> int: + if weight_bytes <= 0: + return 0 + scaled = weight_bytes / float(_EDGE_WEIGHT_SCALE_BYTES) + if scaled < 1.0: + scaled = 1.0 + return max(1, int(round(scaled))) + + +def _write_mapping_file(path: str, maptab: np.ndarray) -> None: + with open(path, "w", encoding="ascii") as handle: + handle.write(f"{len(maptab)}\n") + for idx, target in enumerate(maptab): + handle.write(f"{idx} {int(target)}\n") + + +def _write_metrics_report(path: str, before: str, after: str) -> None: + with open(path, "w", encoding="ascii") as handle: + handle.write("SCOTCH Mapping Metrics (graphMapView output)\n\n") + handle.write("Before:\n") + handle.write(before.rstrip() + "\n\n") + handle.write("After:\n") + handle.write(after.rstrip() + "\n") + + +def _collect_map_metrics( + graph: scotchpy.Graph, arch: scotchpy.Arch, parttab: np.ndarray, *, run_map: bool +) -> str: + mapping = ScotchMapping() + graph.map_init(mapping, arch, parttab=parttab) + try: + if run_map: + strategy = scotch_strat.Strat() + graph.map_compute(mapping, strategy) + with tempfile.TemporaryFile(mode="w+") as tmp: + graph.map_view(mapping, tmp) + tmp.seek(0) + return tmp.read() + finally: + graph.map_exit(mapping) + + +def _maybe_print_commexpan_delta( + before_val: Optional[float], after_val: Optional[float], applied: bool +) -> None: + if before_val is None or after_val is None or before_val == 0: + return + if not (_env_truthy("DEEPFLOW_VISUALIZE_GRAPHS") or _env_truthy("DEEPFLOW_GMAP_ONLY")): + return + delta_pct = (after_val - before_val) / before_val * 100.0 + status = "applied" if applied else "skipped" + print( + "[GMapDebug] CommExpan before " + f"{before_val:.6f} after {after_val:.6f} change {delta_pct:+.2f}% ({status})" + ) + + +def _extract_commexpan(metrics: str) -> Optional[float]: + for line in metrics.splitlines(): + if "CommExpan=" in line: + try: + prefix = line.split("CommExpan=")[1] + value_str = prefix.split("\t", 1)[0] + return float(value_str) + except (ValueError, IndexError): + return None + return None + + +def _commexpan_values(before: str, after: str) -> Tuple[Optional[float], Optional[float]]: + return _extract_commexpan(before), _extract_commexpan(after) + + +def _env_truthy(name: str) -> bool: + value = os.environ.get(name) + if value is None: + return False + normalized = value.strip().lower() + return normalized not in {"", "0", "false", "no"} diff --git a/config.py b/config.py index 7115fd64..b9b6e77e 100644 --- a/config.py +++ b/config.py @@ -354,6 +354,7 @@ class NetworkDimensionLayout: collective_override: Dict[str, str] = field(default_factory=dict) parallelisms: Tuple[str, ...] = field(default_factory=tuple) energy_per_bit: float = 0.0 + optimize_2dmap: bool = False @classmethod def from_raw( @@ -427,6 +428,20 @@ def from_raw( except (TypeError, ValueError) as exc: raise ValueError(f"network dimension '{label}' latency must be numeric") from exc + raw_optimize = topo_dict.get("optimize_2dmap", False) + if raw_optimize not in (True, False): + raise ValueError( + f"network dimension '{label}' topology optimize_2dmap must be a boolean when provided" + ) + optimize_2dmap = bool(raw_optimize) + if optimize_2dmap: + normalized_topo = topo_type.lower() + if normalized_topo not in {"mesh2d", "torus2d", "kingmesh2d"}: + raise ValueError( + f"network dimension '{label}' sets optimize_2dmap but topology type '{topo_type}'" + " is not Mesh2D/Torus2D/KingMesh2D" + ) + collectives_raw = raw.get("collective_override") if collectives_raw is None and "collectives" in raw: collectives_raw = raw.get("collectives") @@ -486,6 +501,7 @@ def from_raw( collective_override=collective_override, parallelisms=tuple(normalized_parallelisms), energy_per_bit=energy_per_bit, + optimize_2dmap=optimize_2dmap, ) @property diff --git a/configs/hardware-config/gmap_mesh2d_demo.yaml b/configs/hardware-config/gmap_mesh2d_demo.yaml new file mode 100644 index 00000000..25708ae0 --- /dev/null +++ b/configs/hardware-config/gmap_mesh2d_demo.yaml @@ -0,0 +1,167 @@ +sw_param: + kernel_launch_overhead: 5e-6 + dp_zero_stage: 0 + precision: + tensor_format: bf16 + mixed_precision: true + param_storage_mode: as_tensor_format + kv_cache: as_tensor_format + h2d_bandwidth: -1 + +tech_param: + core: + operating_frequency: 1.41e9 + num_bundles: 108 + nominal_flop_rate_per_mcu: 512 + nominal_energy_per_flop: 1.8e-13 + nominal_voltage: 0.8 + threshold_voltage: 0.6 + margin_voltage: 0.1 + operating_area_per_mcu: 1.0 + num_mcu_per_bundle: 4 + FMA_d1: 8 + FMA_d2: 4 + dataflow: best + util: 1.0 + DRAM: + size: 80 GB + bandwidth: 1986 GB + dynamic_energy_per_bit: 4.5e-12 + static_power_per_bit: 0.1 + area_per_bit: 0.1 + stack_capacity: 80 + area_per_stack: 0.1 + latency: 100e-9 + mem_ctrl_area: 0.1 + nominal_voltage: 0.8 + threshold_voltage: 0.6 + margin_voltage: 0.1 + num_links_per_mm: 1 + num_links_per_stack: 1 + max_voltage: 0.8 + util: 1 + SRAM-L2: + size: 40 MB + bandwidth: 7050 GB + dynamic_energy_per_bit: 130e-15 + static_power_per_bit: 0.1 + area_per_bit: 0.1 + bank_capacity: 1.0 + controller_area_per_link: 0.1 + latency: 0 + overhead: 0.1 + util: 1 + SRAM-L1: + size: 20736 KB + bandwidth: 18 TB + dynamic_energy_per_bit: 130e-15 + static_power_per_bit: 0.1 + area_per_bit: 0.1 + bank_capacity: 1.0 + controller_area_per_link: 0.1 + latency: 0 + overhead: 0.1 + util: 1 + SRAM-R: + size: 27 MB + bandwidth: 122 TB + dynamic_energy_per_bit: 110e-15 + static_power_per_bit: 0.1 + area_per_bit: 0.1 + bank_capacity: 1.0 + controller_area_per_link: 0.1 + latency: 0 + overhead: 0.1 + util: 1 + +memory_hierarchy: + l0: + type: SRAM-R + scope: mcu + l1: + type: SRAM-L1 + scope: mcu-bundle + l2: + type: SRAM-L2 + scope: global + l3: + type: DRAM + scope: global + +parallelism: + dp: 1 + tp: 4 + tp_sp: false + cp: 4 + lp: 1 + mb: 1 + +network: + dimensions: + - id: dim0 + label: mesh2d_cluster + size: auto + topology: + type: Torus2D + bandwidth: 100 GB + latency: 5e-6 + energy_per_bit: 8e-12 + util: 1.0 + optimize_2dmap: true + parallelisms: [tp, cp, lp] + - id: dim1 + label: pipeline_ring + size: auto + topology: + type: Ring + bandwidth: 25 GB + latency: 5e-6 + energy_per_bit: 8e-12 + util: 1.0 + parallelisms: [dp] + +power_breakdown: + TDP: 400 + core: 200 + DRAM: 100 + L2: 30 + L1: 20 + reg_mem: 10 + network: + inter_node: 5 + intra_node: 5 + +area_breakdown: + proc_chip_area_budget: 1.0 + core: 0.5 + DRAM: 0.1 + L2: 0.1 + L1: 0.1 + reg_mem: 0.1 + device_area_budget: 1.0 + network: + inter_node: 0.1 + intra_node: 0.1 + +perimeter_breakdown: + DRAM: 0.1 + inter_node: 0.1 + intra_node: 0.1 + +execution_backend: + model: astra + astra: + backend: astra + mode: full_astrasim_flattened + collectives: + all_gather: auto + all_reduce: auto + reduce_scatter: auto + all_to_all: auto + sys_options: + endpoint_delay: 5 + active_chunks_per_dimension: 8 + preferred_dataset_splits: 2 + +inference: + kvcache_type: hbm_only diff --git a/configs/model-config/gmap_mesh2d_demo.yaml b/configs/model-config/gmap_mesh2d_demo.yaml new file mode 100644 index 00000000..e1b4e406 --- /dev/null +++ b/configs/model-config/gmap_mesh2d_demo.yaml @@ -0,0 +1,20 @@ +model_param: + mode: "LLM" + run_type: "training" + tied_embeddings: true + model_type: "llama" + global_batch_size: 8 + gradient_accumulation_steps: 1 + seq_len: 32 + hidden_dim: 512 + attention: + attention_type: "mha" + num_heads: 8 + kv_heads: 8 + use_flashattention: false + attention_tile_size: 16 + num_experts: 1 + top_k: 1 + intermediate_size: 2048 + vocab_size: 8192 + num_layers: 2 diff --git a/pyproject.toml b/pyproject.toml index cde7ab12..ac04a1ec 100644 --- a/pyproject.toml +++ b/pyproject.toml @@ -15,6 +15,7 @@ dependencies = [ "python-dateutil==2.9.0.post0", "pytz==2025.2", "pyyaml==6.0.2", + "scotchpy64>=1.0.1", "seaborn>=0.13.2", "six==1.17.0", "tqdm>=4.67.1", diff --git a/run_perf.py b/run_perf.py index 246279e7..65ee8a9a 100755 --- a/run_perf.py +++ b/run_perf.py @@ -19,11 +19,6 @@ from time_calculation_LLM import TimeCalculationLLM from time_calculation_inf import TimeCalculationLLMInference -algByte = False # algorithmic ops false -proj = False # consider projection layer, turn off for end-2-end validation, as baeline model does not have projection layer -validating_v100 = True - - # Cache handling policy for AstraSim integration. # Options: "NO CACHE", "CACHE READONLY", "CACHE READWRITE" cache_handling = "CACHE READWRITE" diff --git a/time_calculation.py b/time_calculation.py index 3d9deb1b..a275dd20 100644 --- a/time_calculation.py +++ b/time_calculation.py @@ -33,6 +33,8 @@ def _astra_collective(self, kind: str, participants: int, size_bytes: int, axis: return 0.0 byte_count = int(math.ceil(size_bytes)) axes_filter = [str(axis).lower()] if axis else None + if os.environ.get("DEEPFLOW_ASTRA_SKIP_EXEC"): + return 0.0 # for collectives ONLY, we cannot use 2D topologies (Mesh2D, Torus2D, KingMesh2D) # transform them to their 1D equivalents (Ring, Mesh, HyperCube (?)) _, max_sec = run_cache_astrasim( diff --git a/tools/gmap_sweep.py b/tools/gmap_sweep.py new file mode 100644 index 00000000..e0a506ab --- /dev/null +++ b/tools/gmap_sweep.py @@ -0,0 +1,582 @@ +#!/usr/bin/env python3 +""" +GMap sweep utility. + +This tool explores tensor/context/pipeline parallelism combinations that live in +the first network dimension (TP/CP/LP) for a square, power-of-two GPU count. +Each configuration is evaluated by running DeepFlow with the +DEEPFLOW_GMAP_ONLY=1 environment variable so execution stops immediately after +SCOTCH emits the mapping artefacts. The script parses the resulting +`first_dim_comm.metrics` file (and the `[GMapDebug]` printouts) to recover the +before/after CommExpan values, computes the percentage delta, and produces both + a tab-separated report and a scatter plot coloured by % change. + +""" + +from __future__ import annotations + +import argparse +import copy +import io +import math +import os +import re +import shutil +import sys +import tempfile +from concurrent.futures import ProcessPoolExecutor, as_completed +from contextlib import redirect_stdout, redirect_stderr +from typing import Dict, Iterable, List, Optional, Sequence, Tuple + +import matplotlib + +matplotlib.use("Agg") +import matplotlib.pyplot as plt # noqa: E402 +import numpy as np # noqa: E402 +import yaml # noqa: E402 +from tqdm import tqdm # noqa: E402 + +sys.path.append(os.path.dirname(os.path.dirname(os.path.abspath(__file__)))) + +import config # noqa: E402 +from time_calculation_LLM import TimeCalculationLLM # noqa: E402 +from parallelism_sweep import set_astrasim_cache_mode # type: ignore + +# -----------------------------------------------------------------------------# +# Global configuration # +# -----------------------------------------------------------------------------# + +# Total GPU count in the first dimension (must be a square, power-of-two). +NUM_GPUS = 256 +NUM_WORKERS = 95 +ASTRA_CACHE_MODE = "NO_CACHE" +TOPOLOGY_VARIANTS = ("Mesh2D", "Torus2D") +# TOPOLOGY_VARIANTS = ("Mesh2D",) +# Hard bounds (inclusive) on each parallelism axis. +TP_BOUNDS = (1, 128) +CP_BOUNDS = (1, 128) +LP_BOUNDS = (1, 128) + +# -----------------------------------------------------------------------------# +# Helpers lifted / adapted from existing sweep tools # +# -----------------------------------------------------------------------------# + + +def read_yaml(path: str) -> Dict[str, object]: + with open(path, "r") as handle: + return yaml.safe_load(handle) + + +def determine_model_mode(model_config_path: str) -> str: + model_dict = read_yaml(model_config_path) + model_param = model_dict.get("model_param") or {} + mode = model_param.get("mode") + if not mode: + raise ValueError(f"model_param.mode must be defined in {model_config_path}") + return str(mode) + + +def make_temp_hw_config( + base_hw_dict: Dict[str, object], + parallel_settings: Dict[str, int], + hw_mutator=None, +): + updated = copy.deepcopy(base_hw_dict) + parallel_block = updated.setdefault("parallelism", {}) + for key, value in parallel_settings.items(): + parallel_block[key] = int(value) + + if hw_mutator: + hw_mutator(updated) + + tmp_file = tempfile.NamedTemporaryFile("w", suffix=".yaml", delete=False) + try: + yaml.safe_dump(updated, tmp_file, default_flow_style=False, sort_keys=False) + tmp_file.flush() + tmp_file.close() + hw_config = config.parse_config(tmp_file.name, config_type="hardware") + return hw_config + finally: + try: + tmp_file.close() + except Exception: + pass + try: + os.unlink(tmp_file.name) + except Exception: + pass + + +# -----------------------------------------------------------------------------# +# Core sweep logic # +# -----------------------------------------------------------------------------# + + +def is_power_of_two(value: int) -> bool: + return value > 0 and (value & (value - 1) == 0) + + +def is_perfect_square(value: int) -> bool: + root = math.isqrt(value) + return root * root == value + + +def enumerate_parallelisms(num_gpus: int) -> List[Dict[str, int]]: + limit = max(1, num_gpus // 2) + def _powers_of_two_between(lo: int, hi: int) -> List[int]: + lo = max(1, lo) + if hi < lo: + return [] + vals: List[int] = [] + val = 1 + while val < lo: + val <<= 1 + while val <= hi: + vals.append(val) + val <<= 1 + return vals + + tp_min, tp_max = TP_BOUNDS + cp_min, cp_max = CP_BOUNDS + lp_min, lp_max = LP_BOUNDS + tp_values = _powers_of_two_between(tp_min, min(tp_max, limit)) + cp_values = _powers_of_two_between(cp_min, min(cp_max, limit)) + lp_values = _powers_of_two_between(lp_min, min(lp_max, limit)) + combos: List[Dict[str, int]] = [] + for tp in tp_values: + for cp in cp_values: + for lp in lp_values: + if tp * cp * lp == num_gpus: + combos.append({"tp": tp, "cp": cp, "lp": lp}) + combos.sort(key=lambda item: (item["tp"], item["cp"], item["lp"])) + return combos + + +def update_first_dimension_size(hw_dict: Dict[str, object], num_gpus: int, topology: str) -> None: + network = hw_dict.get("network") + if not isinstance(network, dict): + return + dimensions = network.get("dimensions") + if isinstance(dimensions, list) and dimensions: + first = dimensions[0] + if isinstance(first, dict): + first["size"] = int(num_gpus) + topology_block = first.setdefault("topology", {}) + topology_block["type"] = topology + + +def parse_metrics_file(path: str) -> Tuple[Optional[float], Optional[float]]: + if not os.path.exists(path): + return None, None + section: Optional[str] = None + before = after = None + with open(path, "r") as handle: + for raw in handle: + line = raw.strip() + if not line: + continue + if line.startswith("Before"): + section = "before" + continue + if line.startswith("After"): + section = "after" + continue + if "CommExpan=" in line: + try: + value = float(line.split("CommExpan=")[1].split("\t", 1)[0]) + except (ValueError, IndexError): + value = None + if value is None: + continue + if section == "before" and before is None: + before = value + elif section == "after" and after is None: + after = value + return before, after + + +_COMMEXPAN_REGEX = re.compile( + r"CommExpan before\s+([0-9eE+\-\.]+)\s+after\s+([0-9eE+\-\.]+)" +) + + +def parse_stdout_for_commexpan(stdout_text: str) -> Tuple[Optional[float], Optional[float]]: + before = after = None + for line in stdout_text.splitlines(): + if "CommExpan" not in line: + continue + match = _COMMEXPAN_REGEX.search(line) + if not match: + continue + try: + before = float(match.group(1)) + after = float(match.group(2)) + except ValueError: + continue + return before, after + + +def compute_delta(before: Optional[float], after: Optional[float]) -> Optional[float]: + if before is None or after is None or before == 0: + return None + return (after - before) / before * 100.0 + + +def _format_delta(delta: Optional[float]) -> str: + if delta is None or delta != delta: + return "nan" + return f"{delta:+.2f}" + + +def _format_float(value: Optional[float]) -> str: + if value is None or value != value: + return "nan" + return f"{value:.6f}" + + +def run_single_configuration( + base_hw_dict: Dict[str, object], + model_config, + mode: str, + parallel_settings: Dict[str, int], + num_gpus: int, + keep_artifacts: bool, + topology: str, + capture_output: bool = True, +) -> Dict[str, object]: + hw_config = make_temp_hw_config( + base_hw_dict, + parallel_settings, + hw_mutator=lambda cfg: update_first_dimension_size(cfg, num_gpus, topology), + ) + if keep_artifacts: + base_dir = os.path.join("tools", "gmap_artifacts") + os.makedirs(base_dir, exist_ok=True) + dir_name = "gmap_{topology}_tp{tp}_cp{cp}_lp{lp}".format( + topology=topology.lower(), + tp=parallel_settings["tp"], + cp=parallel_settings["cp"], + lp=parallel_settings["lp"], + ) + temp_dir = os.path.join(base_dir, dir_name) + shutil.rmtree(temp_dir, ignore_errors=True) + os.makedirs(temp_dir, exist_ok=True) + else: + temp_dir = tempfile.mkdtemp(prefix="gmap_sweep_") + prev_flag = os.environ.get("DEEPFLOW_GMAP_ONLY") + prev_persist = os.environ.get("DEEPFLOW_PERSIST_ASTRASIM_ARTIFACTS") + os.environ["DEEPFLOW_GMAP_ONLY"] = "1" + if keep_artifacts: + os.environ["DEEPFLOW_PERSIST_ASTRASIM_ARTIFACTS"] = "1" + else: + os.environ.pop("DEEPFLOW_PERSIST_ASTRASIM_ARTIFACTS", None) + stdout_buffer = io.StringIO() + try: + calculator = TimeCalculationLLM(hw_config, model_config, mode, output_dir=temp_dir) + if capture_output: + with redirect_stdout(stdout_buffer), redirect_stderr(stdout_buffer): + try: + calculator.calc_time_llm() + except SystemExit as exc: + if exc.code not in (0, None): + raise + else: + try: + calculator.calc_time_llm() + except SystemExit as exc: + if exc.code not in (0, None): + raise + finally: + if prev_flag is None: + os.environ.pop("DEEPFLOW_GMAP_ONLY", None) + else: + os.environ["DEEPFLOW_GMAP_ONLY"] = prev_flag + if prev_persist is None: + os.environ.pop("DEEPFLOW_PERSIST_ASTRASIM_ARTIFACTS", None) + else: + os.environ["DEEPFLOW_PERSIST_ASTRASIM_ARTIFACTS"] = prev_persist + stdout_text = stdout_buffer.getvalue() if capture_output else "" + metrics_path = os.path.join(temp_dir, "first_dim_comm.metrics") + before, after = parse_metrics_file(metrics_path) + if before is None or after is None: + alt_before, alt_after = parse_stdout_for_commexpan(stdout_text) + if before is None: + before = alt_before + if after is None: + after = alt_after + delta_pct = compute_delta(before, after) + applied = bool( + before is not None and after is not None and before > 0 and after < before + ) + artifact_dir = temp_dir if keep_artifacts else None + if not keep_artifacts: + shutil.rmtree(temp_dir, ignore_errors=True) + else: + print(f"Saved artifacts to {temp_dir}") + return { + "parallelism": dict(parallel_settings), + "before": before, + "after": after, + "delta_pct": delta_pct, + "applied": applied, + "stdout": stdout_text, + "artifact_dir": artifact_dir, + "topology": topology, + } + + +# -----------------------------------------------------------------------------# +# Multiprocessing glue # +# -----------------------------------------------------------------------------# + +_GLOBAL_HW_DICT = None +_GLOBAL_MODEL_CONFIG = None +_GLOBAL_MODE = None +_GLOBAL_NUM_GPUS = None +_GLOBAL_KEEP_ARTIFACTS = False + + +def _worker_init(hw_dict, model_config_path, mode, num_gpus, keep_artifacts): + global _GLOBAL_HW_DICT, _GLOBAL_MODEL_CONFIG, _GLOBAL_MODE, _GLOBAL_NUM_GPUS, _GLOBAL_KEEP_ARTIFACTS + _GLOBAL_HW_DICT = hw_dict + _GLOBAL_MODE = mode + _GLOBAL_NUM_GPUS = num_gpus + _GLOBAL_KEEP_ARTIFACTS = bool(keep_artifacts) + _GLOBAL_MODEL_CONFIG = config.parse_config(model_config_path, config_type=mode) + + +def _worker_task(parallel_items: Tuple[Tuple[str, int], ...]) -> Dict[str, object]: + parallel_settings: Dict[str, int] = {} + topology: Optional[str] = None + for key, value in parallel_items: + if key == "__topology": + topology = str(value) + else: + parallel_settings[key] = int(value) + if topology is None: + topology = TOPOLOGY_VARIANTS[0] + try: + result = run_single_configuration( + _GLOBAL_HW_DICT, + _GLOBAL_MODEL_CONFIG, + _GLOBAL_MODE, + parallel_settings, + _GLOBAL_NUM_GPUS, + _GLOBAL_KEEP_ARTIFACTS, + topology, + ) + return {"status": "ok", "result": result} + except Exception as exc: + return { + "status": "error", + "error": str(exc), + "parallelism": parallel_settings, + "topology": topology, + } + + +# -----------------------------------------------------------------------------# +# Plotting / reporting # +# -----------------------------------------------------------------------------# + + +def plot_results(results: Sequence[Dict[str, object]], num_gpus: int, topology: str, path: str) -> None: + if not results: + return + xs = np.arange(len(results)) + ys = np.array( + [ + entry["delta_pct"] if entry["delta_pct"] == entry["delta_pct"] else float("nan") + for entry in results + ], + dtype=float, + ) + labels = [ + "tp={tp}, cp={cp}, lp={lp}".format( + tp=entry["parallelism"]["tp"], + cp=entry["parallelism"]["cp"], + lp=entry["parallelism"]["lp"], + ) + for entry in results + ] + applied_flags = [bool(entry.get("applied")) for entry in results] + applied_x = [x for x, flag in zip(xs, applied_flags) if flag] + applied_y = [y for y, flag in zip(ys, applied_flags) if flag] + skipped_x = [x for x, flag in zip(xs, applied_flags) if not flag] + skipped_y = [y for y, flag in zip(ys, applied_flags) if not flag] + + plt.figure(figsize=(12, 6)) + if applied_x: + plt.scatter(applied_x, applied_y, c="tab:green", s=80, label="Applied", edgecolors="black", linewidths=0.5) + if skipped_x: + plt.scatter(skipped_x, skipped_y, c="tab:red", s=80, label="Skipped", edgecolors="black", linewidths=0.5) + plt.axhline(0.0, color="gray", linestyle="--", linewidth=1) + plt.ylabel("Δ CommExpan (%) (negative is better)") + plt.xlabel("Parallelism configuration") + plt.title(f"GMap sweep ({topology}, num_gpus={num_gpus})") + plt.grid(alpha=0.3, axis="y") + plt.xticks(xs, labels, rotation=45, ha="right") + plt.legend() + plt.tight_layout() + plt.savefig(path, dpi=200) + plt.close() + print(f"Saved scatter plot to {path}") + + +def write_report(results: Sequence[Dict[str, object]], path: str, topology: str) -> None: + header = [ + "tp", + "cp", + "lp", + "before_commexpan", + "after_commexpan", + "delta_percent", + "applied", + "artifact_dir", + "topology", + ] + with open(path, "w") as handle: + handle.write("\t".join(header) + "\n") + for entry in results: + parallelism = entry["parallelism"] + row = [ + str(parallelism["tp"]), + str(parallelism["cp"]), + str(parallelism["lp"]), + _format_float(entry.get("before")), + _format_float(entry.get("after")), + _format_delta(entry.get("delta_pct")), + "yes" if entry.get("applied") else "no", + entry.get("artifact_dir") or "", + topology, + ] + handle.write("\t".join(row) + "\n") + print(f"Wrote report to {path}") + + +# -----------------------------------------------------------------------------# +# CLI # +# -----------------------------------------------------------------------------# + + +def parse_args(): + parser = argparse.ArgumentParser(description="Sweep TP/CP/LP combinations for GMap cost deltas.") + parser.add_argument("--hardware-config", default="configs/hardware-config/gmap_mesh2d_demo.yaml", help="Base hardware config path.") + parser.add_argument("--model-config", default="configs/model-config/Llama3.1-405B.yaml", help="Model config path.") + parser.add_argument("--max-workers", type=int, default=None, help="Maximum concurrent worker processes.") + parser.add_argument("--plot-output", default="tools/gmap_sweep.png", help="Path for the scatter plot.") + parser.add_argument("--report-output", default="tools/gmap_sweep.tsv", help="Path for the TSV report.") + parser.add_argument("--keep-artifacts", action="store_true", help="Preserve per-run artefacts (first_dim_comm.*) on disk.") + return parser.parse_args() + + +def main(): + args = parse_args() + set_astrasim_cache_mode(ASTRA_CACHE_MODE) + if args.max_workers is None or args.max_workers <= 0: + args.max_workers = max(1, (os.cpu_count() or 2) - 1) + num_gpus = int(NUM_GPUS) + if not is_power_of_two(num_gpus): + raise ValueError("num_gpus must be a power of two.") + if not is_perfect_square(num_gpus): + raise ValueError("num_gpus must be a perfect square to satisfy the Mesh2D constraint.") + + base_hw_dict = read_yaml(args.hardware_config) + mode = determine_model_mode(args.model_config) + cases = enumerate_parallelisms(num_gpus) + if not cases: + print("No TP/CP/LP combinations satisfy the constraints.") + return + print(f"Evaluating {len(cases)} configuration(s) for num_gpus={num_gpus}") + + def _suffix_path(base: str, suffix: str) -> str: + root, ext = os.path.splitext(base) + return f"{root}_{suffix.lower()}{ext or ''}" + + task_specs: List[Dict[str, object]] = [] + for topology in TOPOLOGY_VARIANTS: + for combo in cases: + entry = dict(combo) + entry["__topology"] = topology + task_specs.append(entry) + + results_by_topology: Dict[str, List[Dict[str, object]]] = {topo: [] for topo in TOPOLOGY_VARIANTS} + errors_by_topology: Dict[str, List[str]] = {topo: [] for topo in TOPOLOGY_VARIANTS} + + if len(task_specs) == 1: + spec = task_specs[0] + topology = spec["__topology"] + combo = {k: v for k, v in spec.items() if k != "__topology"} + model_config_obj = config.parse_config(args.model_config, config_type=mode) + try: + result = run_single_configuration( + base_hw_dict, + model_config_obj, + mode, + combo, + num_gpus, + args.keep_artifacts, + topology, + capture_output=False, + ) + results_by_topology[topology].append(result) + except Exception as exc: + errors_by_topology[topology].append(f"{combo}: {exc}") + else: + args.max_workers = NUM_WORKERS + max_workers = args.max_workers if args.max_workers and args.max_workers > 0 else 1 + worker_count = min(max_workers, max(1, (os.cpu_count() or 1) - 1), len(task_specs)) + task_items = [tuple(sorted(spec.items())) for spec in task_specs] + with ProcessPoolExecutor( + max_workers=worker_count, + initializer=_worker_init, + initargs=(base_hw_dict, args.model_config, mode, num_gpus, args.keep_artifacts), + ) as executor: + futures = {executor.submit(_worker_task, items): items for items in task_items} + with tqdm(total=len(task_items), desc="Evaluating", unit="config") as progress: + for future in as_completed(futures): + progress.update(1) + outcome = future.result() + topo = outcome.get("result", {}).get("topology") if outcome.get("status") == "ok" else outcome.get("topology") + if topo is None: + topo = TOPOLOGY_VARIANTS[0] + if outcome["status"] == "ok": + results_by_topology[topo].append(outcome["result"]) + else: + errors_by_topology[topo].append(f"{outcome['parallelism']}: {outcome.get('error')}") + + for topology in TOPOLOGY_VARIANTS: + print(f"\n=== Topology: {topology} ===") + results = results_by_topology[topology] + errors = errors_by_topology[topology] + if errors: + print("Encountered errors:") + for msg in errors: + print(" -", msg) + if not results: + print("No successful evaluations for this topology; skipping.") + continue + results.sort(key=lambda entry: (entry["parallelism"]["tp"], entry["parallelism"]["cp"], entry["parallelism"]["lp"])) + best = min( + (entry for entry in results if entry.get("delta_pct") == entry.get("delta_pct")), + key=lambda e: e.get("delta_pct"), + default=None, + ) + if best: + best_delta = best.get("delta_pct") + print( + "Best improvement: tp={tp}, cp={cp}, lp={lp}, Δ={delta}".format( + tp=best["parallelism"]["tp"], + cp=best["parallelism"]["cp"], + lp=best["parallelism"]["lp"], + delta=_format_delta(best_delta), + ) + ) + report_path = _suffix_path(args.report_output, topology) + plot_path = _suffix_path(args.plot_output, topology) + write_report(results, report_path, topology) + plot_results(results, num_gpus, topology, plot_path) + + +if __name__ == "__main__": + main() diff --git a/uv.lock b/uv.lock index 01fc09c5..8eb1a13b 100644 --- a/uv.lock +++ b/uv.lock @@ -133,6 +133,7 @@ dependencies = [ { name = "python-dateutil" }, { name = "pytz" }, { name = "pyyaml" }, + { name = "scotchpy64" }, { name = "seaborn" }, { name = "six" }, { name = "tqdm" }, @@ -151,6 +152,7 @@ requires-dist = [ { name = "python-dateutil", specifier = "==2.9.0.post0" }, { name = "pytz", specifier = "==2025.2" }, { name = "pyyaml", specifier = "==6.0.2" }, + { name = "scotchpy64", specifier = ">=1.0.1" }, { name = "seaborn", specifier = ">=0.13.2" }, { name = "six", specifier = "==1.17.0" }, { name = "tqdm", specifier = ">=4.67.1" }, @@ -687,6 +689,18 @@ wheels = [ { url = "https://files.pythonhosted.org/packages/fa/de/02b54f42487e3d3c6efb3f89428677074ca7bf43aae402517bc7cca949f3/PyYAML-6.0.2-cp313-cp313-win_amd64.whl", hash = "sha256:8388ee1976c416731879ac16da0aff3f63b286ffdd57cdeb95f3f2e085687563", size = 156446 }, ] +[[package]] +name = "scotchpy64" +version = "1.0.1" +source = { registry = "https://pypi.org/simple" } +wheels = [ + { url = "https://files.pythonhosted.org/packages/bd/b3/5e171536989e4fa154b22a480f256dd5fada8d9ce47614c77b747adb250c/scotchpy64-1.0.1-cp311-cp311-manylinux2014_x86_64.manylinux_2_17_x86_64.manylinux_2_28_x86_64.whl", hash = "sha256:78b17354dafe8d99471ff01fc54ce2d26fca2061ea27db95c8f8bb73813fc579", size = 2123702 }, + { url = "https://files.pythonhosted.org/packages/5f/be/3b344a1beba0abd011d0193f151d510a286787d3910c3ae3d7940dd70931/scotchpy64-1.0.1-cp312-cp312-manylinux2014_x86_64.manylinux_2_17_x86_64.manylinux_2_28_x86_64.whl", hash = "sha256:87705f9c4105740306d5048c7c44a51e62b2f03c9a95fa7610d6a3386989e8c5", size = 2123702 }, + { url = "https://files.pythonhosted.org/packages/3b/ea/12491425a0148b8b0ccbe3694a24d2b498a99c21084c8cec09e3e94dd78f/scotchpy64-1.0.1-cp313-cp313-manylinux2014_x86_64.manylinux_2_17_x86_64.manylinux_2_28_x86_64.whl", hash = "sha256:e2ea7f50b3a7797e0cd71537b2c072114dd9bb6ca48dc5f826d555370eb3150a", size = 2123703 }, + { url = "https://files.pythonhosted.org/packages/82/cd/b2f1f81a7adbec6e4c70cea0b518f89c5b2650e2fdc669aed380bcbb2bb4/scotchpy64-1.0.1-cp314-cp314-manylinux2014_x86_64.manylinux_2_17_x86_64.manylinux_2_28_x86_64.whl", hash = "sha256:3cd52831289d592cca3ca9ec4093467855f3c77d0af0966ba182fe49ceadc4d2", size = 2123703 }, + { url = "https://files.pythonhosted.org/packages/e4/f0/25613cd7f193b23dd5f886555b04012816fcebe0589c39d51e75f9fcb687/scotchpy64-1.0.1-cp314-cp314t-manylinux2014_x86_64.manylinux_2_17_x86_64.manylinux_2_28_x86_64.whl", hash = "sha256:e0c8c77a8aa5cc75d59b40cfb47f02814c5bdfb605fc82deaa57cac4f486f125", size = 2123708 }, +] + [[package]] name = "seaborn" version = "0.13.2" From b7f519a467082bdec97ab2913254560df5b613d4 Mon Sep 17 00:00:00 2001 From: gkarfakis19 Date: Sat, 8 Nov 2025 16:28:49 -0800 Subject: [PATCH 29/77] enabled pipeline parallelism for non-divisor num layers/pp --- simulate_LLM.py | 36 +++++++++++++++++++++++++++++++----- 1 file changed, 31 insertions(+), 5 deletions(-) diff --git a/simulate_LLM.py b/simulate_LLM.py index 353acc70..81cdf694 100644 --- a/simulate_LLM.py +++ b/simulate_LLM.py @@ -133,18 +133,44 @@ def __init__( self.num_batch = self.misc_metadata.get("num_batch", 0) self.num_layer = self.misc_metadata.get("num_layer", 0) self.layer_per_device = self.num_layer / self.lp if self.lp else self.num_layer - # if self layers per device is not an integer, raise Exception complaining about divisors! - if self.layer_per_device != int(self.layer_per_device): - raise Exception(f"The number of layers {self.num_layer} is not divisible by the degree of pipeline parallelism {self.lp}") self.all_reduce = self.misc_metadata.get("all_reduce", "every layer") self.transformer_cfg = self.comp_times.get("transformer", {}) self.T_grad_transformer = self.comp_times.get("grad_transformer", 0.0) + self.layers_per_stage = self._compute_layers_per_stage() + self.layer_to_stage = self._build_layer_to_stage() def _time(self, key: str, default: float = 0.0) -> float: value = self.comp_times.get(key) return float(value) if value is not None else default + def _compute_layers_per_stage(self) -> List[int]: + stage_count = max(1, int(self.lp) if self.lp else 1) + if self.num_layer <= 0: + return [0 for _ in range(stage_count)] + base = self.num_layer // stage_count + remainder = self.num_layer % stage_count + layers_per_stage: List[int] = [] + for stage_idx in range(stage_count): + count = base + (1 if stage_idx < remainder else 0) + layers_per_stage.append(count) + return layers_per_stage + + def _build_layer_to_stage(self) -> List[int]: + mapping: List[int] = [] + for stage_idx, count in enumerate(self.layers_per_stage): + if count <= 0: + continue + mapping.extend([stage_idx] * count) + return mapping + + def _stage_for_layer(self, layer_idx: int) -> int: + if layer_idx < 0 or layer_idx >= self.num_layer: + raise IndexError(f"Layer index {layer_idx} is out of bounds for {self.num_layer} layers.") + if self.layer_to_stage: + return self.layer_to_stage[layer_idx] + return 0 + def create_comm_edge(self, name, op_id, comm_key, is_dp=False, local_hw_id=None): """Create a communication edge with optional local computation node. @@ -622,7 +648,7 @@ def attach_parallel_edge(target, gather_edge, skip_non_comm_children=None, skip_ layer_exit_nodes[b] = [] for l in range(self.num_layer): - hw_id = min(l // self.layer_per_device , self.lp - 1)#assign hw_id for transformer + hw_id = self._stage_for_layer(l) transformer_node = Node(f"transformer_layer{l}", op_id, hw_id, transformer_f_time) transformer_node.micro_batch_index = b transformer_node.layer_index = l @@ -711,7 +737,7 @@ def attach_parallel_edge(target, gather_edge, skip_non_comm_children=None, skip_ for l in reversed(range(self.num_layer)): - hw_id = min(l // self.layer_per_device , self.lp - 1) + hw_id = self._stage_for_layer(l) transformer_node_b = Node(f"transformer_layer{l}_b", op_id, hw_id, transformer_b_time, fwd=False) transformer_node_b.micro_batch_index = b transformer_node_b.layer_index = l From adf3d2983dd3df2b0c12fdb40024977b630f46c0 Mon Sep 17 00:00:00 2001 From: gkarfakis19 Date: Sat, 8 Nov 2025 16:45:12 -0800 Subject: [PATCH 30/77] removed network/system cache functionality for astrasim integration to fix hybrid hierarchy bug --- astrasim_lib/config_generation.py | 61 ++++++------------------------- astrasim_lib/executor.py | 1 - astrasim_lib/integration.py | 51 ++++++++++++++++++++++++-- 3 files changed, 59 insertions(+), 54 deletions(-) diff --git a/astrasim_lib/config_generation.py b/astrasim_lib/config_generation.py index 1a9ccf64..806503d1 100644 --- a/astrasim_lib/config_generation.py +++ b/astrasim_lib/config_generation.py @@ -3,6 +3,7 @@ from __future__ import annotations import os +import uuid from typing import Any, Dict, List, Optional, Sequence, Tuple from types import SimpleNamespace @@ -17,15 +18,9 @@ ASTRA_DEBUG = False -_NET_YAML_CACHE: Dict[tuple, str] = {} -_JSON_WRITTEN_BY_NPUS: set[object] = set() - -def _save_json(path: str, data: Dict[str, Any], npus_key: Optional[int] = None) -> None: - """Write ``data`` to ``path`` once per ``npus_key`` per process.""" - key = npus_key if npus_key is not None else path - if key in _JSON_WRITTEN_BY_NPUS and os.path.exists(path): - return +def _save_json(path: str, data: Dict[str, Any]) -> None: + """Write ``data`` to ``path``.""" os.makedirs(os.path.dirname(path), exist_ok=True) tmp_path = path + ".tmp" with open(tmp_path, "w", encoding="utf-8") as handle: @@ -33,7 +28,6 @@ def _save_json(path: str, data: Dict[str, Any], npus_key: Optional[int] = None) _json.dump(data, handle, indent=2) os.replace(tmp_path, path) - _JSON_WRITTEN_BY_NPUS.add(key) def _gbps_from_bps(bps: float) -> float: @@ -114,7 +108,7 @@ def generate_astrasim_configs_from_hw( axes_filter: Optional[Sequence[str]] = None, transform_2d_to_1d: bool = False, faulty_links_override: Optional[Sequence[Tuple[int, int, float]]] = None, - reuse_cached_network: bool = True, + ephemeral_outputs: bool = False, ) -> Dict[str, str]: """Write AstraSim network/system configs derived from ``hw_obj``.""" if npus_count is None: @@ -311,8 +305,9 @@ def _safe_int(value: Any, default: int = 1) -> int: signature_parts = [str(size) for _dim, _axes, size, _ in selected_dims] dim_signature = "_".join(signature_parts) if signature_parts else f"{target}" - net_yaml = os.path.join(out_dir, f"network_analytical_{dim_signature}.yml") - sys_json = os.path.join(out_dir, f"system_native_collectives_{dim_signature}.json") + unique_suffix = f"_{uuid.uuid4().hex}" if ephemeral_outputs else "" + net_yaml = os.path.join(out_dir, f"network_analytical_{dim_signature}{unique_suffix}.yml") + sys_json = os.path.join(out_dir, f"system_native_collectives_{dim_signature}{unique_suffix}.json") topo_str = ", ".join(topo_list) npus_str = ", ".join(str(v) for v in npus_list) @@ -336,31 +331,10 @@ def _safe_int(value: Any, default: int = 1) -> int: net_content += f"faulty_links: {faulty_links_str}\n" os.makedirs(os.path.dirname(net_yaml), exist_ok=True) - cache_key = ( - tuple(topo_list), - tuple(npus_list), - tuple(bw_list), - tuple(lat_list), - faulty_links_tuple, - ) - cached_path = _NET_YAML_CACHE.get(cache_key) if reuse_cached_network else None - need_write = True - if cached_path and os.path.exists(cached_path): - try: - with open(cached_path, "r", encoding="utf-8") as handle: - existing = handle.read() - if existing == net_content: - net_yaml = cached_path - need_write = False - except Exception: - need_write = True - if need_write: - tmp_path = net_yaml + ".tmp" - with open(tmp_path, "w", encoding="utf-8") as handle: - handle.write(net_content) - os.replace(tmp_path, net_yaml) - if reuse_cached_network: - _NET_YAML_CACHE[cache_key] = net_yaml + tmp_path = net_yaml + ".tmp" + with open(tmp_path, "w", encoding="utf-8") as handle: + handle.write(net_content) + os.replace(tmp_path, net_yaml) if astra_cfg: coll = astra_cfg.collectives @@ -426,18 +400,7 @@ def _collective_for_dimension(dim, topo_name: str, op: str, default_alg: str) -> pds = pds[0] if pds else 1 system["preferred-dataset-splits"] = int(pds) - _save_json( - sys_json, - system, - npus_key=( - int(npus_count), - dim_signature, - tuple(ar_impl), - tuple(ag_impl), - tuple(rs_impl), - tuple(a2a_impl), - ), - ) + _save_json(sys_json, system) return { "network_yaml": net_yaml, diff --git a/astrasim_lib/executor.py b/astrasim_lib/executor.py index e9d2acdb..81b803cb 100644 --- a/astrasim_lib/executor.py +++ b/astrasim_lib/executor.py @@ -1738,7 +1738,6 @@ def run_astra_simulation_only_onepath( rank_count, axes_filter=axes_filter, faulty_links_override=faulty_links_override, - reuse_cached_network=not persist, ) # Run AstraSim simulation on forward graph (cached via manifest) diff --git a/astrasim_lib/integration.py b/astrasim_lib/integration.py index 8108dfc2..22bd1997 100644 --- a/astrasim_lib/integration.py +++ b/astrasim_lib/integration.py @@ -36,6 +36,8 @@ ensure_chakra_available() +_DEFAULT_ASTRA_CACHE_DIR = os.path.abspath("./astra_cache") + _CACHE_MODES = { "NO_CACHE", "CACHE_READONLY", @@ -135,6 +137,19 @@ def _hash_file_bundle(paths: Iterable[str]) -> str: return h.hexdigest() +def _path_within_dir(path: Optional[str], directory: str) -> bool: + """Return True when ``path`` is within ``directory``.""" + + if not path: + return False + try: + path_abs = os.path.abspath(path) + dir_abs = os.path.abspath(directory) + return os.path.commonpath([path_abs, dir_abs]) == dir_abs + except (OSError, ValueError): + return False + + def _load_cache(cache_path: str) -> Dict[str, Any]: """Load a JSON cache manifest if available, otherwise return an empty dict.""" @@ -326,6 +341,8 @@ def run_cache_astrasim( if isinstance(hw_obj, str): raise TypeError("run_cache_astrasim expects a parsed HWConfig object, not a path") + comm_lower = str(comm).lower() + cache_mode = _cache_mode() allow_read = cache_mode in {"CACHE_READONLY", "CACHE_READWRITE"} allow_write = cache_mode == "CACHE_READWRITE" @@ -343,6 +360,8 @@ def run_cache_astrasim( if allow_read or allow_write: ensure_cache_file_exists(cache_path) + ephemeral_paths: List[str] = [] + try: (intra_ib_bps, intra_ll_s), _ = compute_intra_inter_ib_ll_from_hw(hw_obj) if not intra_ib_bps or intra_ib_bps <= 0 or not intra_ll_s or intra_ll_s <= 0: @@ -354,8 +373,15 @@ def run_cache_astrasim( sys_opts_sig = _sys_options_from_hw(hw_obj) axes_filter_tuple: Optional[Tuple[str, ...]] = tuple(axes_filter) if axes_filter else None + cache_dir_abs = os.path.abspath(astra_config_dir) + ephemeral_configs = ( + files is None + and comm_lower != "graph" + and cache_dir_abs == _DEFAULT_ASTRA_CACHE_DIR + ) + sig: Dict[str, Any] = { - "comm": comm.lower(), + "comm": comm_lower, "npus": int(npus_count), "size_bytes": int(size_bytes), "topology": topo, @@ -378,7 +404,15 @@ def run_cache_astrasim( npus_count=npus_count, axes_filter=axes_filter_tuple, transform_2d_to_1d=transform_2d_to_1d, + ephemeral_outputs=ephemeral_configs, ) + if ephemeral_configs: + ephemeral_paths.extend( + [ + files.get("system_json"), + files.get("network_yaml"), + ] + ) topo_list = files.get("topology_list") npus_per_dim = files.get("npus_per_dim") if topo_list: @@ -393,8 +427,8 @@ def run_cache_astrasim( pass elif not manifest_json_path: label = size_label(size_bytes) - base_dir = os.path.join(astra_config_dir, "workload", comm.lower(), f"{npus_count}npus_{label}") - workload_prefix = os.path.join(base_dir, f"{comm.lower()}_{label}") + base_dir = os.path.join(astra_config_dir, "workload", comm_lower, f"{npus_count}npus_{label}") + workload_prefix = os.path.join(base_dir, f"{comm_lower}_{label}") expected = [f"{workload_prefix}.{r}.et" for r in range(npus_count)] if not all(os.path.exists(p) for p in expected): generate_workload_et(comm, npus_count, size_bytes, astra_config_dir=astra_config_dir) @@ -423,7 +457,7 @@ def run_cache_astrasim( cached_per = entry.get("per_node_sec", []) cached_max = float(entry.get("max_sec", 0.0)) if cached_per and cached_max > 0 and len(cached_per) == int(npus_count) and all((t > 0 for t in cached_per)): - if str(comm).lower() == "graph": + if comm_lower == "graph": try: print( f"[AstraSim] Cache HIT: comm={comm}, npus={npus_count}, size={size_bytes}, total={cached_max:.6f}s" @@ -497,6 +531,15 @@ def run_cache_astrasim( return per_node_sec, max_sec finally: + if ephemeral_paths: + for cleanup_path in {p for p in ephemeral_paths if p}: + if _path_within_dir(cleanup_path, astra_config_dir): + try: + os.remove(cleanup_path) + except FileNotFoundError: + pass + except OSError: + pass if temp_cache_dir is not None: shutil.rmtree(temp_cache_dir, ignore_errors=True) From e733920bac5dc274886a93fc261151e2bf9795d8 Mon Sep 17 00:00:00 2001 From: gkarfakis19 Date: Sat, 8 Nov 2025 16:47:18 -0800 Subject: [PATCH 31/77] fixed many configs --- configs/hardware-config/H100_SXM5_80GB.yaml | 23 ++--- configs/hardware-config/a100_80GB.yaml | 91 ++++++++++--------- .../a100_80GB_no_parallelism.yaml | 23 ++--- configs/hardware-config/v100.yaml | 23 ++--- .../waferscale_20v100_80hbm.yaml | 23 ++--- 5 files changed, 95 insertions(+), 88 deletions(-) diff --git a/configs/hardware-config/H100_SXM5_80GB.yaml b/configs/hardware-config/H100_SXM5_80GB.yaml index b6a642c2..5c2b8e53 100644 --- a/configs/hardware-config/H100_SXM5_80GB.yaml +++ b/configs/hardware-config/H100_SXM5_80GB.yaml @@ -91,17 +91,18 @@ parallelism: cp : 1 # context parallelism network: - - id: dim0 - label: device_fabric - size: 1 - topology: - type: Mesh - bandwidth: 400 Gb - latency: 5e-6 - energy_per_bit: 8e-12 # (J/bit) - util: 1.0 - collective_override: {} # overrides execution_backend.collectives when provided - parallelisms: [tp, cp, lp, dp, kp1, kp2] + dimensions: + - id: dim0 + label: device_fabric + size: 1 + topology: + type: Mesh + bandwidth: 400 Gb + latency: 5e-6 + energy_per_bit: 8e-12 # (J/bit) + util: 1.0 + collective_override: {} # overrides execution_backend.collectives when provided + parallelisms: [tp, cp, lp, dp] execution_backend: model: astra # analytical | astra diff --git a/configs/hardware-config/a100_80GB.yaml b/configs/hardware-config/a100_80GB.yaml index d48cf3f2..7e27d5a5 100644 --- a/configs/hardware-config/a100_80GB.yaml +++ b/configs/hardware-config/a100_80GB.yaml @@ -2,7 +2,7 @@ sw_param: kernel_launch_overhead: 9e-6 # (s) # dp_zero_stage selects ZeRO-style data-parallel sharding: 0=DDP, 1=ZeRO-1, 2=ZeRO-2, 3=ZeRO-3 # only used if data parallelism (dp) > 1 - dp_zero_stage: 2 + dp_zero_stage: 1 precision: # tensor_format accepts numeric byte counts or dtype strings (fp16, bf16, fp32). # Take care that the number of flops defined below has been carefully tuned for FP16/BF16 formats. @@ -13,7 +13,7 @@ sw_param: # -> false forces all activations/weights/gradients/optimizer states/stats to use tensor_format bytes. # -> true enables AMP-style FP32 accumulations/optimizer states depending on param_storage_mode. # if true, optimizer states and stats are always FP32. - mixed_precision: false + mixed_precision: true # param_storage_mode: as_tensor_format | tensor_plus_fp32_master | fp32_params. # -> as_tensor_format makes weights/gradients use tensor_format bytes. # -> tensor_plus_fp32_master makes weights/gradients use tensor_format bytes and @@ -33,10 +33,10 @@ tech_param: FMA_d1: 8 # tensor core can compute the results for an d1 x d2 x d1 mixed-precision matrix multiplication per clock FMA_d2: 4 dataflow: best # {wst | ast | ost | best}: wst: weight stationary, ast: activation stationary, ost: output stationary, best: best of all 3 - util: 100.0 + util: 1.0 DRAM: size: 80 GB - bandwidth: 198600 GB # (/s) + bandwidth: 1986 GB # (/s) dynamic_energy_per_bit: 4.5e-12 # (J) latency: 100e-9 # (s) util: 1 @@ -75,52 +75,55 @@ memory_hierarchy: parallelism: dp: 2 # data parallelism. Acts as replica count during inference - tp: 4 # tensor parallelism + tp: 2 # tensor parallelism tp_sp: true # tensor-sequence parallelism - cp: 1 # context parallelism. Not supported for inference. - lp: 2 # pipeline parallelism. Number of pipeline stages + cp: 2 # context parallelism. Not supported for inference. + lp: 4 # pipeline parallelism. Number of pipeline stages # mb is not a true parallelism parameter, does not affect num of GPUs. - mb: 2 # micro-batches (for best results set to >= lp). + mb: 4 # micro-batches (for best results set to >= lp). network: - # currently only support astrasim integration with all devices in one node. - - id: dim0 - label: intra_wafer_mesh_silicon - size: 4 - topology: - type: Ring - bandwidth: 1000 Gb - latency: 5e-6 - energy_per_bit: 8e-12 # (J/bit) - util: 1.0 - collective_override: {} # overrides execution_backend.collectives when provided - parallelisms: [tp, cp] - - id: dim1 - label: inter_wafer_ring_infiniband - size: 2 - topology: - type: Ring - bandwidth: 200 Gb - latency: 5e-6 - energy_per_bit: 8e-12 # (J/bit) - util: 1.0 - collective_override: {} # overrides execution_backend.collectives when provided - parallelisms: [lp] - - id: dim2 - label: wafer_dp_chains_ring_infiniband - size: 2 - topology: - type: ring - bandwidth: 200 Gb - latency: 5e-6 - energy_per_bit: 8e-12 # (J/bit) - util: 0.96 - collective_override: {} # overrides execution_backend.collectives when provided - parallelisms: [dp] - + # Global faulty links apply to the active AstraSim dimension: [[src, dst, weight]] + # faulty_links: [[0,1,0.6],[16,17,0.1]] + dimensions: + # currently only support astrasim integration with all devices in one node. + - id: dim0 + label: torus2d_tp + size: auto + topology: + type: Ring + bandwidth: 100 GB # We assume 8 25GB/s NVLINKs in the torus. Torus2D is hardcoded bidirectional, so it's actually 4 links, 25GB/s each (because they get replicated). So 100GB/s total. + latency: 5e-6 + energy_per_bit: 8e-12 # (J/bit) + util: 1.0 + optimize_2dmap: false + collective_override: {} # overrides execution_backend.collectives when provided + parallelisms: [tp, cp] + - id: dim1 + label: inter_torus_infiniband + size: auto + topology: + type: Ring + bandwidth: 25 GB + latency: 5e-6 + energy_per_bit: 8e-12 # (J/bit) + util: 1.0 + collective_override: {} # overrides execution_backend.collectives when provided + parallelisms: [lp,dp] + - id: dim2 + label: inter_torus_infiniband_dp + size: auto + topology: + type: Ring + bandwidth: 25 GB + latency: 5e-6 + energy_per_bit: 8e-12 # (J/bit) + util: 1.0 + parallelisms: [] + collective_override: {} # overrides execution_backend.collectives when provided execution_backend: - model: analytical # analytical | astra + model: astra # analytical | astra astra: backend: astra mode: full_astrasim_flattened # hybrid | full_astrasim_hierarchical | full_astrasim_flattened diff --git a/configs/hardware-config/a100_80GB_no_parallelism.yaml b/configs/hardware-config/a100_80GB_no_parallelism.yaml index ef10ba61..798101a2 100644 --- a/configs/hardware-config/a100_80GB_no_parallelism.yaml +++ b/configs/hardware-config/a100_80GB_no_parallelism.yaml @@ -95,17 +95,18 @@ parallelism: cp : 1 # context parallelism network: - - id: dim0 - label: device_fabric - size: 1 - topology: - type: Ring - bandwidth: 400 Gb - latency: 5e-6 - energy_per_bit: 8e-12 # (J/bit) - util: 1.0 - collective_override: {} # overrides execution_backend.collectives when provided - parallelisms: [tp, cp, lp, dp, kp1, kp2] + dimensions: + - id: dim0 + label: device_fabric + size: 1 + topology: + type: Ring + bandwidth: 400 Gb + latency: 5e-6 + energy_per_bit: 8e-12 # (J/bit) + util: 1.0 + collective_override: {} # overrides execution_backend.collectives when provided + parallelisms: [tp, cp, lp, dp] diff --git a/configs/hardware-config/v100.yaml b/configs/hardware-config/v100.yaml index 5f7837cf..a90f36a8 100644 --- a/configs/hardware-config/v100.yaml +++ b/configs/hardware-config/v100.yaml @@ -95,17 +95,18 @@ parallelism: cp : 1 # context parallelism network: - - id: dim0 - label: device_fabric - size: 1 - topology: - type: Mesh - bandwidth: 300 Gb - latency: 5e-6 - energy_per_bit: 8e-12 # (J/bit) - util: 1.0 - collective_override: {} # overrides execution_backend.collectives when provided - parallelisms: [tp, cp, lp, dp, kp1, kp2] + dimensions: + - id: dim0 + label: device_fabric + size: 1 + topology: + type: Mesh + bandwidth: 300 Gb + latency: 5e-6 + energy_per_bit: 8e-12 # (J/bit) + util: 1.0 + collective_override: {} # overrides execution_backend.collectives when provided + parallelisms: [tp, cp, lp, dp] execution_backend: diff --git a/configs/hardware-config/waferscale_20v100_80hbm.yaml b/configs/hardware-config/waferscale_20v100_80hbm.yaml index b7d00522..cfa0ecd0 100644 --- a/configs/hardware-config/waferscale_20v100_80hbm.yaml +++ b/configs/hardware-config/waferscale_20v100_80hbm.yaml @@ -112,17 +112,18 @@ parallelism: cp: 1 network: - - id: dim0 - label: wafer_mesh - size: 20 - topology: - type: Mesh - bandwidth: 5000 Gb - latency: 5e-6 - energy_per_bit: 8e-12 # (J/bit) - util: 0.95 - collective_override: {} # overrides execution_backend.collectives when provided - parallelisms: [tp, cp, lp, dp] + dimensions: + - id: dim0 + label: wafer_mesh + size: 20 + topology: + type: Mesh + bandwidth: 5000 Gb + latency: 5e-6 + energy_per_bit: 8e-12 # (J/bit) + util: 0.95 + collective_override: {} # overrides execution_backend.collectives when provided + parallelisms: [tp, cp, lp, dp] execution_backend: model: astra # analytical | astra From b9cd28d3edc38a145c8d7ed8e260a6de8e0f8a0e Mon Sep 17 00:00:00 2001 From: gkarfakis19 Date: Mon, 10 Nov 2025 17:45:55 -0800 Subject: [PATCH 32/77] added pytest, fixed hybrid/hierarchical bug --- LLM_execution.py | 11 - pyproject.toml | 6 + tests/test_koyeb_validation.py | 18 ++ uv.lock | 45 +++ validation_scripts/__init__.py | 1 + validation_scripts/koyeb.py | 273 ++++++++++-------- validation_scripts/validation_helpers.py | 336 +++++++++++++++++++++++ 7 files changed, 560 insertions(+), 130 deletions(-) create mode 100644 tests/test_koyeb_validation.py create mode 100644 validation_scripts/__init__.py create mode 100644 validation_scripts/validation_helpers.py diff --git a/LLM_execution.py b/LLM_execution.py index 963c00d7..17d26b19 100644 --- a/LLM_execution.py +++ b/LLM_execution.py @@ -765,17 +765,6 @@ def _safe_int(value: Any, default: int = 1) -> int: "stage_span": span, } - return descriptor - - def _attach_optimize_hint(self, root: Any) -> None: - if root is None: - return - if self._first_dim_optimize_cfg: - setattr(root, "_optimize_2dmap", dict(self._first_dim_optimize_cfg)) - elif hasattr(root, "_optimize_2dmap"): - delattr(root, "_optimize_2dmap") - - def _subset_layout(allowed: Sequence[str]) -> Optional[Dict[str, Any]]: subset = [axis for axis in axis_order if axis in allowed and axis_sizes.get(axis, 1) > 1] if not subset: diff --git a/pyproject.toml b/pyproject.toml index ac04a1ec..ecfe9c9c 100644 --- a/pyproject.toml +++ b/pyproject.toml @@ -12,6 +12,7 @@ dependencies = [ "pandas==2.3.1", "protobuf==5.28.2", "pydot>=3.0.0", + "pytest>=9.0.0", "python-dateutil==2.9.0.post0", "pytz==2025.2", "pyyaml==6.0.2", @@ -21,3 +22,8 @@ dependencies = [ "tqdm>=4.67.1", "tzdata==2025.2", ] + +[tool.pytest.ini_options] +testpaths = [ + "tests", +] diff --git a/tests/test_koyeb_validation.py b/tests/test_koyeb_validation.py new file mode 100644 index 00000000..e9c526a0 --- /dev/null +++ b/tests/test_koyeb_validation.py @@ -0,0 +1,18 @@ +import math +import sys +from pathlib import Path + +PROJECT_ROOT = Path(__file__).resolve().parents[1] +if str(PROJECT_ROOT) not in sys.path: + sys.path.insert(0, str(PROJECT_ROOT)) + +from validation_scripts import koyeb + + +def test_koyeb_average_absolute_error_below_threshold(): + results = koyeb.run_all_and_plot(enable_plot=False, verbose=False) + avg_abs_error = results["avg_abs_error"] + assert not math.isnan(avg_abs_error), "Koyeb validation produced no valid measurements" + # Historically the average absolute error bottoms out at roughly 5.97%, so we allow + # a little slack (6.5%) to absorb minor DeepFlow changes while still flagging regressions. + assert avg_abs_error <= 6.5, f"Average absolute error {avg_abs_error:.2f}% exceeds 6.5%" diff --git a/uv.lock b/uv.lock index 8eb1a13b..2142aa36 100644 --- a/uv.lock +++ b/uv.lock @@ -130,6 +130,7 @@ dependencies = [ { name = "pandas" }, { name = "protobuf" }, { name = "pydot" }, + { name = "pytest" }, { name = "python-dateutil" }, { name = "pytz" }, { name = "pyyaml" }, @@ -149,6 +150,7 @@ requires-dist = [ { name = "pandas", specifier = "==2.3.1" }, { name = "protobuf", specifier = "==5.28.2" }, { name = "pydot", specifier = ">=3.0.0" }, + { name = "pytest", specifier = ">=9.0.0" }, { name = "python-dateutil", specifier = "==2.9.0.post0" }, { name = "pytz", specifier = "==2025.2" }, { name = "pyyaml", specifier = "==6.0.2" }, @@ -217,6 +219,15 @@ wheels = [ { url = "https://files.pythonhosted.org/packages/91/4c/e0ce1ef95d4000ebc1c11801f9b944fa5910ecc15b5e351865763d8657f8/graphviz-0.21-py3-none-any.whl", hash = "sha256:54f33de9f4f911d7e84e4191749cac8cc5653f815b06738c54db9a15ab8b1e42", size = 47300 }, ] +[[package]] +name = "iniconfig" +version = "2.3.0" +source = { registry = "https://pypi.org/simple" } +sdist = { url = "https://files.pythonhosted.org/packages/72/34/14ca021ce8e5dfedc35312d08ba8bf51fdd999c576889fc2c24cb97f4f10/iniconfig-2.3.0.tar.gz", hash = "sha256:c76315c77db068650d49c5b56314774a7804df16fee4402c1f19d6d15d8c4730", size = 20503 } +wheels = [ + { url = "https://files.pythonhosted.org/packages/cb/b1/3846dd7f199d53cb17f49cba7e651e9ce294d8497c8c150530ed11865bb8/iniconfig-2.3.0-py3-none-any.whl", hash = "sha256:f631c04d2c48c52b84d0d0549c99ff3859c98df65b3101406327ecc7d53fbf12", size = 7484 }, +] + [[package]] name = "joblib" version = "1.5.2" @@ -598,6 +609,15 @@ wheels = [ { url = "https://files.pythonhosted.org/packages/95/7e/f896623c3c635a90537ac093c6a618ebe1a90d87206e42309cb5d98a1b9e/pillow-12.0.0-pp311-pypy311_pp73-win_amd64.whl", hash = "sha256:b290fd8aa38422444d4b50d579de197557f182ef1068b75f5aa8558638b8d0a5", size = 6997850 }, ] +[[package]] +name = "pluggy" +version = "1.6.0" +source = { registry = "https://pypi.org/simple" } +sdist = { url = "https://files.pythonhosted.org/packages/f9/e2/3e91f31a7d2b083fe6ef3fa267035b518369d9511ffab804f839851d2779/pluggy-1.6.0.tar.gz", hash = "sha256:7dcc130b76258d33b90f61b658791dede3486c3e6bfb003ee5c9bfb396dd22f3", size = 69412 } +wheels = [ + { url = "https://files.pythonhosted.org/packages/54/20/4d324d65cc6d9205fabedc306948156824eb9f0ee1633355a8f7ec5c66bf/pluggy-1.6.0-py3-none-any.whl", hash = "sha256:e920276dd6813095e9377c0bc5566d94c932c33b27a3e3945d8389c374dd4746", size = 20538 }, +] + [[package]] name = "protobuf" version = "5.28.2" @@ -624,6 +644,15 @@ wheels = [ { url = "https://files.pythonhosted.org/packages/7e/32/a7125fb28c4261a627f999d5fb4afff25b523800faed2c30979949d6facd/pydot-4.0.1-py3-none-any.whl", hash = "sha256:869c0efadd2708c0be1f916eb669f3d664ca684bc57ffb7ecc08e70d5e93fee6", size = 37087 }, ] +[[package]] +name = "pygments" +version = "2.19.2" +source = { registry = "https://pypi.org/simple" } +sdist = { url = "https://files.pythonhosted.org/packages/b0/77/a5b8c569bf593b0140bde72ea885a803b82086995367bf2037de0159d924/pygments-2.19.2.tar.gz", hash = "sha256:636cb2477cec7f8952536970bc533bc43743542f70392ae026374600add5b887", size = 4968631 } +wheels = [ + { url = "https://files.pythonhosted.org/packages/c7/21/705964c7812476f378728bdf590ca4b771ec72385c533964653c68e86bdc/pygments-2.19.2-py3-none-any.whl", hash = "sha256:86540386c03d588bb81d44bc3928634ff26449851e99741617ecb9037ee5ec0b", size = 1225217 }, +] + [[package]] name = "pyparsing" version = "3.2.5" @@ -633,6 +662,22 @@ wheels = [ { url = "https://files.pythonhosted.org/packages/10/5e/1aa9a93198c6b64513c9d7752de7422c06402de6600a8767da1524f9570b/pyparsing-3.2.5-py3-none-any.whl", hash = "sha256:e38a4f02064cf41fe6593d328d0512495ad1f3d8a91c4f73fc401b3079a59a5e", size = 113890 }, ] +[[package]] +name = "pytest" +version = "9.0.0" +source = { registry = "https://pypi.org/simple" } +dependencies = [ + { name = "colorama", marker = "sys_platform == 'win32'" }, + { name = "iniconfig" }, + { name = "packaging" }, + { name = "pluggy" }, + { name = "pygments" }, +] +sdist = { url = "https://files.pythonhosted.org/packages/da/1d/eb34f286b164c5e431a810a38697409cca1112cee04b287bb56ac486730b/pytest-9.0.0.tar.gz", hash = "sha256:8f44522eafe4137b0f35c9ce3072931a788a21ee40a2ed279e817d3cc16ed21e", size = 1562764 } +wheels = [ + { url = "https://files.pythonhosted.org/packages/72/99/cafef234114a3b6d9f3aaed0723b437c40c57bdb7b3e4c3a575bc4890052/pytest-9.0.0-py3-none-any.whl", hash = "sha256:e5ccdf10b0bac554970ee88fc1a4ad0ee5d221f8ef22321f9b7e4584e19d7f96", size = 373364 }, +] + [[package]] name = "python-dateutil" version = "2.9.0.post0" diff --git a/validation_scripts/__init__.py b/validation_scripts/__init__.py new file mode 100644 index 00000000..36426389 --- /dev/null +++ b/validation_scripts/__init__.py @@ -0,0 +1 @@ +"""Validation helpers for DeepFlow.""" diff --git a/validation_scripts/koyeb.py b/validation_scripts/koyeb.py index 107a8446..86b496a1 100644 --- a/validation_scripts/koyeb.py +++ b/validation_scripts/koyeb.py @@ -1,16 +1,17 @@ #!/usr/bin/env python3 import os -import sys -import re -import tempfile -import subprocess -import yaml import math -from typing import Dict, Tuple, List +from typing import Dict, Tuple, List, Any import seaborn as sns import matplotlib.pyplot as plt +from validation_scripts.validation_helpers import ( + ValidationSpec, + parse_decode_time, + run_validation_suite, +) + # Llama-3.1-8B Instruct - A100 SXM (from Koyeb matrix, using actual Input/Output tokens) # https://www.koyeb.com/docs/hardware/gpu-benchmarks @@ -73,62 +74,6 @@ HARDWARE_CONFIG = os.path.join(PROJECT_ROOT, "configs/hardware-config/a100_80GB_no_parallelism.yaml") BASE_MODEL_CONFIG = os.path.join(PROJECT_ROOT, "configs/model-config/Llama3.1-8B_inf.yaml") -def _load_yaml(path: str) -> dict: - with open(path, "r") as f: - return yaml.safe_load(f) - -def _write_yaml(path: str, data: dict) -> None: - with open(path, "w") as f: - yaml.safe_dump(data, f, sort_keys=False) - - -def _update_model_config(base_cfg: dict, *, prefill_len: int, decode_len: int, batch_size: int) -> dict: - cfg = dict(base_cfg) # shallow copy is fine; we reassign primitives - model_param = dict(cfg.get("model_param", {})) - # seq_len = prefill + decode - model_param["seq_len"] = int(prefill_len + decode_len) - model_param["decode_len"] = int(decode_len) - model_param["batch_size"] = int(batch_size) - cfg["model_param"] = model_param - return cfg - - -def _run_deepflow_with_temp_model(prefill_len: int, decode_len: int, batch_size: int) -> Tuple[str, float]: - base_cfg = _load_yaml(BASE_MODEL_CONFIG) - updated_cfg = _update_model_config(base_cfg, prefill_len=prefill_len, decode_len=decode_len, batch_size=batch_size) - # if .tmp does not exist, make it - if not os.path.exists("./tmp"): - os.makedirs("./tmp") - fd, tmp_path = tempfile.mkstemp(prefix="koyeb_model_", suffix=".yaml", dir="./tmp") - os.close(fd) - - try: - _write_yaml(tmp_path, updated_cfg) - cmd = [ - sys.executable, - RUN_PERF, - "--hardware_config", HARDWARE_CONFIG, - "--model_config", tmp_path, - ] - - proc = subprocess.run(cmd, cwd=PROJECT_ROOT, stdout=subprocess.PIPE, stderr=subprocess.STDOUT, text=True, check=False) - output = proc.stdout - - # Parse decode time from line: [prefill] time: Xs, [decode] time: Ys, [total] time: Zs - decode_time = None - m = re.search(r"\[prefill\]\s*time:\s*[0-9.]+s,\s*\[decode\]\s*time:\s*([0-9.]+)s,\s*\[total\]\s*time:\s*[0-9.]+s", output) - if m: - decode_time = float(m.group(1)) - - if decode_time is None: - raise RuntimeError(f"Failed to parse decode time from DeepFlow output: {output}") - return output, decode_time - finally: - try: - os.remove(tmp_path) - except Exception: - pass - def _collect_points_in_order(data: Dict[str, Dict[Tuple[int, int, int], Tuple[float, float]]]) -> List[Tuple[str, Tuple[int, int, int], Tuple[float, float]]]: ordered = [] @@ -142,83 +87,173 @@ def _collect_points_in_order(data: Dict[str, Dict[Tuple[int, int, int], Tuple[fl return ordered -def run_all_and_plot() -> None: +def run_all_and_plot(*, enable_plot: bool = True, verbose: bool = True) -> Dict[str, Any]: experiments = _collect_points_in_order(KOYEB_DATA) if not experiments: - print("No experiments to run.") - return - + if verbose: + print("No experiments to run.") + return { + "labels": [], + "expected_tokps": [], + "predicted_tokps": [], + "percent_errors": [], + "avg_error": float("nan"), + "avg_abs_error": float("nan"), + "plot_path": None, + } + + log = print if verbose else (lambda *args, **kwargs: None) labels: List[str] = [] expected_tokps: List[float] = [] predicted_tokps: List[float] = [] percent_errors: List[float] = [] + specs: List[ValidationSpec] = [] - for size, (prefill_len, decode_len, batch_size), (expected_decode_time_s, expected_tokps_agg) in experiments: + for idx, (size, (prefill_len, decode_len, batch_size), (expected_decode_time_s, expected_tokps_agg)) in enumerate(experiments): label = f"{size} bs={batch_size}" - print(f"\n=== Running {label} (prefill={prefill_len}, decode={decode_len}, batch={batch_size}) ===") - try: - console, decode_time_s = _run_deepflow_with_temp_model(prefill_len, decode_len, batch_size) - # Aggregate tokens per second using overall definition: total tokens / decode_time - total_tokens_agg = int(decode_len) * int(batch_size) # dp=1 for no-parallelism hardware config - predicted_tokps_agg = float(total_tokens_agg) / float(decode_time_s) if decode_time_s > 0 else 0.0 - - # Calculate percentage error - if expected_tokps_agg > 0: - pct_error = (predicted_tokps_agg - expected_tokps_agg) / expected_tokps_agg * 100.0 - else: - pct_error = float("nan") - - print(f"Predicted agg tok/s = {predicted_tokps_agg:.2f} (decode_time={decode_time_s:.2f}s)") - print(f"Expected agg tok/s = {expected_tokps_agg:.2f}") - print(f"Error = {pct_error:+.2f}%") - except Exception as e: - print(f"ERROR running DeepFlow for {label}: {e}") + metadata = { + "size": size, + "prefill_len": int(prefill_len), + "decode_len": int(decode_len), + "batch_size": int(batch_size), + "expected_tokps": float(expected_tokps_agg), + "expected_decode_time_s": float(expected_decode_time_s), + } + model_overrides = { + "model_param": { + "seq_len": int(prefill_len + decode_len), + "decode_len": int(decode_len), + "global_batch_size": int(batch_size), + }, + } + specs.append( + ValidationSpec( + label=label, + model_overrides=model_overrides, + hardware_overrides=None, + metadata=metadata, + order=idx, + ) + ) + + validation_results = run_validation_suite( + specs, + base_model_config_path=BASE_MODEL_CONFIG, + base_hardware_config_path=HARDWARE_CONFIG, + result_parser=parse_decode_time, + run_perf_path=RUN_PERF, + ) + + for result in validation_results: + meta = result.spec.metadata + label = result.spec.label + prefill_len = int(meta["prefill_len"]) + decode_len = int(meta["decode_len"]) + batch_size = int(meta["batch_size"]) + expected_tokps_agg = float(meta["expected_tokps"]) + labels.append(label) + expected_tokps.append(expected_tokps_agg) + + if result.success: + decode_time_s = float(result.metrics.get("decode_time_s", float("nan"))) + else: + decode_time_s = float("nan") + + total_tokens_agg = int(decode_len) * int(batch_size) + if decode_time_s and decode_time_s > 0: + predicted_tokps_agg = float(total_tokens_agg) / float(decode_time_s) + else: predicted_tokps_agg = float("nan") + + if expected_tokps_agg > 0 and not math.isnan(predicted_tokps_agg): + pct_error = (predicted_tokps_agg - expected_tokps_agg) / expected_tokps_agg * 100.0 + else: pct_error = float("nan") - labels.append(label) - expected_tokps.append(expected_tokps_agg) predicted_tokps.append(predicted_tokps_agg) percent_errors.append(pct_error) + block_lines = [ + f"\n=== Result {label} (prefill={prefill_len}, decode={decode_len}, batch={batch_size}) ===", + ] + if result.success and not math.isnan(predicted_tokps_agg): + block_lines.append(f"Predicted agg tok/s = {predicted_tokps_agg:.2f} (decode_time={decode_time_s:.2f}s)") + block_lines.append(f"Expected agg tok/s = {expected_tokps_agg:.2f}") + block_lines.append(f"Error = {pct_error:+.2f}%") + else: + err_detail = result.error or "DeepFlow run failed" + block_lines.append(f"ERROR: {err_detail}") + if result.raw_output: + block_lines.append(result.raw_output.strip()) + log("\n".join(block_lines)) + # Print summary statistics - print("\n" + "="*70) - print("SUMMARY") - print("="*70) + log("\n" + "="*70) + log("SUMMARY") + log("="*70) valid_errors = [e for e in percent_errors if not math.isnan(e)] if valid_errors: avg_error = sum(valid_errors) / len(valid_errors) avg_abs_error = sum(abs(e) for e in valid_errors) / len(valid_errors) - print(f"Average error: {avg_error:+.2f}%") - print(f"Average absolute error: {avg_abs_error:.2f}%") + log(f"Average error: {avg_error:+.2f}%") + log(f"Average absolute error: {avg_abs_error:.2f}%") else: - print("No valid error measurements.") - print("="*70) - - # Plot - fig = plt.figure(figsize=(12, 6)) - x = list(range(len(labels))) - plt.plot(x, predicted_tokps, marker="o", linestyle="-", color="#1f77b4", label="DeepFlow predicted tok/s (aggregate)") - plt.plot(x, expected_tokps, marker="o", linestyle="--", color="#ff7f0e", alpha=0.7, label="Expected tok/s (Koyeb)") - plt.xticks(x, labels, rotation=30, ha="right") - plt.ylabel("Tokens per second (aggregate)") - plt.title("Koyeb vs DeepFlow (single A100 SXM, Llama3.1-8B)") - plt.grid(True, linestyle=":", alpha=0.4) - plt.legend() - - # Add note about large bs=32 datapoint - note_text = ( "X-axis is prefill length + batch size following Koyebs format (small = 512x512, medium = 1024x1024, large = 4096x1024)\n" - "For large bs=32 datapoint, we use PCIE data point, as SXM data is very noisy \n" - "on the official website (it's 2x slower than PCIE). You can change this in the code.") - fig.text(0.5, 0.02, note_text, ha='center', fontsize=9, style='italic', color='#555555', wrap=True) - - out_dir = os.path.join(PROJECT_ROOT, "output", "validation") - os.makedirs(out_dir, exist_ok=True) - out_path = os.path.join(out_dir, "koyeb_a100_sxm_no_parallelism.png") - plt.tight_layout(rect=[0, 0.06, 1, 1]) # Leave space at bottom for note - plt.savefig(out_path, dpi=160) - print(f"Saved plot to: {out_path}") + avg_error = float("nan") + avg_abs_error = float("nan") + log("No valid error measurements.") + log("="*70) + + plot_path = None + if enable_plot: + # Plot + fig = plt.figure(figsize=(12, 6)) + x = list(range(len(labels))) + plt.plot(x, predicted_tokps, marker="o", linestyle="-", color="#1f77b4", label="DeepFlow predicted tok/s (aggregate)") + plt.plot(x, expected_tokps, marker="o", linestyle="--", color="#ff7f0e", alpha=0.7, label="Expected tok/s (Koyeb)") + plt.xticks(x, labels, rotation=30, ha="right") + plt.ylabel("Tokens per second (aggregate)") + plt.title("Koyeb vs DeepFlow (single A100 SXM, Llama3.1-8B)") + plt.grid(True, linestyle=":", alpha=0.4) + plt.legend() + + # Add note about large bs=32 datapoint + note_text = ( "X-axis is prefill length + batch size following Koyebs format (small = 512x512, medium = 1024x1024, large = 4096x1024)\n" + "For large bs=32 datapoint, we use PCIE data point, as SXM data is very noisy \n" + "on the official website (it's 2x slower than PCIE). You can change this in the code.") + fig.text(0.5, 0.02, note_text, ha='center', fontsize=9, style='italic', color='#555555', wrap=True) + + out_dir = os.path.join(PROJECT_ROOT, "output", "validation") + os.makedirs(out_dir, exist_ok=True) + plot_path = os.path.join(out_dir, "koyeb_a100_sxm_no_parallelism.png") + plt.tight_layout(rect=[0, 0.06, 1, 1]) # Leave space at bottom for note + plt.savefig(plot_path, dpi=160) + log(f"Saved plot to: {plot_path}") + + return { + "labels": labels, + "expected_tokps": expected_tokps, + "predicted_tokps": predicted_tokps, + "percent_errors": percent_errors, + "avg_error": avg_error, + "avg_abs_error": avg_abs_error, + "plot_path": plot_path, + } + + +def _parse_args(): + import argparse + parser = argparse.ArgumentParser(description="Validate DeepFlow against Koyeb benchmarks.") + parser.add_argument("--no-plot", action="store_true", help="Skip generating the matplotlib plot/output file.") + parser.add_argument("--quiet", action="store_true", help="Suppress detailed per-experiment logs.") + return parser.parse_args() if __name__ == "__main__": - run_all_and_plot() + args = _parse_args() + results = run_all_and_plot(enable_plot=not args.no_plot, verbose=not args.quiet) + if args.quiet: + avg_abs = results["avg_abs_error"] + if math.isnan(avg_abs): + print("Average absolute error: NaN (no valid measurements)") + else: + print(f"Average absolute error: {avg_abs:.2f}%") diff --git a/validation_scripts/validation_helpers.py b/validation_scripts/validation_helpers.py new file mode 100644 index 00000000..cf538e4b --- /dev/null +++ b/validation_scripts/validation_helpers.py @@ -0,0 +1,336 @@ +from __future__ import annotations + +import copy +import os +import re +import shutil +import subprocess +import sys +import tempfile +import time +from concurrent.futures import ProcessPoolExecutor, as_completed +from dataclasses import dataclass, field, replace +from typing import Any, Callable, Dict, List, Mapping, Optional, Sequence, Tuple + +import yaml + +PROJECT_ROOT = os.path.dirname(os.path.dirname(os.path.abspath(__file__))) +if PROJECT_ROOT not in sys.path: + sys.path.insert(0, PROJECT_ROOT) + +from tools.parallelism_sweep import set_astrasim_cache_mode # type: ignore + +VALIDATION_WORKERS_ENV = "DEEPFLOW_VALIDATION_WORKERS" +DEFAULT_WORKER_COUNT = 8 + +ResultParser = Callable[[str, "ValidationSpec"], Dict[str, Any]] + + +@dataclass(frozen=True) +class ValidationSpec: + label: str + model_overrides: Optional[Mapping[str, Any]] = None + hardware_overrides: Optional[Mapping[str, Any]] = None + metadata: Dict[str, Any] = field(default_factory=dict) + model_config_path: Optional[str] = None + hardware_config_path: Optional[str] = None + order: int = 0 + + +@dataclass +class ValidationResult: + spec: ValidationSpec + success: bool + metrics: Dict[str, Any] + raw_output: str + error: Optional[str] = None + returncode: Optional[int] = None + duration_s: float = 0.0 + model_config_used: Optional[str] = None + hardware_config_used: Optional[str] = None + + +@dataclass +class _WorkerContext: + run_perf_path: str + python_executable: str + tmp_root: str + env_overrides: Dict[str, str] + model_config_cache: Dict[str, Dict[str, Any]] + hardware_config_cache: Dict[str, Dict[str, Any]] + default_model_config_path: str + default_hardware_config_path: str + result_parser: ResultParser + cache_mode: str + extra_run_perf_args: Tuple[str, ...] + + +_WORKER_CONTEXT: Optional[_WorkerContext] = None + + +def _worker_init(context: _WorkerContext) -> None: + global _WORKER_CONTEXT + _WORKER_CONTEXT = context + set_astrasim_cache_mode(context.cache_mode) + + +def _deep_update(target: Dict[str, Any], overrides: Mapping[str, Any]) -> Dict[str, Any]: + for key, value in overrides.items(): + if isinstance(value, Mapping): + existing = target.get(key) + if not isinstance(existing, dict): + existing = {} + target[key] = _deep_update(existing, value) + else: + target[key] = value + return target + + +def _load_yaml(path: str) -> Dict[str, Any]: + with open(path, "r") as handle: + data = yaml.safe_load(handle) + if not isinstance(data, dict): + raise ValueError(f"{path} did not contain a YAML dictionary.") + return data + + +def _write_yaml(path: str, data: Dict[str, Any]) -> None: + with open(path, "w") as handle: + yaml.safe_dump(data, handle, sort_keys=False) + + +def _prepare_config( + tmp_dir: str, + base_path: str, + base_dict: Dict[str, Any], + overrides: Optional[Mapping[str, Any]], + filename: str, +) -> str: + if not overrides: + return base_path + updated = copy.deepcopy(base_dict) + _deep_update(updated, overrides) + config_path = os.path.join(tmp_dir, filename) + _write_yaml(config_path, updated) + return config_path + + +def _read_env_worker_setting() -> Tuple[int, str]: + env_raw = os.environ.get(VALIDATION_WORKERS_ENV) + if env_raw is None: + return DEFAULT_WORKER_COUNT, f"{VALIDATION_WORKERS_ENV} unset (default {DEFAULT_WORKER_COUNT})" + try: + value = int(env_raw) + except (TypeError, ValueError): + print(f"[validation] WARNING: {VALIDATION_WORKERS_ENV} must be a positive integer (got {env_raw!r}); using default {DEFAULT_WORKER_COUNT}.") + return DEFAULT_WORKER_COUNT, f"{VALIDATION_WORKERS_ENV} invalid ({env_raw!r}) -> default {DEFAULT_WORKER_COUNT}" + if value <= 0: + print(f"[validation] WARNING: {VALIDATION_WORKERS_ENV} must be > 0 (got {env_raw!r}); using default {DEFAULT_WORKER_COUNT}.") + return DEFAULT_WORKER_COUNT, f"{VALIDATION_WORKERS_ENV} <= 0 ({env_raw!r}) -> default {DEFAULT_WORKER_COUNT}" + return value, f"{VALIDATION_WORKERS_ENV}={value}" + + +def _determine_worker_count(spec_count: int, explicit_max: Optional[int]) -> int: + env_workers, env_desc = _read_env_worker_setting() + requested = explicit_max if explicit_max is not None else env_workers + request_desc = f"max_workers={explicit_max}" if explicit_max is not None else env_desc + if requested <= 0: + print(f"[validation] WARNING: requested worker budget {requested} is not positive; bumping to 1.") + requested = 1 + cpu_limit = max(1, os.cpu_count() or 1) + if explicit_max is None and requested > cpu_limit: + print(f"[validation] WARNING: {VALIDATION_WORKERS_ENV} requested {requested} worker(s) but only {cpu_limit} CPU core(s) detected. Capping to {cpu_limit}.") + requested = cpu_limit + elif explicit_max is not None and requested > cpu_limit: + print(f"[validation] WARNING: max_workers requested {requested} worker(s) but only {cpu_limit} CPU core(s) detected. Capping to {cpu_limit}.") + requested = cpu_limit + worker_count = max(1, min(spec_count, requested)) + print(f"[validation] Using {worker_count} worker(s) for {spec_count} experiment(s) " + f"(requested {request_desc}, cpu_limit={cpu_limit}).") + return worker_count + + +def _execute_spec(spec: ValidationSpec) -> ValidationResult: + if _WORKER_CONTEXT is None: + raise RuntimeError("Worker context is not initialised.") + ctx = _WORKER_CONTEXT + tmp_dir = tempfile.mkdtemp(prefix="validation_run_", dir=ctx.tmp_root) + env = os.environ.copy() + env.update(ctx.env_overrides) + env.setdefault("ASTRA_CACHE_DIR", os.path.join(tmp_dir, "astra_cache")) + start_time = time.time() + base_model_path = os.path.abspath(spec.model_config_path or ctx.default_model_config_path) + base_hw_path = os.path.abspath(spec.hardware_config_path or ctx.default_hardware_config_path) + if base_model_path not in ctx.model_config_cache: + raise ValueError(f"Model config {base_model_path} was not preloaded for validation.") + if base_hw_path not in ctx.hardware_config_cache: + raise ValueError(f"Hardware config {base_hw_path} was not preloaded for validation.") + model_config_path = base_model_path + hardware_config_path = base_hw_path + try: + model_config_path = _prepare_config( + tmp_dir, + base_model_path, + ctx.model_config_cache[base_model_path], + spec.model_overrides, + "model.yaml", + ) + hardware_config_path = _prepare_config( + tmp_dir, + base_hw_path, + ctx.hardware_config_cache[base_hw_path], + spec.hardware_overrides, + "hardware.yaml", + ) + cmd = [ + ctx.python_executable, + ctx.run_perf_path, + "--hardware_config", + hardware_config_path, + "--model_config", + model_config_path, + ] + if ctx.extra_run_perf_args: + cmd.extend(ctx.extra_run_perf_args) + proc = subprocess.run( + cmd, + cwd=tmp_dir, + stdout=subprocess.PIPE, + stderr=subprocess.STDOUT, + text=True, + check=False, + env=env, + ) + duration = time.time() - start_time + output = proc.stdout or "" + if proc.returncode != 0: + return ValidationResult( + spec=spec, + success=False, + metrics={}, + raw_output=output, + error=f"DeepFlow exited with status {proc.returncode}", + returncode=proc.returncode, + duration_s=duration, + model_config_used=model_config_path, + hardware_config_used=hardware_config_path, + ) + try: + metrics = ctx.result_parser(output, spec) + except Exception as exc: + return ValidationResult( + spec=spec, + success=False, + metrics={}, + raw_output=output, + error=f"Failed to parse DeepFlow output: {exc}", + returncode=proc.returncode, + duration_s=duration, + model_config_used=model_config_path, + hardware_config_used=hardware_config_path, + ) + return ValidationResult( + spec=spec, + success=True, + metrics=metrics, + raw_output=output, + error=None, + returncode=proc.returncode, + duration_s=duration, + model_config_used=model_config_path, + hardware_config_used=hardware_config_path, + ) + finally: + shutil.rmtree(tmp_dir, ignore_errors=True) + + +def run_validation_suite( + specs: Sequence[ValidationSpec], + *, + base_model_config_path: str, + base_hardware_config_path: str, + result_parser: ResultParser, + run_perf_path: Optional[str] = None, + python_executable: Optional[str] = None, + tmp_root: Optional[str] = None, + max_workers: Optional[int] = None, + env_overrides: Optional[Dict[str, str]] = None, + cache_mode: str = "NO_CACHE", + extra_run_perf_args: Optional[Sequence[str]] = None, +) -> List[ValidationResult]: + if not specs: + return [] + run_perf_path = os.path.abspath(run_perf_path or os.path.join(PROJECT_ROOT, "run_perf.py")) + python_executable = python_executable or sys.executable + tmp_root = os.path.abspath(tmp_root or os.path.join(PROJECT_ROOT, "tmp", "validation_runs")) + os.makedirs(tmp_root, exist_ok=True) + set_astrasim_cache_mode(cache_mode) + env_map = dict(env_overrides or {}) + env_map.setdefault("DEEPFLOW_ASTRA_CACHE_MODE", os.environ.get("DEEPFLOW_ASTRA_CACHE_MODE", cache_mode)) + model_paths = {os.path.abspath(base_model_config_path)} + hardware_paths = {os.path.abspath(base_hardware_config_path)} + for spec in specs: + if spec.model_config_path: + model_paths.add(os.path.abspath(spec.model_config_path)) + if spec.hardware_config_path: + hardware_paths.add(os.path.abspath(spec.hardware_config_path)) + model_cache = {path: _load_yaml(path) for path in model_paths} + hardware_cache = {path: _load_yaml(path) for path in hardware_paths} + extra_args: Tuple[str, ...] = tuple(extra_run_perf_args or ()) + assigned_specs = [replace(spec, order=idx) for idx, spec in enumerate(specs)] + worker_count = _determine_worker_count(len(assigned_specs), max_workers) + context = _WorkerContext( + run_perf_path=run_perf_path, + python_executable=python_executable, + tmp_root=tmp_root, + env_overrides=env_map, + model_config_cache=model_cache, + hardware_config_cache=hardware_cache, + default_model_config_path=os.path.abspath(base_model_config_path), + default_hardware_config_path=os.path.abspath(base_hardware_config_path), + result_parser=result_parser, + cache_mode=os.environ.get("DEEPFLOW_ASTRA_CACHE_MODE", cache_mode), + extra_run_perf_args=extra_args, + ) + results: List[ValidationResult] = [] + with ProcessPoolExecutor(max_workers=worker_count, initializer=_worker_init, initargs=(context,)) as executor: + future_map = {executor.submit(_execute_spec, spec): spec for spec in assigned_specs} + for future in as_completed(future_map): + spec = future_map[future] + try: + result = future.result() + except Exception as exc: + result = ValidationResult( + spec=spec, + success=False, + metrics={}, + raw_output="", + error=str(exc), + returncode=None, + duration_s=0.0, + model_config_used=None, + hardware_config_used=None, + ) + results.append(result) + results.sort(key=lambda item: item.spec.order) + return results + + +_DECODE_TIME_REGEX = re.compile( + r"\[prefill\]\s*time:\s*[0-9.]+s,\s*\[decode\]\s*time:\s*([0-9.]+)s,\s*\[total\]\s*time:\s*[0-9.]+s" +) + + +def parse_decode_time(output: str, spec: ValidationSpec) -> Dict[str, Any]: + match = _DECODE_TIME_REGEX.search(output) + if not match: + raise ValueError(f"Failed to parse decode time for experiment '{spec.label}'.") + return {"decode_time_s": float(match.group(1))} + + +__all__ = [ + "ValidationSpec", + "ValidationResult", + "run_validation_suite", + "parse_decode_time", +] From b4f75d7b4c8a7d78691e11626a34f684eb03ea42 Mon Sep 17 00:00:00 2001 From: gkarfakis19 Date: Mon, 10 Nov 2025 17:47:57 -0800 Subject: [PATCH 33/77] small fix --- validation_scripts/koyeb.py | 2 +- 1 file changed, 1 insertion(+), 1 deletion(-) diff --git a/validation_scripts/koyeb.py b/validation_scripts/koyeb.py index 86b496a1..90be11c0 100644 --- a/validation_scripts/koyeb.py +++ b/validation_scripts/koyeb.py @@ -6,7 +6,7 @@ import matplotlib.pyplot as plt -from validation_scripts.validation_helpers import ( +from validation_helpers import ( ValidationSpec, parse_decode_time, run_validation_suite, From 65de98a80a6170535fe09048b1066b51d8d1d497 Mon Sep 17 00:00:00 2001 From: NanoMembers Date: Sun, 9 Nov 2025 01:07:00 -0800 Subject: [PATCH 34/77] implemented expert parallelism --- LLM_execution.py | 21 ++- LLM_util.py | 4 +- config.py | 3 +- model.py | 2 +- time_calculation.py | 2 - time_calculation_LLM.py | 301 +++++++++++++++++++++++++++++++++++----- 6 files changed, 291 insertions(+), 42 deletions(-) diff --git a/LLM_execution.py b/LLM_execution.py index 17d26b19..9c9a655d 100644 --- a/LLM_execution.py +++ b/LLM_execution.py @@ -1036,11 +1036,28 @@ def _run_hybrid(self) -> ExecutionResult: generate_graphs = _env_flag("DEEPFLOW_VISUALIZE_GRAPHS") transformer_time = self._run_transformer_astrasim(ExecutionMode.HYBRID) + + if generate_graphs: - self.transformer_graph.save_graph( + transformer_timed_forward_root = self.transformer_graph.convert_comm_sizes_to_times( self.transformer_forward_root, + self.time_calc.network_model, + self.interconnect_params, + ) + transformer_timed_backward_root = self.transformer_graph.convert_comm_sizes_to_times( + self.transformer_backward_root, + self.time_calc.network_model, + self.interconnect_params, + ) + self.transformer_graph.save_graph( + transformer_timed_forward_root, + self.time_calc.output_dir, + "/hybrid_graph_transformer_forward", + ) + self.transformer_graph.save_graph( + transformer_timed_backward_root, self.time_calc.output_dir, - "/hybrid_graph_transformer", + "/hybrid_graph_transformer_backward", ) if transformer_time is not None: diff --git a/LLM_util.py b/LLM_util.py index 6b10f6b2..29851fd4 100644 --- a/LLM_util.py +++ b/LLM_util.py @@ -77,12 +77,13 @@ def multihead_decoder_gemm(self, batch_size, seq_len, d_model, num_heads, kv_hea # assuming equal load balancing here num_experts = max(1, int(getattr(self, "moe_num_experts", 1))) top_k = max(1, int(getattr(self, "moe_top_k", 1))) - gemms["ffn1"] = (batch_size, seq_len * top_k / num_experts, d_model, projected_dim ) #gemm shape per expert + gemms["ffn1"] = (batch_size, seq_len * top_k / num_experts, d_model, projected_dim ) #gemm shape per expert assuming each dp group has all experts gemms["ffn2"] = (batch_size, seq_len * top_k / num_experts, intermediate_size, d_model) else: gemms["ffn1"] = (batch_size, seq_len, d_model, projected_dim) gemms["ffn2"] = (batch_size, seq_len, intermediate_size, d_model) gemms["linear"] = (batch_size, seq_len, d_model, vocab_size) + gemms["gate"] = (batch_size, seq_len, d_model, self.moe_num_experts) return gemms @@ -136,7 +137,6 @@ def get_transformer_mem_layer( dp, tp, lp, mb, batch_size, hidden_dim, seq_len, act_memory_layer_inf = seq_len * batch_size * intermediate_size / tp * precision.activations #inference max activation memory, no need to store for backpropagation ffn_proj_factor = 3 if str(model_type).lower() == "llama" else 2 - transformer_param_layer = 4* hidden_dim * hidden_dim + intermediate_size * ffn_proj_factor * hidden_dim # weights Wq,Wk,Wv,Wo,ffn optimizer_mem = (precision.optimizer_states * 2 + precision.activations) * transformer_param_layer / tp # don't divide by dp for DDP, NeMo ZeRO-1 optimizer style diff --git a/config.py b/config.py index b9b6e77e..ce785786 100644 --- a/config.py +++ b/config.py @@ -1299,8 +1299,7 @@ def parse_config(filename, config_type): ) from exc if top_k <= 0: raise ValueError("model_param.top_k must be a positive integer when provided") - if top_k > num_experts: - raise ValueError("model_param.top_k cannot exceed model_param.num_experts") + if "run_type" not in mp: raise ValueError("model_param.run_type must be specified") diff --git a/model.py b/model.py index c395e755..e913e67b 100644 --- a/model.py +++ b/model.py @@ -42,7 +42,7 @@ def __init__(self, exp_config): self.moe_num_experts = int(getattr(exp_config.model_config, "num_experts", 1)) self.moe_top_k = int(getattr(exp_config.model_config, "top_k", 1)) - if self.moe_top_k > self.moe_num_experts: + if self.moe_top_k > self.moe_num_experts and self.moe_num_experts > 1: raise ValueError("model_param.top_k cannot exceed model_param.num_experts") self.use_moe = self.moe_num_experts > 1 diff --git a/time_calculation.py b/time_calculation.py index a275dd20..46db0544 100644 --- a/time_calculation.py +++ b/time_calculation.py @@ -364,8 +364,6 @@ def __init__(self, hw_config, model_config, mode, *, astra_policy_override: Opti raw_top_k = getattr(self.model, "moe_top_k", 1) self.moe_num_experts = max(1, int(raw_num_experts)) self.moe_top_k = max(1, int(raw_top_k)) - if self.moe_top_k > self.moe_num_experts: - raise ValueError("model_param.top_k cannot exceed model_param.num_experts") self.use_moe = self.moe_num_experts > 1 diff --git a/time_calculation_LLM.py b/time_calculation_LLM.py index 20c7218a..d8f3cb42 100644 --- a/time_calculation_LLM.py +++ b/time_calculation_LLM.py @@ -1,6 +1,7 @@ import math import os import json +import warnings from enum import Enum from typing import Any, Dict, Tuple, Optional, List, Mapping, Sequence, Set from collections.abc import Mapping as MappingABC, Sequence as SequenceABC @@ -110,6 +111,53 @@ class ParallelismMode(Enum): COMM_RULE_DEFAULT_KEY: {'forward': None, 'backward': None}, }, } + +_MOE_GATE_RULE = { + 'forward': { + 'kind': 'all_to_all', + 'participants': 'moe', + 'interconnect': 'tp', + 'suffix': 'all_to_all', + }, + 'backward': { + 'kind': 'reduce_scatter', + 'participants': 'moe', + 'interconnect': 'tp', + 'suffix': 'reduce_scatter', + }, +} +_MOE_MLP_RULE = { + 'forward': { + 'kind': 'all_to_all', + 'participants': 'moe', + 'interconnect': 'tp', + 'suffix': 'all_to_all', + }, + 'backward': { + 'kind': 'all_to_all', + 'participants': 'moe', + 'interconnect': 'tp', + 'suffix': 'all_to_all', + }, +} +_MOE_LAYER_NORM1_RULE = { + 'backward': { + 'kind': 'all_to_all', + 'participants': 'moe', + 'interconnect': 'tp', + 'suffix': 'all_to_all', + }, +} +MOE_COMMUNICATION_RULES: Dict[ + ParallelismMode, Dict[str, Dict[str, Optional[Dict[str, str]]]] +] = { + mode: { + 'gate': _MOE_GATE_RULE, + 'MLP': _MOE_MLP_RULE, + 'layernorm1': _MOE_LAYER_NORM1_RULE, + } + for mode in ParallelismMode +} class GemmType(Enum): ATTENTION_SCORE = "attention_score" ATTENTION_OUTPUT = "attention_output" @@ -118,6 +166,7 @@ class GemmType(Enum): FFN1 = "ffn1" FFN2 = "ffn2" LINEAR_SOFTMAX = "linear_softmax" + LAYER_NORM_1 = "layer_norm1" class TimeCalculationLLM(TimeCalculation): @@ -132,8 +181,7 @@ def __init__(self, hw_config, model_config, mode, output_dir: Optional[str] = No mode, astra_policy_override=astra_policy, ) - if self.num_workers > 1 and model_config.model_config.num_experts > 1: - raise ValueError("MoE is only supported for single-worker runs (total worker count must be 1).") + self.output_dir = os.path.abspath(output_dir) if output_dir else os.getcwd() os.makedirs(self.output_dir, exist_ok=True) @@ -193,6 +241,31 @@ def get_parallelism_mode(self): else: return ParallelismMode.SINGLE + @property + def experts_per_gpu(self) -> int: + """ + Return the number of MoE experts assigned to each GPU assuming each dp group has all experts. + """ + if not self.use_moe or self.moe_num_experts <= 1: + return 0 + moe_ranks = max(1, self.tp * self.cp) + if moe_ranks > self.moe_num_experts: + warnings.warn( + "MoE configuration has fewer experts than tp * cp; DeepFlow only supports at least one expert per GPU (no expert slicing).", + RuntimeWarning, + ) + raise ValueError("Number of MoE experts must be at least equal to the number of parallel ranks.") + if self.moe_num_experts % moe_ranks != 0: + raise ValueError("Number of MoE experts must be divisible by the total number of parallel ranks (tp * cp ).") + return max(1, math.ceil(self.moe_num_experts / moe_ranks)) + + @experts_per_gpu.setter + def experts_per_gpu(self, value: int) -> None: + # Preserve the base class assignment for debugging, but the getter drives usage. + self._experts_per_gpu_base_value = value + + + def _param_stats_per_rank( self, hidden_dim: int, @@ -371,12 +444,11 @@ def flash_attention_kernel_backward(self) -> Tuple[Dict[str, float], Dict[str, f #TODO load 3 tiled gemm inputs and 1 output from HBM to SRAM #TODO disable HBM access - attention_forward_reduction_time = 0 + attention_forward_reduction_time = 1 attention_size_f = 0 attn_score_time = self.getGEMMTime(Br, d, Bc, "attention_score_f")[0] * Tc * Tr #attention score gemm time for one head number_flops = Br * Bc * d * Tc * Tr * batch_size * num_heads - print(f"Flash Attention: attention score gemm flops={number_flops/1e9} GFLOPS") # Softmax time elements = Br * Bc @@ -437,9 +509,6 @@ def single_gpu_gemm_forward(self, gemm: Tuple[int, ...], name: str, gemm_type: O gemm_type = self._normalize_gemm_type(gemm_type) if gemm_type in (GemmType.ATTENTION_SCORE, GemmType.ATTENTION_OUTPUT): # attention gemm gemm_time = self.getGEMMTime(m, k, n, name, disable_overhead=True)[0] * batch + self.O - elif (gemm_type == GemmType.FFN1 or gemm_type == GemmType.FFN2) and self.use_moe: - gemm_time,_,_, mem_accesses = self.getGEMMTime(m, k, n, name) - gemm_time *= self.moe_num_experts else : gemm_time,_,_, mem_accesses = self.getGEMMTime(m, k, n, name) return gemm_time, 0, 0, total_flops, mem_accesses @@ -450,9 +519,6 @@ def single_gpu_gemm_backward(self, gemm: Tuple[int, ...], name: str, gemm_type: if gemm_type in (GemmType.ATTENTION_SCORE, GemmType.ATTENTION_OUTPUT): # attention gemm grad_time_act = self.getGEMMTime(m, k, n, name, disable_overhead=True)[0] * batch + self.O grad_time_wt = self.getGEMMTime(k, m, n, name, disable_overhead=True)[0] * batch + self.O - elif (gemm_type == GemmType.FFN1 or gemm_type == GemmType.FFN2) and self.use_moe: - grad_time_act = self.getGEMMTime(m, n, k, name)[0] * self.moe_num_experts - grad_time_wt = self.getGEMMTime(k, m, n, name)[0] * self.moe_num_experts else : grad_time_act = self.getGEMMTime(m, n, k, name)[0] grad_time_wt = self.getGEMMTime(k, m, n, name)[0] @@ -675,10 +741,6 @@ def _tensor_parallelism_gemm_forward(self, gemm: Tuple[int, ...], name: str, gem if gemm_type in (GemmType.ATTENTION_SCORE, GemmType.ATTENTION_OUTPUT): # attention gemm gemm_time,_,_, mem_accesses = self.getGEMMTime(m, k, n, name, disable_overhead=True) gemm_time = gemm_time * batch / max(1, self.tp) + self.O - elif gemm_type in (GemmType.FFN1, GemmType.FFN2) and self.use_moe: - # MoE FFN layers replicate experts across all TP ranks, so no TP sharding here. - gemm_time,_,_, mem_accesses = self.getGEMMTime(m, k, n, name) - gemm_time *= self.moe_num_experts elif gemm_type in (GemmType.QKV, GemmType.FFN1): # column wise shard_n = math.ceil(n / max(1, self.tp)) gemm_time,_,_, mem_accesses = self.getGEMMTime(m, k, shard_n, name) @@ -837,7 +899,7 @@ def _context_parallelism_gemm_backward(self, gemm: Tuple[int, ...], name: str, g if gemm_type == GemmType.ATTENTION_SCORE: grad_time_act = self.getGEMMTime(shard_m, n, k, name, disable_overhead=True)[0] * batch + self.O grad_time_wt = self.getGEMMTime(k, shard_m, n, name, disable_overhead=True)[0] * batch + self.O - total_bytes = self.precision.grad_communication * k * n * batch * 2 # account for both Q and K + total_bytes = self.precision.grad_communication * k * n * batch * 2 # account for both K and V kind = "reduce_scatter" elif gemm_type == GemmType.ATTENTION_OUTPUT: # attention gemm grad_time_act = self.getGEMMTime(shard_m, n, k, name, disable_overhead=True)[0] * batch + self.O @@ -865,7 +927,53 @@ def _context_parallelism_gemm_backward(self, gemm: Tuple[int, ...], name: str, g axis="cp", ) return gemm_time, reduction_time, total_bytes - + def get_moe_ffn_f(self, gemm, name, gemm_type: Optional[GemmType] = None): + batch, m, k, n = self._expand_gemm_descriptor(gemm) + reduction_time = 0 + total_bytes = 0 + total_flops = 2 * batch * m * k * n * self.experts_per_gpu + gemm_time, _, _ ,mem_accesses= self.getGEMMTime(m, k, n, name) + gemm_time *= self.experts_per_gpu + + if gemm_type == GemmType.FFN2: + per_rank_bytes = math.ceil(self.precision.activations * m * n * self.experts_per_gpu) + total_bytes = int(math.ceil(per_rank_bytes * self.cp * self.tp )) + reduction_time = self.network_model.collective( + kind="all_to_all", + size_bytes=total_bytes, + participants=self.cp * self.tp, + ib=self.IBTP, + ll=self.LLTP, + local_bytes=0, + debug_label="moe_ffn_f_all_to_all", + ) + return gemm_time, reduction_time, total_bytes, total_flops, mem_accesses + + def get_moe_ffn_b(self, gemm, name, gemm_type): + + batch, m, k, n = self._expand_gemm_descriptor(gemm) + reduction_time = 0 + total_bytes = 0 + grad_time_act = self.getGEMMTime(m, n, k, name)[0] * self.experts_per_gpu + grad_time_wt = self.getGEMMTime(k, m, n, name)[0] * self.experts_per_gpu + gemm_time = grad_time_act + grad_time_wt + + if gemm_type == GemmType.FFN1: + per_rank_bytes = math.ceil(self.precision.activations * m * k * self.experts_per_gpu) + total_bytes = int(math.ceil(per_rank_bytes * self.cp * self.tp )) + reduction_time = self.network_model.collective( + kind="all_to_all", + size_bytes=total_bytes, + participants=self.tp * self.cp, + ib=self.IBTP, + ll=self.LLTP, + local_bytes=0, + debug_label="moe_ffn_b_all_to_all", + + ) + + return gemm_time, reduction_time, total_bytes + def get_embedding_f(self, vocab_size, seq_len, hidden_dim): @@ -1156,7 +1264,7 @@ def get_layernorm_f(self, batch, seq_len, d_model, comm_after=False): - def get_layernorm_b(self, batch, seq_len, d_model, comm_after=False): + def get_layernorm_b(self, batch, seq_len, d_model, type = Optional): tp_mode = self.get_parallelism_mode() seq_degree = self._sequence_parallel_degree() if tp_mode == ParallelismMode.TENSOR_CONTEXT_HYBRID: @@ -1172,7 +1280,21 @@ def get_layernorm_b(self, batch, seq_len, d_model, comm_after=False): name="pointwise-layernorm-b", mem_level=self.num_levels - 1, ) + 4 * self.O - if tp_mode in (ParallelismMode.TENSOR_SEQUENCE, ParallelismMode.TENSOR_CONTEXT_HYBRID): # communication after layernorm + if self.use_moe and type == GemmType.LAYER_NORM_1: # communication after layernorm + per_rank_bytes = self.precision.grad_communication * elements * self.moe_top_k + total_bytes = int(math.ceil(per_rank_bytes * self.tp * self.cp)) + + reduction_time = self.network_model.collective( + kind="all_to_all", + size_bytes=total_bytes, + participants=self.tp * self.cp, + ib=self.IBTP, + ll=self.LLTP, + local_bytes=0, + debug_label="layernorm_b_all_to_all", + ) + + elif tp_mode in (ParallelismMode.TENSOR_SEQUENCE, ParallelismMode.TENSOR_CONTEXT_HYBRID) : per_rank_bytes = self.precision.grad_communication * elements total_bytes = int(math.ceil(per_rank_bytes * self.tp)) reduction_time = self.network_model.collective( @@ -1192,7 +1314,57 @@ def get_layernorm_b(self, batch, seq_len, d_model, comm_after=False): return compute_time, reduction_time, total_bytes + def get_gate_f(self, gemm_gate, gemm_ffn1): + #TODO gate overhead for moe is very minimal, can be ignored, implemented for completeness + _, effective_m, k, n = self._effective_dims(gemm_gate) + effective_m = effective_m / (self.cp) + gemm_time = self.single_gpu_gemm_forward(gemm_gate, "gate_f")[0] + elements = effective_m * n + flops = elements * SOFTMAX_FORWARD_FLOPS_PER_ELEMENT + mem_bytes = self.precision.activations * elements * SOFTMAX_FORWARD_MEM_ACCESSES + + compute_time = gemm_time + self.roofline(flops, mem_bytes, name="pointwise-gate-f", mem_level=self.num_levels - 1) + self.O + + per_rank_bytes = math.ceil(self.precision.activations * gemm_ffn1[0] * gemm_ffn1[1] * self.experts_per_gpu) + total_bytes = int(math.ceil(per_rank_bytes * self.cp * self.tp )) + reduction_time = self.network_model.collective( + kind="all_to_all", + size_bytes=total_bytes, + participants = self.tp * self.cp, + ib=self.IBTP, + ll=self.LLTP, + local_bytes=0, + debug_label="gate_f_all_to_all", + ) + + + return compute_time, reduction_time, total_bytes + def get_gate_b(self, gemm_gate, gemm_ffn1): + #TODO gate overhead for moe is very minimal, can be ignored, implemented for completeness + _, effective_m, k, n = self._effective_dims(gemm_gate) + elements = effective_m * k / (self.tp * self.cp) + + gemm_time = self.single_gpu_gemm_backward(gemm_gate, "gate_f")[0] + + flops = elements * SOFTMAX_BACKWARD_FLOPS_PER_ELEMENT + mem_bytes = self.precision.activations * elements * SOFTMAX_BACKWARD_MEM_ACCESSES + + compute_time = gemm_time + self.roofline(flops, mem_bytes, name="pointwise-gate-f", mem_level=self.num_levels - 1) + self.O + bytes_per_rank = math.ceil(self.precision.activations * elements) + + total_bytes = int(math.ceil(bytes_per_rank * self.tp)) + reduction_time = self.network_model.collective( + kind="reduce_scatter", + size_bytes=total_bytes, + participants=self.tp , + ib=self.IBTP, + ll=self.LLTP, + local_bytes=0, + debug_label="gate_b_reduce_scatter", + ) + + return compute_time, reduction_time, total_bytes def get_embedding_b(self, vocab_size, seq_len, hidden_dim): batch = self._effective_transformer_batch() @@ -1447,6 +1619,7 @@ def compute_all_gemm_and_node_times(self, batch_size, vocab_size, hidden_dim, se gemm_ffn1 = gemm_shapes["ffn1"] gemm_ffn2 = gemm_shapes["ffn2"] gemm_linear = gemm_shapes["linear"] + gemm_gate = gemm_shapes["gate"] transformer_timings: Dict[str, OperationTiming] = {} @@ -1672,13 +1845,52 @@ def _extract_forward(ret: Sequence[Any]) -> Tuple[float, float, float, float, An comm_bytes=out_proj_size_b, ), ) - - ffn1_gemm_f, ffn1_reduction_f, ffn1_size_f, ffn1_flops_f, ffn1_mem_f = _extract_forward( - self.parallelism_gemm_forward(gemm_ffn1, "ffn_f", gemm_type=GemmType.FFN1) - ) - ffn1_gemm_b, ffn1_reduction_b, ffn1_size_b = self.parallelism_gemm_backward( - gemm_ffn1, "ffn_b", gemm_type=GemmType.FFN1 - ) + if not self.use_moe: + ffn1_gemm_f, ffn1_reduction_f, ffn1_size_f, ffn1_flops_f, ffn1_mem_f = _extract_forward( + self.parallelism_gemm_forward(gemm_ffn1, "ffn_f", gemm_type=GemmType.FFN1) + ) + ffn1_gemm_b, ffn1_reduction_b, ffn1_size_b = self.parallelism_gemm_backward( + gemm_ffn1, "ffn_b", gemm_type=GemmType.FFN1 + ) + ffn2_gemm_f, ffn2_reduction_f, ffn2_size_f, ffn2_flops_f, ffn2_mem_f = _extract_forward( + self.parallelism_gemm_forward(gemm_ffn2, "ffn2_f", gemm_type=GemmType.FFN2) + ) + ffn2_gemm_b, ffn2_reduction_b, ffn2_size_b = self.parallelism_gemm_backward( + gemm_ffn2, "ffn2_b", gemm_type=GemmType.FFN2 + ) + + + + else: + gate_gemm_time_f , gate_reduction_f, gate_size_f, _, _ = _extract_forward(self.get_gate_f(gemm_gate, gemm_ffn1)) + gate_gemm_time_b , gate_reduction_b, gate_size_b= self.get_gate_b(gemm_gate, gemm_ffn1) + ffn1_gemm_f, ffn1_reduction_f, ffn1_size_f, ffn1_flops_f, ffn1_mem_f = _extract_forward( + self.get_moe_ffn_f(gemm_ffn1, "ffn1_f", gemm_type=GemmType.FFN1) + ) + ffn1_gemm_b, ffn1_reduction_b, ffn1_size_b = self.get_moe_ffn_b(gemm_ffn1, "ffn1_b", gemm_type=GemmType.FFN1) + ffn2_gemm_f, ffn2_reduction_f, ffn2_size_f, ffn2_flops_f, ffn2_mem_f = _extract_forward( + self.get_moe_ffn_f(gemm_ffn2, "ffn2_f", gemm_type=GemmType.FFN2) + ) + ffn2_gemm_b, ffn2_reduction_b, ffn2_size_b = self.get_moe_ffn_b(gemm_ffn2, "ffn2_b", gemm_type=GemmType.FFN2) + if self.use_moe: + transformer_timings["gate"] = OperationTiming( + "gate", + forward=_make_direction( + "gate", + "forward", + compute_time=gate_gemm_time_f, + comm_time=gate_reduction_f, + comm_bytes=gate_size_f, + ), + backward=_make_direction( + "gate", + "backward", + compute_time=gate_gemm_time_b, + comm_time=gate_reduction_b, + comm_bytes=gate_size_b, + ), + ) + transformer_timings["ffn1"] = OperationTiming( "ffn1", forward=_make_direction( @@ -1699,12 +1911,7 @@ def _extract_forward(ret: Sequence[Any]) -> Tuple[float, float, float, float, An ), ) - ffn2_gemm_f, ffn2_reduction_f, ffn2_size_f, ffn2_flops_f, ffn2_mem_f = _extract_forward( - self.parallelism_gemm_forward(gemm_ffn2, "ffn2_f", gemm_type=GemmType.FFN2) - ) - ffn2_gemm_b, ffn2_reduction_b, ffn2_size_b = self.parallelism_gemm_backward( - gemm_ffn2, "ffn2_b", gemm_type=GemmType.FFN2 - ) + transformer_timings["ffn2"] = OperationTiming( "ffn2", forward=_make_direction( @@ -1757,6 +1964,7 @@ def _extract_forward(ret: Sequence[Any]) -> Tuple[float, float, float, float, An batch=batch_size, seq_len=seq_len, d_model=hidden_dim, + type=GemmType.LAYER_NORM_1, ) transformer_timings["layernorm1"] = OperationTiming( "layernorm1", @@ -2080,13 +2288,36 @@ def _prepare_execution_graphs( transformer_operation_entries: List[Dict[str, Any]] = [] transformer_comm_metadata: Dict[str, Dict[str, Any]] = {} parallelism_mode = self.get_parallelism_mode() - rules_by_mode = COMMUNICATION_RULES.get(parallelism_mode) or {} + def _clone_comm_rules(rules: Dict[str, Dict[str, Optional[Dict[str, str]]]]) -> Dict[str, Dict[str, Optional[Dict[str, str]]]]: + cloned: Dict[str, Dict[str, Optional[Dict[str, str]]]] = {} + for op_name, directions in rules.items(): + cloned[op_name] = {} + for direction, spec in directions.items(): + cloned[op_name][direction] = dict(spec) if spec else None + return cloned + + def _resolve_comm_rules(mode: ParallelismMode) -> Dict[str, Dict[str, Optional[Dict[str, str]]]]: + base_rules = COMMUNICATION_RULES.get(mode) or {} + if not self.use_moe: + return base_rules + overrides = MOE_COMMUNICATION_RULES.get(mode) + if not overrides: + return base_rules + merged = _clone_comm_rules(base_rules) + for op_name, directions in overrides.items(): + merged.setdefault(op_name, {}) + for direction, spec in directions.items(): + merged[op_name][direction] = dict(spec) if spec else None + return merged + + rules_by_mode = _resolve_comm_rules(parallelism_mode) participants_lookup = { 'tp': int(getattr(self, 'tp', 0) or 0), 'cp': int(getattr(self, 'cp', 0) or 0), 'dp': int(getattr(self, 'dp', 0) or 0), 'lp': int(getattr(self, 'lp', 0) or 0), } + participants_lookup['moe'] = max(1, participants_lookup['tp'] * participants_lookup['cp']) group_members = { "MHA": ("qkv_proj", "attention", "output_proj"), "MLP": ("ffn1", "gelu", "ffn2"), @@ -2197,11 +2428,15 @@ def _resolve_operation(name: str) -> OperationTiming: return transformer_timings[name] if parallelism_mode in (ParallelismMode.CONTEXT, ParallelismMode.TENSOR_CONTEXT_HYBRID): - op_names = ("layernorm1", "qkv_proj", "attention", "output_proj", "layernorm2", "MLP") + op_names: Sequence[str] = ("layernorm1", "qkv_proj", "attention", "output_proj", "layernorm2", "MLP") elif parallelism_mode in (ParallelismMode.TENSOR, ParallelismMode.TENSOR_SEQUENCE, ParallelismMode.SINGLE): op_names = ("layernorm1", "MHA", "layernorm2", "MLP") else: raise ValueError(f"Unsupported parallelism mode: {parallelism_mode}") + op_names = list(op_names) + if self.use_moe and "MLP" in op_names and "gate" not in op_names: + mlp_index = op_names.index("MLP") + op_names.insert(mlp_index, "gate") for key in op_names: try: From 6139dea50eb597c9a69f22f4f05b27e26c96d7d7 Mon Sep 17 00:00:00 2001 From: NanoMembers Date: Sun, 9 Nov 2025 21:14:38 -0800 Subject: [PATCH 35/77] flash attention forward and backward implemented --- time_calculation_LLM.py | 128 ++++++++++++++++++++++++++-------------- 1 file changed, 84 insertions(+), 44 deletions(-) diff --git a/time_calculation_LLM.py b/time_calculation_LLM.py index d8f3cb42..58942278 100644 --- a/time_calculation_LLM.py +++ b/time_calculation_LLM.py @@ -394,7 +394,7 @@ def flash_attention_kernel_forward(self, batch_size, hidden_dim, seq_len, num_he # assuming key and value are preloaded into shared memory before attention computation load_kv_bytes = Bc * d * 2 * num_SMs * self.precision.activations #load key and value for one tile from HBM to SRAM initial_load_time = self.roofline(0, load_kv_bytes, "flash_attention_initial_load", mem_level=self.num_levels - 1) #assume key and value of one attention head is loaded from HBM to SRAM - initial_load_time_per_head = initial_load_time * math.ceil(Tc/ num_SMs) + # attention score gemm load_q_bytes = Br * d * self.precision.activations #load query for one tile assuming k is already in shared memory @@ -411,12 +411,13 @@ def flash_attention_kernel_forward(self, batch_size, hidden_dim, seq_len, num_he attn_output_time_per_tile = self.getGEMMTime(Br, Bc, d, "attention_output", read_bytes_l2=output_bytes, write_bytes_l2=output_bytes, flashattn_enable=True)[0] attn_output_time = attn_output_time_per_tile * Tc * Tr #attention output gemm time for one head - load_time = initial_load_time_per_head * batch_size * kv_heads / max(1, self.tp) #load time for key and value are needed once per kv head + attn_score_time *= batch_size * num_heads / max(1, self.tp) + self.O attn_scale_softmax_time *= batch_size * num_heads / max(1, self.tp) attn_output_time *= batch_size * num_heads / max(1, self.tp) + self.O - - attention_forward_gemm_time = load_time + attn_score_time + attn_scale_softmax_time + attn_output_time + + + attention_forward_gemm_time = initial_load_time + attn_score_time + attn_scale_softmax_time + attn_output_time attention_forward_time = attention_forward_gemm_time + attention_forward_reduction_time # HBM traffic consists of only reading Q, K, V once and writing output once @@ -425,47 +426,81 @@ def flash_attention_kernel_forward(self, batch_size, hidden_dim, seq_len, num_he return attention_forward_time, attention_forward_gemm_time, attention_forward_reduction_time, attention_size_f, attention_mem - def flash_attention_kernel_backward(self) -> Tuple[Dict[str, float], Dict[str, float], Dict[str, float]]: - #TODO: should be repleced with a modified version of getgemmtime that considers flash attention tiling - """Return time for flash attention kernel.""" - #FIXME not finished - #TODO gqa support - batch_size = self.microB - seq_len = self.seq_len - hidden_dim = self.hidden_dim - num_heads = self.num_heads + def flash_attention_kernel_backward(self, batch_size, hidden_dim, seq_len, num_heads, kv_heads, num_SMs) -> Tuple[Dict[str, float], Dict[str, float], Dict[str, float]]: + """Return time for flash attention backward kernel.""" + shard_seq = math.ceil(seq_len / max(1, self.cp)) - d = hidden_dim // num_heads #gemm shape for one head is (seq_len, d) x (d, seq_len) - Bc = self.attention_tile_size - Br = min(Bc, d) - Tr = math.ceil(shard_seq / Br) #number of row tiles - Tc = math.ceil(shard_seq / Bc) #number of column tiles - print(f"Flash Attention: Br={Br}, Tr={Tr}, d={d}, Bc={Bc}, Tc={Tc}") - #TODO load 3 tiled gemm inputs and 1 output from HBM to SRAM + d = hidden_dim // num_heads # gemm shape for one head is (seq_len, d) x (d, seq_len) + Bc = self.attention_tile_size # kv tile size + Br = min(Bc, d) # q tile size + Tr = math.ceil(shard_seq / Br) #number of q tiles + Tc = math.ceil(shard_seq / Bc) #number of kv tiles - #TODO disable HBM access - attention_forward_reduction_time = 1 - attention_size_f = 0 + attention_backward_reduction_time = 0 + attention_size_b = 0 + # assuming key and value are preloaded into shared memory before attention backward computation + load_kv_bytes = Bc * d * 2 * num_SMs * self.precision.activations #load key and value for one tile from HBM to SRAM + initial_load_time = self.roofline(0, load_kv_bytes, "flash_attention_initial_load_b", mem_level=self.num_levels - 1) #assume key and value of one attention head is loaded from HBM to SRAM + - attn_score_time = self.getGEMMTime(Br, d, Bc, "attention_score_f")[0] * Tc * Tr #attention score gemm time for one head - number_flops = Br * Bc * d * Tc * Tr * batch_size * num_heads - - # Softmax time + # attention score recompute + load_q_bytes = Br * d * self.precision.activations #load query for one tile assuming k is already in shared memory + attn_score_time_per_tile = self.getGEMMTime(Br, d, Bc, "attention_score_b",read_bytes_l2=load_q_bytes, flashattn_enable=True)[0] #recompute S = QK^T for backward [Br, Bc] + attn_score_time = attn_score_time_per_tile * Tc * Tr #attention score gemm time for one head + + # Softmax recompute time elements = Br * Bc - flops = SOFTMAX_FORWARD_FLOPS_PER_ELEMENT * elements # exponentiation + subtract max + dropout - attn_scale_softmax_time = self.roofline(flops, 1, "attention_scale_softmax", mem_level=self.num_levels - 1) * Tc * Tr #use roofline model for softmax time with no memory access, memory access set to 1 because roofline does not accept 0 memory access - attn_output_time = self.getGEMMTime(Br, Bc, d, "attention_output")[0] * Tc * Tr #attention output gemm time for one head - attn_score_time *= batch_size * num_heads / max(1, self.tp) + flops = SOFTMAX_FORWARD_FLOPS_PER_ELEMENT * elements + attn_scale_softmax_time = self.roofline(flops, 1, "attention_scale_softmax_recompute", mem_level=self.num_levels - 1) * Tc * Tr #use roofline model for softmax time with no memory access, memory access set to 1 because roofline does not accept 0 memory access + + # dV dP + load_o_bytes = Br * d * self.precision.activations #load output for one tile assuming v is already in shared memory + act_dO_time_per_tile = self.getGEMMTime(Bc, Br, d, "attention_output_b", read_bytes_l2=load_o_bytes, flashattn_enable=True)[0] #dV = P^T * dO [Bc, d] + act_dP_time_per_tile = self.getGEMMTime(Br, d, Bc, "attention_score_b", read_bytes_l2=0, flashattn_enable=True)[0] #dP = dO * V^T [Br, Bc] + act_dO_time = act_dO_time_per_tile * Tc * Tr #dV time for one head + act_dP_time = act_dP_time_per_tile * Tc * Tr #dP time for one head + + # compute dS + elements = Br * Bc + flops = (SOFTMAX_BACKWARD_FLOPS_PER_ELEMENT ) * elements + softmax_time_backward = self.roofline(flops, 1, "attention_scale_softmax_b", mem_level=self.num_levels - 1) * Tc * Tr #use roofline model for softmax time with no memory access, memory access set to 1 because roofline does not accept 0 memory access + + # dQ + act_dQ_time_per_tile = self.getGEMMTime(Br, Bc, d, "attention_score_b", read_bytes_l2=0, flashattn_enable=True)[0] #dQ = dS * K [Br, d] + act_dQ_time = act_dQ_time_per_tile * Tc * Tr #dQ time for one head + + + attn_score_time *= batch_size * num_heads / max(1, self.tp) + self.O attn_scale_softmax_time *= batch_size * num_heads / max(1, self.tp) - attn_output_time *= batch_size * num_heads / max(1, self.tp) + act_dO_time *= batch_size * num_heads / max(1, self.tp) + self.O + act_dP_time *= batch_size * num_heads / max(1, self.tp) + self.O + softmax_time_backward *= batch_size * num_heads / max(1, self.tp) + act_dQ_time *= batch_size * num_heads / max(1, self.tp) + self.O - attention_forward_gemm_time = attn_score_time + attn_scale_softmax_time + attn_output_time + attention_backward_gemm_time = initial_load_time + attn_score_time + attn_scale_softmax_time + act_dO_time + act_dP_time + softmax_time_backward + act_dQ_time + attention_backward_time = attention_backward_gemm_time + attention_backward_reduction_time + + + attention_size_b = self.precision.grad_communication * batch_size * seq_len * hidden_dim * 2 / self.tp # weight gradient of K V need to be reduce scattered *2 account for both attn key and value + kind = "reduce_scatter" + participants = self.cp + axis_hint = "cp" + if attention_size_b > 0: + attention_backward_reduction_time = self.network_model.collective( + kind=kind, + size_bytes=attention_size_b, + participants=participants, + ib=self.IBTP, + ll=self.LLTP, + local_bytes=0, + debug_label= "attention_backward_reduction", + axis=axis_hint, + ) + - attention_forward_time = attention_forward_gemm_time + attention_forward_reduction_time - attention_size_f = 0 - return attention_forward_time, attention_forward_gemm_time, attention_forward_reduction_time, attention_size_f + return attention_backward_time, attention_backward_gemm_time, attention_backward_reduction_time, attention_size_b @staticmethod @@ -1705,9 +1740,11 @@ def _extract_forward(ret: Sequence[Any]) -> Tuple[float, float, float, float, An forward=attn_score_forward, backward=attn_score_backward, ) - + + attention_scale_softmax_f = self.get_scale_softmax_f(gemm=gemm_attention_score) attention_scale_softmax_b = self.get_scale_softmax_b(gemm=gemm_attention_score) + transformer_timings["attention_scale_softmax"] = OperationTiming( name="attention_scale_softmax", forward=_make_direction( @@ -1797,10 +1834,15 @@ def _extract_forward(ret: Sequence[Any]) -> Tuple[float, float, float, float, An kv_heads, num_SMs, ) - attention_b = attention_f * 2 - attention_gemm_b = attention_gemm_f * 2 - attention_reduction_b = attention_reduction_f * 2 - attention_size_b = attention_size_f * 2 + attention_b, attention_gemm_b, attention_reduction_b, attention_size_b = self.flash_attention_kernel_backward( + batch_size, + hidden_dim, + seq_len, + num_heads, + kv_heads, + num_SMs, + ) + transformer_timings["attention"] = OperationTiming( "attention", forward=_make_direction( @@ -1858,9 +1900,7 @@ def _extract_forward(ret: Sequence[Any]) -> Tuple[float, float, float, float, An ffn2_gemm_b, ffn2_reduction_b, ffn2_size_b = self.parallelism_gemm_backward( gemm_ffn2, "ffn2_b", gemm_type=GemmType.FFN2 ) - - - + else: gate_gemm_time_f , gate_reduction_f, gate_size_f, _, _ = _extract_forward(self.get_gate_f(gemm_gate, gemm_ffn1)) gate_gemm_time_b , gate_reduction_b, gate_size_b= self.get_gate_b(gemm_gate, gemm_ffn1) From 7c57274ea59ea86871ead42c92e72895c7b81eef Mon Sep 17 00:00:00 2001 From: NanoMembers Date: Mon, 10 Nov 2025 21:26:45 -0800 Subject: [PATCH 36/77] minor clean up and fix --- LLM_util.py | 2 +- time_calculation_LLM.py | 70 ++++++++++++++++++------------------- validation_scripts/koyeb.py | 2 +- 3 files changed, 37 insertions(+), 37 deletions(-) diff --git a/LLM_util.py b/LLM_util.py index 29851fd4..110af2bf 100644 --- a/LLM_util.py +++ b/LLM_util.py @@ -83,7 +83,7 @@ def multihead_decoder_gemm(self, batch_size, seq_len, d_model, num_heads, kv_hea gemms["ffn1"] = (batch_size, seq_len, d_model, projected_dim) gemms["ffn2"] = (batch_size, seq_len, intermediate_size, d_model) gemms["linear"] = (batch_size, seq_len, d_model, vocab_size) - gemms["gate"] = (batch_size, seq_len, d_model, self.moe_num_experts) + gemms["router"] = (batch_size, seq_len, d_model, self.moe_num_experts) return gemms diff --git a/time_calculation_LLM.py b/time_calculation_LLM.py index 58942278..e76783ea 100644 --- a/time_calculation_LLM.py +++ b/time_calculation_LLM.py @@ -111,8 +111,8 @@ class ParallelismMode(Enum): COMM_RULE_DEFAULT_KEY: {'forward': None, 'backward': None}, }, } - -_MOE_GATE_RULE = { +# MoE-specific overrides for router, MLP, and layer norm collectives. +_MOE_ROUTER_RULE = { 'forward': { 'kind': 'all_to_all', 'participants': 'moe', @@ -152,7 +152,7 @@ class ParallelismMode(Enum): ParallelismMode, Dict[str, Dict[str, Optional[Dict[str, str]]]] ] = { mode: { - 'gate': _MOE_GATE_RULE, + 'router': _MOE_ROUTER_RULE, 'MLP': _MOE_MLP_RULE, 'layernorm1': _MOE_LAYER_NORM1_RULE, } @@ -1349,16 +1349,16 @@ def get_layernorm_b(self, batch, seq_len, d_model, type = Optional): return compute_time, reduction_time, total_bytes - def get_gate_f(self, gemm_gate, gemm_ffn1): - #TODO gate overhead for moe is very minimal, can be ignored, implemented for completeness - _, effective_m, k, n = self._effective_dims(gemm_gate) + def get_router_f(self, gemm_router, gemm_ffn1): + #TODO router overhead for moe is very minimal, can be ignored, implemented for completeness + _, effective_m, k, n = self._effective_dims(gemm_router) effective_m = effective_m / (self.cp) - gemm_time = self.single_gpu_gemm_forward(gemm_gate, "gate_f")[0] + gemm_time = self.single_gpu_gemm_forward(gemm_router, "router_f")[0] elements = effective_m * n flops = elements * SOFTMAX_FORWARD_FLOPS_PER_ELEMENT mem_bytes = self.precision.activations * elements * SOFTMAX_FORWARD_MEM_ACCESSES - compute_time = gemm_time + self.roofline(flops, mem_bytes, name="pointwise-gate-f", mem_level=self.num_levels - 1) + self.O + compute_time = gemm_time + self.roofline(flops, mem_bytes, name="pointwise-router-f", mem_level=self.num_levels - 1) + self.O per_rank_bytes = math.ceil(self.precision.activations * gemm_ffn1[0] * gemm_ffn1[1] * self.experts_per_gpu) total_bytes = int(math.ceil(per_rank_bytes * self.cp * self.tp )) @@ -1369,22 +1369,22 @@ def get_gate_f(self, gemm_gate, gemm_ffn1): ib=self.IBTP, ll=self.LLTP, local_bytes=0, - debug_label="gate_f_all_to_all", + debug_label="router_f_all_to_all", ) return compute_time, reduction_time, total_bytes - def get_gate_b(self, gemm_gate, gemm_ffn1): - #TODO gate overhead for moe is very minimal, can be ignored, implemented for completeness - _, effective_m, k, n = self._effective_dims(gemm_gate) + def get_router_b(self, gemm_router, gemm_ffn1): + #TODO router overhead for moe is very minimal, can be ignored, implemented for completeness + _, effective_m, k, n = self._effective_dims(gemm_router) elements = effective_m * k / (self.tp * self.cp) - gemm_time = self.single_gpu_gemm_backward(gemm_gate, "gate_f")[0] + gemm_time = self.single_gpu_gemm_backward(gemm_router, "router_f")[0] flops = elements * SOFTMAX_BACKWARD_FLOPS_PER_ELEMENT mem_bytes = self.precision.activations * elements * SOFTMAX_BACKWARD_MEM_ACCESSES - compute_time = gemm_time + self.roofline(flops, mem_bytes, name="pointwise-gate-f", mem_level=self.num_levels - 1) + self.O + compute_time = gemm_time + self.roofline(flops, mem_bytes, name="pointwise-router-f", mem_level=self.num_levels - 1) + self.O bytes_per_rank = math.ceil(self.precision.activations * elements) total_bytes = int(math.ceil(bytes_per_rank * self.tp)) @@ -1395,7 +1395,7 @@ def get_gate_b(self, gemm_gate, gemm_ffn1): ib=self.IBTP, ll=self.LLTP, local_bytes=0, - debug_label="gate_b_reduce_scatter", + debug_label="router_b_reduce_scatter", ) @@ -1654,7 +1654,7 @@ def compute_all_gemm_and_node_times(self, batch_size, vocab_size, hidden_dim, se gemm_ffn1 = gemm_shapes["ffn1"] gemm_ffn2 = gemm_shapes["ffn2"] gemm_linear = gemm_shapes["linear"] - gemm_gate = gemm_shapes["gate"] + gemm_router = gemm_shapes["router"] transformer_timings: Dict[str, OperationTiming] = {} @@ -1887,7 +1887,7 @@ def _extract_forward(ret: Sequence[Any]) -> Tuple[float, float, float, float, An comm_bytes=out_proj_size_b, ), ) - if not self.use_moe: + if not self.use_moe: ffn1_gemm_f, ffn1_reduction_f, ffn1_size_f, ffn1_flops_f, ffn1_mem_f = _extract_forward( self.parallelism_gemm_forward(gemm_ffn1, "ffn_f", gemm_type=GemmType.FFN1) ) @@ -1901,9 +1901,9 @@ def _extract_forward(ret: Sequence[Any]) -> Tuple[float, float, float, float, An gemm_ffn2, "ffn2_b", gemm_type=GemmType.FFN2 ) - else: - gate_gemm_time_f , gate_reduction_f, gate_size_f, _, _ = _extract_forward(self.get_gate_f(gemm_gate, gemm_ffn1)) - gate_gemm_time_b , gate_reduction_b, gate_size_b= self.get_gate_b(gemm_gate, gemm_ffn1) + else: # MOE expert parallelism + router_gemm_time_f , router_reduction_f, router_size_f, _, _ = _extract_forward(self.get_router_f(gemm_router, gemm_ffn1)) + router_gemm_time_b , router_reduction_b, router_size_b= self.get_router_b(gemm_router, gemm_ffn1) ffn1_gemm_f, ffn1_reduction_f, ffn1_size_f, ffn1_flops_f, ffn1_mem_f = _extract_forward( self.get_moe_ffn_f(gemm_ffn1, "ffn1_f", gemm_type=GemmType.FFN1) ) @@ -1913,21 +1913,21 @@ def _extract_forward(ret: Sequence[Any]) -> Tuple[float, float, float, float, An ) ffn2_gemm_b, ffn2_reduction_b, ffn2_size_b = self.get_moe_ffn_b(gemm_ffn2, "ffn2_b", gemm_type=GemmType.FFN2) if self.use_moe: - transformer_timings["gate"] = OperationTiming( - "gate", + transformer_timings["router"] = OperationTiming( + "router", forward=_make_direction( - "gate", + "router", "forward", - compute_time=gate_gemm_time_f, - comm_time=gate_reduction_f, - comm_bytes=gate_size_f, + compute_time=router_gemm_time_f, + comm_time=router_reduction_f, + comm_bytes=router_size_f, ), backward=_make_direction( - "gate", + "router", "backward", - compute_time=gate_gemm_time_b, - comm_time=gate_reduction_b, - comm_bytes=gate_size_b, + compute_time=router_gemm_time_b, + comm_time=router_reduction_b, + comm_bytes=router_size_b, ), ) @@ -2474,9 +2474,9 @@ def _resolve_operation(name: str) -> OperationTiming: else: raise ValueError(f"Unsupported parallelism mode: {parallelism_mode}") op_names = list(op_names) - if self.use_moe and "MLP" in op_names and "gate" not in op_names: + if self.use_moe and "MLP" in op_names and "router" not in op_names: # if MOE is enabled, ensure router is included in the graph mlp_index = op_names.index("MLP") - op_names.insert(mlp_index, "gate") + op_names.insert(mlp_index, "router") for key in op_names: try: @@ -2552,7 +2552,7 @@ def _resolve_operation(name: str) -> OperationTiming: graph_root_no_dp = None need_no_dp_variant = getattr(self, "gradient_accumulation_steps", 1) > 1 if need_no_dp_variant: - pipeline_graph_obj_no_dp = simulate_LLM.Graph( + pipeline_graph_obj_no_dp = simulate_LLM.Graph( # if gradient accumulation is used, we need a no-dp variant of the graph for the non-last step, since dp all-reduce is only done on the last step mode="pipeline", dp=1, lp=self.lp, @@ -2782,10 +2782,10 @@ def calc_time_llm(self): print(f"Actual transformer forward time: {self.transformer_astrasim_time_forward:.4f}s") print(f"Actual transformer backward time: {self.transformer_astrasim_time_backward:.4f}s") - if self.gradient_accumulation_steps > 1: #if gradient accumulation is enabled, non-last steps have no dp overhead, only last step has dp overhead + if self.gradient_accumulation_steps > 1: if time_fw_bw_no_dp is None: raise RuntimeError("Missing no-DP timing for gradient accumulation computation") - self.tot_time = time_fw_bw_no_dp * (self.gradient_accumulation_steps - 1) + time_fw_bw + self.tot_time = time_fw_bw_no_dp * (self.gradient_accumulation_steps - 1) + time_fw_bw #total time is (time for no-dp steps) + (time for last step with dp) else: self.tot_time = time_fw_bw return self.tot_time diff --git a/validation_scripts/koyeb.py b/validation_scripts/koyeb.py index 90be11c0..b883fdb7 100644 --- a/validation_scripts/koyeb.py +++ b/validation_scripts/koyeb.py @@ -6,7 +6,7 @@ import matplotlib.pyplot as plt -from validation_helpers import ( +from .validation_helpers import ( ValidationSpec, parse_decode_time, run_validation_suite, From 7db2824bfbe8f0b373931386ac8a93ffffd190fb Mon Sep 17 00:00:00 2001 From: gkarfakis19 Date: Sat, 15 Nov 2025 16:31:43 -0800 Subject: [PATCH 37/77] attempt to add CI --- .github/workflows/ci.yml | 44 +++++++++++++++++++++++++++++++++++++ scripts/install_astrasim.sh | 34 ++++++++++++++++++++++++++++ 2 files changed, 78 insertions(+) create mode 100644 .github/workflows/ci.yml create mode 100755 scripts/install_astrasim.sh diff --git a/.github/workflows/ci.yml b/.github/workflows/ci.yml new file mode 100644 index 00000000..60904ea2 --- /dev/null +++ b/.github/workflows/ci.yml @@ -0,0 +1,44 @@ +name: CI + +on: + push: + branches: + - main + - llm-dev + pull_request: + +jobs: + tests: + runs-on: ubuntu-latest + steps: + - name: Checkout repository + uses: actions/checkout@v4 + with: + submodules: recursive + + - name: Install uv + uses: astral-sh/setup-uv@v2 + + - name: Sync Python dependencies + run: uv sync + + - name: Determine AstraSim commit + run: echo "ASTRASIM_SHA=$(git -C astra-sim rev-parse HEAD)" >> $GITHUB_ENV + + - name: Restore AstraSim build cache + id: ascache + uses: actions/cache@v4 + with: + path: astra-sim/build/astra_analytical/build + key: astra-${{ env.ASTRASIM_SHA }} + + - name: Build AstraSim (cache miss) + if: steps.ascache.outputs.cache-hit != 'true' + run: | + chmod +x scripts/install_astrasim.sh + scripts/install_astrasim.sh + + - name: Run pytest + env: + DEEPFLOW_VALIDATION_WORKERS: "1" + run: uv run pytest diff --git a/scripts/install_astrasim.sh b/scripts/install_astrasim.sh new file mode 100755 index 00000000..5dbf05ef --- /dev/null +++ b/scripts/install_astrasim.sh @@ -0,0 +1,34 @@ +#!/usr/bin/env bash +set -euo pipefail + +SCRIPT_DIR="$(cd "$(dirname "${BASH_SOURCE[0]}")" && pwd)" +REPO_ROOT="$(cd "${SCRIPT_DIR}/.." && pwd)" +ASTRA_SIM_DIR="${REPO_ROOT}/astra-sim" + +cd "${REPO_ROOT}" + +echo "[astrasim] Updating DeepFlow submodules..." +git submodule update --init --recursive +git submodule update --remote -- astra-sim + +if [[ ! -d "${ASTRA_SIM_DIR}" ]]; then + echo "[astrasim] ERROR: expected submodule directory at ${ASTRA_SIM_DIR}" >&2 + exit 1 +fi + +echo "[astrasim] Updating AstraSim nested submodules..." +git -C "${ASTRA_SIM_DIR}" submodule update --init --recursive +git -C "${ASTRA_SIM_DIR}" submodule update --remote -- extern/network_backend/analytical + +BUILD_SCRIPT="${ASTRA_SIM_DIR}/build/astra_analytical/build.sh" +if [[ ! -x "${BUILD_SCRIPT}" ]]; then + echo "[astrasim] ERROR: build script not found or not executable at ${BUILD_SCRIPT}" >&2 + exit 1 +fi + +echo "[astrasim] Building AstraSim (this may take a while)..." +pushd "${ASTRA_SIM_DIR}" >/dev/null +"${BUILD_SCRIPT}" +popd >/dev/null + +echo "[astrasim] AstraSim build completed successfully." From 02abd950cb9701249664a107c6280044dbbc7a67 Mon Sep 17 00:00:00 2001 From: gkarfakis19 Date: Sat, 15 Nov 2025 16:38:55 -0800 Subject: [PATCH 38/77] minor fix --- scripts/install_astrasim.sh | 1 - 1 file changed, 1 deletion(-) diff --git a/scripts/install_astrasim.sh b/scripts/install_astrasim.sh index 5dbf05ef..cc3074ae 100755 --- a/scripts/install_astrasim.sh +++ b/scripts/install_astrasim.sh @@ -17,7 +17,6 @@ if [[ ! -d "${ASTRA_SIM_DIR}" ]]; then fi echo "[astrasim] Updating AstraSim nested submodules..." -git -C "${ASTRA_SIM_DIR}" submodule update --init --recursive git -C "${ASTRA_SIM_DIR}" submodule update --remote -- extern/network_backend/analytical BUILD_SCRIPT="${ASTRA_SIM_DIR}/build/astra_analytical/build.sh" From 65bdd6cd1d66f0f902f1567e9b36dc87ce713774 Mon Sep 17 00:00:00 2001 From: gkarfakis19 Date: Sat, 15 Nov 2025 16:43:15 -0800 Subject: [PATCH 39/77] dependencies for CI --- .github/workflows/ci.yml | 12 +++++++++++- 1 file changed, 11 insertions(+), 1 deletion(-) diff --git a/.github/workflows/ci.yml b/.github/workflows/ci.yml index 60904ea2..b15d11e0 100644 --- a/.github/workflows/ci.yml +++ b/.github/workflows/ci.yml @@ -6,7 +6,7 @@ on: - main - llm-dev pull_request: - + jobs: tests: runs-on: ubuntu-latest @@ -16,6 +16,16 @@ jobs: with: submodules: recursive + - name: Install AstraSim system dependencies + run: | + sudo apt-get update + sudo apt-get install -y \ + gcc-11 g++-11 make cmake \ + clang-format \ + libboost-dev libboost-program-options-dev \ + libprotobuf-dev protobuf-compiler \ + openmpi-bin openmpi-doc libopenmpi-dev + - name: Install uv uses: astral-sh/setup-uv@v2 From acd157d104e3d9a89e16cd5ad690151673e1b20a Mon Sep 17 00:00:00 2001 From: gkarfakis19 Date: Sat, 15 Nov 2025 16:49:20 -0800 Subject: [PATCH 40/77] test if CI fails --- configs/hardware-config/a100_80GB_no_parallelism.yaml | 2 +- 1 file changed, 1 insertion(+), 1 deletion(-) diff --git a/configs/hardware-config/a100_80GB_no_parallelism.yaml b/configs/hardware-config/a100_80GB_no_parallelism.yaml index 798101a2..23260c14 100644 --- a/configs/hardware-config/a100_80GB_no_parallelism.yaml +++ b/configs/hardware-config/a100_80GB_no_parallelism.yaml @@ -28,7 +28,7 @@ sw_param: h2d_bandwidth: -1 # Bytes/s; set -1 to disable H2D timing tech_param: core: - operating_frequency: 1.41e9 + operating_frequency: 1.41e6 num_bundles: 108 nominal_flop_rate_per_mcu: 512 nominal_energy_per_flop: 1.8e-13 From 72ff52ef3869f8d6f905f370428ba7d0d12bf693 Mon Sep 17 00:00:00 2001 From: gkarfakis19 Date: Sat, 15 Nov 2025 16:53:04 -0800 Subject: [PATCH 41/77] revert test if CI fails --- configs/hardware-config/a100_80GB_no_parallelism.yaml | 2 +- 1 file changed, 1 insertion(+), 1 deletion(-) diff --git a/configs/hardware-config/a100_80GB_no_parallelism.yaml b/configs/hardware-config/a100_80GB_no_parallelism.yaml index 23260c14..798101a2 100644 --- a/configs/hardware-config/a100_80GB_no_parallelism.yaml +++ b/configs/hardware-config/a100_80GB_no_parallelism.yaml @@ -28,7 +28,7 @@ sw_param: h2d_bandwidth: -1 # Bytes/s; set -1 to disable H2D timing tech_param: core: - operating_frequency: 1.41e6 + operating_frequency: 1.41e9 num_bundles: 108 nominal_flop_rate_per_mcu: 512 nominal_energy_per_flop: 1.8e-13 From a359b164d0a9d33bc42aa59af1bfc4c18d798ccf Mon Sep 17 00:00:00 2001 From: NanoMembers Date: Tue, 18 Nov 2025 13:38:04 -0800 Subject: [PATCH 42/77] Add NVIDIA validation tests and update validation helpers for inference time parsing --- configs/hardware-config/H100_SXM5_80GB.yaml | 69 +++-- configs/model-config/Llama2-13B_inf.yaml | 21 ++ configs/model-config/Llama2-70B_inf.yaml | 21 ++ tests/test_nvidia_validation.py | 17 ++ validation_scripts/nvidia_inf.csv | 23 ++ validation_scripts/nvidia_inf.py | 285 ++++++++++++++++++++ validation_scripts/validation_helpers.py | 42 +++ 7 files changed, 455 insertions(+), 23 deletions(-) create mode 100644 configs/model-config/Llama2-13B_inf.yaml create mode 100644 configs/model-config/Llama2-70B_inf.yaml create mode 100644 tests/test_nvidia_validation.py create mode 100644 validation_scripts/nvidia_inf.csv create mode 100644 validation_scripts/nvidia_inf.py diff --git a/configs/hardware-config/H100_SXM5_80GB.yaml b/configs/hardware-config/H100_SXM5_80GB.yaml index 5c2b8e53..be3fbdfc 100644 --- a/configs/hardware-config/H100_SXM5_80GB.yaml +++ b/configs/hardware-config/H100_SXM5_80GB.yaml @@ -8,7 +8,7 @@ sw_param: kernel_launch_overhead: 9e-6 # (s) # dp_zero_stage selects ZeRO-style data-parallel sharding: 0=DDP, 1=ZeRO-1, 2=ZeRO-2, 3=ZeRO-3 # only used if data parallelism (dp) > 1 - dp_zero_stage: 0 + dp_zero_stage: 1 precision: # tensor_format accepts numeric byte counts or dtype strings (fp16, bf16, fp32). # Take care that the number of flops defined below has been carefully tuned for FP16/BF16 formats. @@ -19,7 +19,7 @@ sw_param: # -> false forces all activations/weights/gradients/optimizer states/stats to use tensor_format bytes. # -> true enables AMP-style FP32 accumulations/optimizer states depending on param_storage_mode. # if true, optimizer states and stats are always FP32. - mixed_precision: false + mixed_precision: true # param_storage_mode: as_tensor_format | tensor_plus_fp32_master | fp32_params. # -> as_tensor_format makes weights/gradients use tensor_format bytes. # -> tensor_plus_fp32_master makes weights/gradients use tensor_format bytes and @@ -28,8 +28,7 @@ sw_param: param_storage_mode: as_tensor_format # kv_cache precision may be set as bytes, dtype string, or "as_tensor_format". kv_cache: as_tensor_format - h2d_bandwidth: 13297674240 # Bytes/s; set -1 to disable H2D timing - + h2d_bandwidth: -1 # Bytes/s; set -1 to disable H2D timing tech_param: core: operating_frequency: 1.98e9 # (Hz) GPU boost clock frequency [TePU] @@ -37,13 +36,13 @@ tech_param: nominal_flop_rate_per_mcu: 1024 # F16 FLOP / cycle / SM is 4096 [Semi] nominal_energy_per_flop: 1.8e-13 num_mcu_per_bundle: 4 # number of tensorcores per SM [TePU] - FMA_d1: 8 # tensor core can compute the results for an d1 x d2 x d1 mixed-precision matrix multiplication per clock - FMA_d2: 4 - dataflow: best #{wst. ast, ost, best, none}: wst: weight stationary, ast: activation stationary, ost: output stationary - util: 0.75 #util should be 0.75 (12/16) for tensorcore + FMA_d1: 16 # tensor core can compute the results for an d1 x d2 x d1 mixed-precision matrix multiplication per clock + FMA_d2: 8 + dataflow: best + util: 1.0 DRAM: size: 80 GB # [TePU] - bandwidth: 3440 GB # (/s) [TePU] + bandwidth: 3430 GB # (/s) [TePU] dynamic_energy_per_bit: 4.5e-12 # (J) [https://passlab.github.io/mchpc/mchpc2019/presentations/MCHPC_Pawlowski_keynote.pdf, slide 36] latency: 100e-9 # (s) util: 1 @@ -61,7 +60,7 @@ tech_param: util: 1 SRAM-R: size: 33 MB # register file memory per SM * number of SMs [NVIDIA] - bandwidth: 244 TB # traffic for 2 reads and 1 write for GEMMs computed per clock + bandwidth: 122 TB # traffic for 2 reads and 1 write for GEMMs computed per clock dynamic_energy_per_bit: 110e-15 latency: 0 util: 1 @@ -84,31 +83,55 @@ parallelism: auto: false ### LLM PARAMS ### dp: 1 # data parallelism. Acts as replica count during inference + tp: 8 # tensor parallelism + tp_sp: false # tensor-sequence parallelism + cp: 1 # context parallelism. Not supported for inference. lp: 1 # pipeline parallelism. Number of pipeline stages - mb: 1 # micro-batches (for best results set to >= lp) - tp: 1 # tensor parallelism - tp_sp: true # tensor-sequence parallelism - cp : 1 # context parallelism + # mb is not a true parallelism parameter, does not affect num of GPUs. + mb: 1 # micro-batches (for best results set to >= lp). network: dimensions: - id: dim0 - label: device_fabric - size: 1 + label: torus2d_tp + size: auto + topology: + type: Ring + bandwidth: 100 GB # We assume 8 25GB/s NVLINKs in the torus. Torus2D is hardcoded bidirectional, so it's actually 4 links, 25GB/s each (because they get replicated). So 100GB/s total. + latency: 5e-6 + energy_per_bit: 8e-12 # (J/bit) + util: 1.0 + optimize_2dmap: false + collective_override: {} # overrides execution_backend.collectives when provided + parallelisms: [tp, cp] + - id: dim1 + label: inter_torus_infiniband + size: auto + topology: + type: Ring + bandwidth: 25 GB + latency: 5e-6 + energy_per_bit: 8e-12 # (J/bit) + util: 1.0 + collective_override: {} # overrides execution_backend.collectives when provided + parallelisms: [lp,dp] + - id: dim2 + label: inter_torus_infiniband_dp + size: auto topology: - type: Mesh - bandwidth: 400 Gb + type: Ring + bandwidth: 25 GB latency: 5e-6 energy_per_bit: 8e-12 # (J/bit) util: 1.0 + parallelisms: [] collective_override: {} # overrides execution_backend.collectives when provided - parallelisms: [tp, cp, lp, dp] execution_backend: - model: astra # analytical | astra + model: analytical # analytical | astra astra: backend: astra - mode: hybrid # hybrid | full_astrasim_hierarchical | full_astrasim_flattened + mode: full_astrasim_flattened # hybrid | full_astrasim_hierarchical | full_astrasim_flattened collectives: # Default collective algorithms. Per-dimension collective_override entries can override these. # Allowed values per op: auto | ring | direct | halvingDoubling | doubleBinaryTree | mesh | mesh2d | hypercube | torus2d @@ -118,9 +141,9 @@ execution_backend: reduce_scatter: auto all_to_all: auto sys_options: - endpoint_delay: 10 + endpoint_delay: 5 active_chunks_per_dimension: 32 # Higher value means significantly slower astrasim runtime, but more accurate especially for tensor parallelism - preferred_dataset_splits: 32 # Should be same as active_chunks_per_dimension + preferred_dataset_splits: 4 # Should be same as active_chunks_per_dimension inference: kvcache_type: hbm_only diff --git a/configs/model-config/Llama2-13B_inf.yaml b/configs/model-config/Llama2-13B_inf.yaml new file mode 100644 index 00000000..a5105fa0 --- /dev/null +++ b/configs/model-config/Llama2-13B_inf.yaml @@ -0,0 +1,21 @@ +model_param: + mode: "LLM" # Mode of operation + run_type: "inference" # Type of run: "training" or "inference" + tied_embeddings: true # Share input and output embedding weights (WIP: affects memory estimation only) + model_type: "llama" # gpt assumes dense-style architecture with GELU; llama assumes llama-style architecture with SwiGLU + global_batch_size: 16 # Sequences per global batch + gradient_accumulation_steps: 1 # Number of gradient accumulation micro-steps per optimizer update + seq_len: 2048 # Total context length (prefill + decode_len for inference) + decode_len: 1024 # Number of tokens to decode during inference + hidden_dim: 5120 # Transformer hidden dimension + attention: # Attention-related parameters + attention_type: "mha" # Supported attention types: "mha" or "gqa" ("mla" not implemented) + num_heads: 40 # Total number of attention heads + kv_heads: 8 # Number of key/value heads (used only when attention_type = "gqa") + use_flashattention: true # Enable FlashAttention optimization + attention_tile_size: 256 # Tile size used in attention computation + num_experts: 1 # Number of experts in MoE layers. Set to 1 to disable MoE. + top_k: 1 # Number of active experts per token in MoE + intermediate_size: 13824 # Output dimension of the first FFN layer (typically 4×hidden_dim) + vocab_size: 32000 # Vocabulary size + num_layers: 40 # Number of transformer layers diff --git a/configs/model-config/Llama2-70B_inf.yaml b/configs/model-config/Llama2-70B_inf.yaml new file mode 100644 index 00000000..24facac3 --- /dev/null +++ b/configs/model-config/Llama2-70B_inf.yaml @@ -0,0 +1,21 @@ +model_param: + mode: "LLM" # Mode of operation + run_type: "inference" # Type of run: "training" or "inference" + tied_embeddings: true # Share input and output embedding weights (WIP: affects memory estimation only) + model_type: "llama" # gpt assumes dense-style architecture with GELU; llama assumes llama-style architecture with SwiGLU + global_batch_size: 16 # Sequences per global batch + gradient_accumulation_steps: 1 # Number of gradient accumulation micro-steps per optimizer update + seq_len: 2048 # Total context length (prefill + decode_len for inference) + decode_len: 1024 # Number of tokens to decode during inference + hidden_dim: 8192 # Transformer hidden dimension + attention: # Attention-related parameters + attention_type: "gqa" # Supported attention types: "mha" or "gqa" ("mla" not implemented) + num_heads: 64 # Total number of attention heads + kv_heads: 8 # Number of key/value heads (used only when attention_type = "gqa") + use_flashattention: true # Enable FlashAttention optimization + attention_tile_size: 256 # Tile size used in attention computation + num_experts: 1 # Number of experts in MoE layers. Set to 1 to disable MoE. + top_k: 1 # Number of active experts per token in MoE + intermediate_size: 28672 # Output dimension of the first FFN layer (typically 4×hidden_dim) + vocab_size: 32000 # Vocabulary size + num_layers: 80 # Number of transformer layers diff --git a/tests/test_nvidia_validation.py b/tests/test_nvidia_validation.py new file mode 100644 index 00000000..a2ef5095 --- /dev/null +++ b/tests/test_nvidia_validation.py @@ -0,0 +1,17 @@ +import math +import sys +from pathlib import Path + +PROJECT_ROOT = Path(__file__).resolve().parents[1] +if str(PROJECT_ROOT) not in sys.path: + sys.path.insert(0, str(PROJECT_ROOT)) + +from validation_scripts import nvidia_inf + +def test_avg_abs_error_below_threshold_network_ignored(): + results = nvidia_inf.run(enable_plot=True, network_ignored=True) + avg_abs_error = results["avg_abs_error"] + assert not math.isnan(avg_abs_error), "NVIDIA validation produced no valid measurements" + # Historically the average absolute error bottoms out at roughly 5.97%, so we allow + # a little slack (6.5%) to absorb minor DeepFlow changes while still flagging regressions. + assert avg_abs_error <= 17.0, f"Average absolute error {avg_abs_error:.2f}% exceeds 17.0%" diff --git a/validation_scripts/nvidia_inf.csv b/validation_scripts/nvidia_inf.csv new file mode 100644 index 00000000..7dcf5fca --- /dev/null +++ b/validation_scripts/nvidia_inf.csv @@ -0,0 +1,23 @@ +device,model,TP,actual,seconds +A100,Llama 2-7B,1,2.19,2.36 +A100,Llama 2-7B,2,1.544,1.82 +A100,Llama 2-7B,4,1.28,1.77 +A100,Llama 2-7B,8,1.187,2.16 +A100,Llama 2-13B,1,3.884,3.84 +A100,Llama 2-13B,2,2.499,2.75 +A100,Llama 2-13B,4,1.894,2.46 +A100,Llama 2-13B,8,1.693,2.85 +A100,Llama 2-70B,2,10.5,9.55 +A100,Llama 2-70B,4,6.403,6.98 +A100,Llama 2-70B,8,4.735,6.74 +H100,Llama 2-7B,1,1.44,1.82 +H100,Llama 2-7B,2,1.143,1.55 +H100,Llama 2-7B,4,0.924,1.62 +H100,Llama 2-7B,8,0.828,2.06 +H100,Llama 2-13B,1,2.396,2.75 +H100,Llama 2-13B,2,1.717,2.2 +H100,Llama 2-13B,4,1.431,2.17 +H100,Llama 2-13B,8,1.201,2.67 +H100,Llama 2-70B,2,6.267,6.77 +H100,Llama 2-70B,4,4.1226,5.56 +H100,Llama 2-70B,8,3.202,6 diff --git a/validation_scripts/nvidia_inf.py b/validation_scripts/nvidia_inf.py new file mode 100644 index 00000000..9983da22 --- /dev/null +++ b/validation_scripts/nvidia_inf.py @@ -0,0 +1,285 @@ +import argparse +import math +import os +from pathlib import Path +import pandas as pd +from typing import List +import matplotlib.pyplot as plt + +from .validation_helpers import ( + ValidationSpec, + run_validation_suite, + parse_inference_time, +) + +HW_CONFIGS = { + "A100": "a100_80GB.yaml", + "H100": "H100_SXM5_80GB.yaml", +} + +MODEL_CONFIGS = { + "Llama 2-7B": "Llama2-7B_inf.yaml", + "Llama 2-13B": "Llama2-13B_inf.yaml", + "Llama 2-70B": "Llama2-70B_inf.yaml", +} + +PROJECT_ROOT = os.path.dirname(os.path.dirname(os.path.abspath(__file__))) +RUN_PERF = os.path.join(PROJECT_ROOT, "run_perf.py") +HARDWARE_CONFIG_PATH = os.path.join(PROJECT_ROOT, "configs/hardware-config") +MODEL_CONFIG_PATH = os.path.join(PROJECT_ROOT, "configs/model-config") + + +def _load_data(csv_path: Path) -> pd.DataFrame: + # Skip the leading "// filepath: ..." line by treating '/' as a comment char. + df = pd.read_csv(csv_path, comment="/") + # Normalize dtypes + df["TP"] = df["TP"].astype(int) + df["actual"] = pd.to_numeric(df["actual"], errors="coerce") + df["seconds"] = pd.to_numeric(df["seconds"], errors="coerce") + return df.dropna(subset=["device", "model", "TP", "actual", "seconds"]) + + +def _iter_tests(df: pd.DataFrame): + for _, row in df.iterrows(): + yield (row["device"], row["model"], row["TP"]) + + +def _run_test( + device: str, model: str, tp: int, idx: int, network_ignored: bool = True +) -> float: + print(f"Running test {idx + 1}: Device={device}, Model={model}, TP={tp}") + specs: List[ValidationSpec] = [] + label = f"{device} {model} TP={tp}" + + model_overrides = { + "model_param": { + "global_batch_size": 1, + "seq_len": 400, + "decode_len": 200 + } + } + + hw_overrides = { + "parallelism": { + "dp": 1, + "tp": int(tp), + "tp_sp": False, + "cp": 1, + "lp": 1, + "mb": 1, + } + } + if network_ignored: + override_bw = { + "network": { + "dimensions": [ + { + "id": "dim0", + "topology": {"bandwidth": "100000 GB", "latency": "0"}, + } + ] + }, + "execution_backend": { + "model": "analytical", + }, + } + hw_overrides.update(override_bw) + + specs.append( + ValidationSpec( + label=label, + model_overrides=model_overrides, + hardware_overrides=hw_overrides, + order=idx, + ) + ) + + hw_config_path = os.path.join(HARDWARE_CONFIG_PATH, HW_CONFIGS.get(device)) + model_config_path = os.path.join(MODEL_CONFIG_PATH, MODEL_CONFIGS.get(model)) + + validation_results = run_validation_suite( + specs, + base_model_config_path=model_config_path, + base_hardware_config_path=hw_config_path, + result_parser=parse_inference_time, + run_perf_path=RUN_PERF, + ) + + result = validation_results[0] + if result.success: + inf_time_s = float(result.metrics.get("inference_time_s", float("nan"))) + else: + inf_time_s = float("nan") + err_detail = result.error or "DeepFlow run failed" + print(f" ERROR: {err_detail}") + if result.raw_output: + print(result.raw_output.strip()) + + return inf_time_s + + +def plot_device(df: pd.DataFrame, device: str, outdir: Path) -> Path | None: + sub = df[df["device"] == device].copy() + if sub.empty: + return None + sub.sort_values(["model", "TP"], inplace=True) + + # Hierarchical grouping: model cluster (primary) -> TP subgroup (secondary with two bars: seconds & actual) + # Enforce explicit primary model order + desired_order = ["Llama 2-7B", "Llama 2-13B", "Llama 2-70B"] + unique_models = list(pd.unique(sub["model"])) + models = [m for m in desired_order if m in unique_models] + [ + m for m in unique_models if m not in desired_order + ] + + bar_width = 0.28 # width per individual bar + tp_gap = 0.08 # gap between TP subgroups inside a model cluster + model_gap = 0.7 # gap between model clusters + + x_seconds = [] + x_actual = [] + seconds_vals = [] + actual_vals = [] + tp_midpoints = [] + tp_labels = [] + model_centers = [] + model_bounds = [] + + current_x = 0.0 + for m in models: + m_rows = sub[sub["model"] == m] + tps = sorted(m_rows["TP"].unique()) + cluster_start = current_x + for tp in tps: + row = m_rows[m_rows["TP"] == tp].iloc[0] + x_sec = current_x + x_act = current_x + bar_width + x_seconds.append(x_sec) + x_actual.append(x_act) + seconds_vals.append(row["seconds"]) + actual_vals.append(row["actual"]) + # Midpoint of the TP subgroup (between the two bars) + tp_mid = x_sec + bar_width * 0.5 + tp_midpoints.append(tp_mid) + tp_labels.append(f"TP{tp}") + current_x += 2 * bar_width + tp_gap + cluster_end = current_x - tp_gap # last subgroup end (exclude trailing tp_gap) + model_centers.append((cluster_start + cluster_end) / 2) + model_bounds.append((cluster_start, cluster_end)) + current_x += model_gap # gap after cluster + + # Dynamic figure width + fig_w = max(8.0, 0.12 * len(tp_midpoints) + 2.5) + fig, ax = plt.subplots(figsize=(fig_w, 5)) + + bars_seconds = ax.bar( + x_seconds, seconds_vals, bar_width, label="DeepFlow", color="#1f77b4" + ) + bars_actual = ax.bar( + x_actual, actual_vals, bar_width, label="Actual [1]", color="#0bbd37" + ) + + # Primary ticks: model names at cluster centers + ax.set_xticks(model_centers) + ax.set_xticklabels(models, fontsize=11) + ax.set_ylabel("Inference Latency (s)") + ax.set_title( + f"Validation of Inference Latency on Systems of {device} (network ignored)" + ) + + # Secondary TP labels ABOVE bars (centered over subgroup) + ymin, ymax = ax.get_ylim() + pad = 0.02 * (ymax - ymin) + subgroup_max = [max(s, a) for s, a in zip(seconds_vals, actual_vals)] + new_ymax = max(ymax, max(subgroup_max) + 3 * pad) + ax.set_ylim(ymin, new_ymax) + for mid, lbl, v in zip(tp_midpoints, tp_labels, subgroup_max): + ax.text(mid, v + pad, lbl, ha="center", va="bottom", fontsize=8) + + # Draw light separators between model clusters + for start, end in model_bounds: + ax.axvspan( + start - bar_width * 0.5, + end - bar_width * 0.5, + facecolor="#000000", + alpha=0.02, + ) + + # Custom legend + ax.legend() + ax.grid(axis="y", linestyle="--", alpha=0.3) + ax.margins(x=0.02, y=0.1) + + # Add citation below the graph + citation = "[1] J. Kundu et al., “Performance Modeling and Workload Analysis of Distributed Large Language Model Training and Inference,” (2024)" + plt.subplots_adjust(bottom=0.1) + fig.text(0.5, 0.015, citation, ha="center", va="bottom", fontsize=8) + + outpath = outdir / f"inf_{device}.png" + fig.savefig(outpath, dpi=200, bbox_inches="tight") + plt.close(fig) + return outpath + + +def run(enable_plot: bool = True, network_ignored: bool = True): + pct_errors = [] + data = _load_data(Path(__file__).parent / "nvidia_inf.csv") + + for idx, (device, model, tp) in enumerate(_iter_tests(data)): + seconds = _run_test(device, model, tp, idx, network_ignored) + + actual = data.loc[ + (data["device"] == device) & (data["model"] == model) & (data["TP"] == tp), + "actual", + ].values[0] + + pct_error = ( + abs(seconds - actual) / actual * 100.0 if actual != 0 else float("nan") + ) + pct_errors.append(pct_error) + + data.loc[ + (data["device"] == device) & (data["model"] == model) & (data["TP"] == tp), + "seconds", + ] = seconds + + data.loc[ + (data["device"] == device) & (data["model"] == model) & (data["TP"] == tp), + "pct_error", + ] = pct_error + + block_lines = [ + f"\n=== Result (device={device}, model={model}, TP={tp}) ===", + ] + if not math.isnan(seconds): + block_lines.append(f" DeepFlow Inference Time: {seconds:.2f}s") + block_lines.append(f" Actual Inference Time: {actual:.2f}s") + block_lines.append(f" Percent Error: {pct_error:.2f}%") + else: + block_lines.append(" DeepFlow run failed.") + print("\n".join(block_lines)) + + if enable_plot: + out_dir = os.path.join(PROJECT_ROOT, "output", "validation") + os.makedirs(out_dir, exist_ok=True) + out_dir = Path(out_dir) + + outputs = [] + for device in ["A100", "H100"]: + out = plot_device(data, device, out_dir) + if out is not None: + outputs.append(out) + if not outputs: + print("No plots generated (no matching device rows).") + else: + for p in outputs: + print(f"Saved: {p}") + avg_abs_error = sum(e for e in pct_errors if not math.isnan(e)) / len( + [e for e in pct_errors if not math.isnan(e)] + ) + # print("Average absolute percent error across all tests: {:.2f}%".format(avg_abs_error)) + return {"avg_abs_error": avg_abs_error} + + +if __name__ == "__main__": + run(enable_plot=True, network_ignored=True) diff --git a/validation_scripts/validation_helpers.py b/validation_scripts/validation_helpers.py index cf538e4b..3e8a063e 100644 --- a/validation_scripts/validation_helpers.py +++ b/validation_scripts/validation_helpers.py @@ -73,6 +73,33 @@ def _worker_init(context: _WorkerContext) -> None: _WORKER_CONTEXT = context set_astrasim_cache_mode(context.cache_mode) +def _is_list_of_dicts(value: Any) -> bool: + if not isinstance(value, list): + return False + return all(isinstance(item, dict) for item in value) + +def _merge_list_of_dicts(orig: List[Dict[str, Any]], + overrides: List[Dict[str, Any]]) -> List[Dict[str, Any]]: + orig_index = {d["id"]: d for d in orig} + orig_order = [d["id"] for d in orig] + + # Apply deep updates (or add new items) + for d in overrides: + id = d["id"] + if id in orig_index: + orig_index[id] = _deep_update(copy.deepcopy(orig_index[id]), d) + else: + orig_index[id] = copy.deepcopy(d) + + # Rebuild list preserving original order, appending new items at the end + result: List[Dict[str, Any]] = [] + for orig_id in orig_order: + result.append(orig_index[orig_id]) + for d in overrides: + id = d["id"] + if id not in orig_order: + result.append(orig_index[id]) + return result def _deep_update(target: Dict[str, Any], overrides: Mapping[str, Any]) -> Dict[str, Any]: for key, value in overrides.items(): @@ -81,6 +108,12 @@ def _deep_update(target: Dict[str, Any], overrides: Mapping[str, Any]) -> Dict[s if not isinstance(existing, dict): existing = {} target[key] = _deep_update(existing, value) + elif isinstance(value, list): + existing_list = target.get(key) + if isinstance(existing_list, list) and _is_list_of_dicts(existing_list) and _is_list_of_dicts(value): + target[key] = _merge_list_of_dicts(existing_list, value) + else: + target[key] = copy.deepcopy(value) else: target[key] = value return target @@ -320,6 +353,9 @@ def run_validation_suite( r"\[prefill\]\s*time:\s*[0-9.]+s,\s*\[decode\]\s*time:\s*([0-9.]+)s,\s*\[total\]\s*time:\s*[0-9.]+s" ) +_INF_TIME_REGEX = re.compile( + r"LLM inference time:\s*([0-9]+(?:\.[0-9]+)?)s" +) def parse_decode_time(output: str, spec: ValidationSpec) -> Dict[str, Any]: match = _DECODE_TIME_REGEX.search(output) @@ -327,10 +363,16 @@ def parse_decode_time(output: str, spec: ValidationSpec) -> Dict[str, Any]: raise ValueError(f"Failed to parse decode time for experiment '{spec.label}'.") return {"decode_time_s": float(match.group(1))} +def parse_inference_time(output: str, spec: ValidationSpec) -> Dict[str, Any]: + match = _INF_TIME_REGEX.search(output) + if not match: + raise ValueError(f"Failed to parse inference time for experiment '{spec.label}'.") + return {"inference_time_s": float(match.group(1))} __all__ = [ "ValidationSpec", "ValidationResult", "run_validation_suite", "parse_decode_time", + "parse_inference_time" ] From 6c387cd6572a0059f0e5e3e3de15a1359ca59b21 Mon Sep 17 00:00:00 2001 From: NanoMembers Date: Tue, 18 Nov 2025 13:40:28 -0800 Subject: [PATCH 43/77] Fix L2 capacity assignment in TiledGEMM class --- tile.py | 2 +- 1 file changed, 1 insertion(+), 1 deletion(-) diff --git a/tile.py b/tile.py index ff60af5c..06929fa2 100644 --- a/tile.py +++ b/tile.py @@ -308,7 +308,7 @@ def __init__(self, order_dims, tile_dims, core, memLayer, dtype_size=2): self.memLayer = memLayer self.dtype_size = dtype_size self.M, self.K, self.N = tile_dims[3] - self.capacity = memLayer[3].size_per_bundle + self.capacity = memLayer[2].size_per_bundle # L2 capacity in bytes self.order_dims = order_dims if isinstance(order_dims, str) else None # Precompute inner code once (accepts str or int code) self._inner_code = ( From 2460968edea49e7203fa4261e16bfe96bff2d9e9 Mon Sep 17 00:00:00 2001 From: NanoMembers Date: Tue, 18 Nov 2025 13:42:59 -0800 Subject: [PATCH 44/77] Add AccessBytes class and update memory access calculations in TimeCalculation --- tile.py | 139 +++++++++++++++++++++++++------------------- time_calculation.py | 15 +++-- 2 files changed, 90 insertions(+), 64 deletions(-) diff --git a/tile.py b/tile.py index 06929fa2..5fedf9b9 100644 --- a/tile.py +++ b/tile.py @@ -31,7 +31,19 @@ from math import ceil from functools import lru_cache from enum import IntEnum -from typing import Tuple, Union +from typing import Tuple, Union, NamedTuple + +class AccessBytes(NamedTuple): + reads: tuple[int, ...] + writes: tuple[int, ...] + def totals(self) -> tuple[int, ...]: + return tuple(r + w for r, w in zip(self.reads, self.writes)) + +def describe_accesses(profile: AccessBytes) -> str: + return ", ".join( + f"L{lvl}: R={formatBytes(r)} W={formatBytes(w)}" + for lvl, (r, w) in enumerate(zip(profile.reads, profile.writes)) + ) class Dataflow(IntEnum): @@ -152,7 +164,7 @@ def generate_tile_space(memLayer, num_levels: int, dim1: int, dim2: int, dim3: i Largely unchanged from original function as it is not perf-critical. """ - original = True + original = False tile_space = [] tiles = [None] * num_levels for level in range(0, num_levels - 1): @@ -214,18 +226,69 @@ def _sysarray_accesses_sig(M: int, N: int, K: int, return load_bytes + store_bytes +def _l2_accesses_sig(M: int, N: int, K: int, + l1_M: int, l1_K: int, l1_N: int, + dtype_size: int) -> tuple[int, int]: + """Compute L2 memory accesses (reads, writes) in bytes.""" + grid_x = -(-N // l1_N) + grid_y = -(-M // l1_M) + + read_bytes = grid_x * (M * K * dtype_size) + grid_y * (K * N * dtype_size) + write_bytes = M * N * dtype_size + return (read_bytes, write_bytes) + +# Simplified approximate HBM model +@lru_cache(maxsize=65536) +def _hbm_accesses_sig_simplified(M: int, N: int, K: int, + l1_M: int, l1_K: int, l1_N: int, + dtype_size: int, total_bytes: int, l2_read: int, + capacity: int) -> tuple[int, int]: + """Compute HBM/DRAM memory accesses (reads, writes) in bytes. + + Simplified reuse-distance model: + - Keep cold misses (first touches of A/B tiles). + - Approximate capacity pressure via a coarse working-set estimate per k-block. + - Avoid enumerating detailed reuse-distance patterns; replace with a ratio-based + penalty that scales cold misses when working set exceeds capacity. + """ + tile_A = int(l1_M * l1_K) + tile_B = int(l1_K * l1_N) + tile_C = int(l1_M * l1_N) + + capacity -= (108 * tile_C) + + rows = -(-M // l1_M) + cols = -(-N // l1_N) + k_blocks = -(-K // l1_K) + + total_accesses = rows * cols * k_blocks + + cold_misses = rows * k_blocks * tile_A + cols * k_blocks * tile_B + + working_set = rows * tile_A + cols * tile_B + + if working_set > capacity: + excess_ratio = min(1.0, (working_set - capacity) / max(1, working_set)) + capacity_misses = int(cold_misses * excess_ratio) + else: + capacity_misses = 0 + + total_misses = cold_misses + capacity_misses + + write_bytes = M * N * dtype_size if total_bytes > capacity * dtype_size else 0 + + per_transfer_bytes = max(1, (tile_A + tile_B)) + read_bytes = l2_read * total_misses / (total_accesses * per_transfer_bytes) + + return (int(read_bytes), int(write_bytes)) + @lru_cache(maxsize=131072) def _simulate_accesses_sig(inner_code: int, M: int, K: int, N: int, l2_M: int, l2_K: int, l2_N: int, l1_M: int, l1_K: int, l1_N: int, dtype_size: int, FMA_x: int, FMA_y: int, dataflow_code: Union[int, Dataflow], capacity: int, total_bytes: int) -> Tuple[float, int, int, int]: - """Compute memory traffic at each level (sys, L1/shared, L2, DRAM). - - Parameters are integers. The function is pure and amenable to - caching. It uses closed-form counts derived from tile reuse rather than - iterating through actual tiles, which keeps evaluation fast. - """ + """Compute memory traffic at each level (sys, L1/shared, L2, DRAM).""" sys_bytes = _sysarray_accesses_sig(M, N, K, FMA_x, FMA_y, dataflow_code, dtype_size) num_tiles_M = -(-M // l2_M) # == math.ceil(M/l2_M) @@ -240,58 +303,16 @@ def _simulate_accesses_sig(inner_code: int, M: int, K: int, N: int, write_bytes = (l1_M * l1_N * dtype_size) * (reuse_M * reuse_N) l1_shared_total = (read_bytes + write_bytes) * max_reload - # L2 accesses - # LOGIC BELOW THIS LINE CAN BE REWRITTEN IN A MORE EFFICIENT WAY - # IT HAS BEEN LEFT VERBOSE FOR READABILITY - factor_M = reuse_M - factor_K = reuse_K - factor_N = reuse_N - if inner_code == InnerLoop.M: - mk_load = max_reload - kn_load = num_tiles_K * num_tiles_N - mn_load = max_reload - factor_M = 1 - dfk = Dataflow.WST - elif inner_code == InnerLoop.K: - mk_load = max_reload - kn_load = max_reload - mn_load = num_tiles_M * num_tiles_N - factor_K = 1 - dfk = Dataflow.OST - elif inner_code == InnerLoop.N: - mk_load = num_tiles_M * num_tiles_K - kn_load = max_reload - mn_load = max_reload - factor_N = 1 - dfk = Dataflow.AST - - l2_bytes = ( - mk_load * (l2_M * l2_K * dtype_size) * factor_N - + kn_load * (l2_K * l2_N * dtype_size) * factor_M - + 2 * mn_load * (l2_M * l2_N * dtype_size) * factor_K - ) + # L2 accesses (reads+writes lumped to match return contract) + l2_read, l2_write = _l2_accesses_sig(M, N, K, l1_M, l1_K, l1_N, dtype_size) - l2f0 = M // l2_M # are we sure these are floor and should not be ceil? (ceil is -(-M // l2_M)) - l2f1 = K // l2_K - l2f2 = N // l2_N - reload_mk = 1 - reload_kn = 1 - reload_mn = 1 - if dfk == Dataflow.WST: - reload_mk = l2f2 - reload_mn = l2f1 - elif dfk == Dataflow.OST: - reload_mk = l2f2 - reload_kn = l2f0 - elif dfk == Dataflow.AST: - reload_kn = l2f0 - reload_mn = l2f1 - - dram_counts = reload_mk * (M * K * dtype_size) + reload_kn * (K * N * dtype_size) - if total_bytes > capacity: - dram_counts += reload_mn * (M * N * dtype_size) - - return (sys_bytes, l1_shared_total, l2_bytes, dram_counts) + # HBM/DRAM accesses via selected model (reads+writes lumped) + hbm_read, hbm_write = _hbm_accesses_sig_simplified(M, N, K, l1_M, l1_K, l1_N, dtype_size, total_bytes, l2_read, capacity) + + return AccessBytes( + reads=(sys_bytes, read_bytes * max_reload, l2_read, hbm_read), + writes=(0, write_bytes * max_reload, l2_write, hbm_write) + ) diff --git a/time_calculation.py b/time_calculation.py index 46db0544..4890e986 100644 --- a/time_calculation.py +++ b/time_calculation.py @@ -13,7 +13,7 @@ from hw_component import Core, MemoryHierarchy, Network, DRAM from astrasim_lib import run_cache_astrasim from model import Model_LSTM, Model_GEMM, Model_LLM -from tile import TiledGEMM, formatBytes +from tile import AccessBytes, TiledGEMM, formatBytes algByte = False # algorithmic ops false proj = False # consider projection layer, turn off for end-2-end validation, as baeline model does not have projection layer @@ -758,7 +758,7 @@ def roofline(self, flop, mem_access_, name="", info=False, mem_level=None, flash num_level = len(mem_access) try: if not flashattn_enable: - assert mem_access[num_level - 1] > 0, "last_level_mem = 0" + assert mem_access[num_level - 1] > 0, f"mem_access: {mem_access_}" except Exception as e: print( "{}: Number of accesses to the last level of memory hierarchy cannot be zero:\n {}".format( @@ -788,12 +788,14 @@ def roofline(self, flop, mem_access_, name="", info=False, mem_level=None, flash # print("Roofline: exited {}".format(name)) return max_time - def getGEMMTime(self, dim1, dim2, dim3, name, + + def getGEMMTime(self, dim1, dim2, dim3, name="", flashattn_enable=False, disable_overhead=False, read_bytes_l2=0, write_bytes_l2=0, original=False): # Streaming best selection to avoid building large dicts best_time = float("inf") best_choice = None # type: Optional[tuple] best_mem_access = None # type: Optional[tuple] + best_rw_access = None # type: Optional[AccessBytes] best_gemm = None # type: Optional[TiledGEMM] best_metric = float("inf") @@ -805,8 +807,9 @@ def getGEMMTime(self, dim1, dim2, dim3, name, print("===============================================================") GEMM_flop = gemm.GEMM_flop - mem_access = gemm.mem_accesses + rw_accesses = gemm.mem_accesses + mem_access = rw_accesses.totals() # (L0, L1, L2, DRAM) if flashattn_enable: mem_access = list(mem_access) mem_access[3] = 0 # no HBM accesses @@ -837,12 +840,14 @@ def getGEMMTime(self, dim1, dim2, dim3, name, key = (gemm._inner_code, tile_dims) # Tie-breaker metric identical to previous selection: hypot(dram, l2) + metric = math.hypot(mem_access[3], mem_access[2]) if (GEMM_time < best_time) or (GEMM_time == best_time and metric < best_metric): best_time = GEMM_time best_choice = key best_mem_access = mem_access + best_rw_access = rw_accesses best_gemm = gemm best_metric = metric @@ -861,7 +866,7 @@ def getGEMMTime(self, dim1, dim2, dim3, name, # 2 -> inner 'n' (activation stationary) best_inner_code = best_choice[0] # type: ignore[index] best_tile_dims = best_choice[1] # type: ignore[index] - return best_time, best_inner_code, best_tile_dims, mem_access + return best_time, best_inner_code, best_tile_dims, mem_access #, best_rw_access def generateTileSpace(self, dim1=None, dim2=None, dim3=None, original=False): tile_space = [] From dcdf1ece331aa7a3b7b1f619bac7950984d36899 Mon Sep 17 00:00:00 2001 From: NanoMembers Date: Tue, 18 Nov 2025 13:44:51 -0800 Subject: [PATCH 45/77] Consider SM utilization based on wave calculation method --- hw_component.py | 23 +++++++++++++++++++++++ time_calculation.py | 17 +++++++++++++---- 2 files changed, 36 insertions(+), 4 deletions(-) diff --git a/hw_component.py b/hw_component.py index 73abb13b..602a5685 100644 --- a/hw_component.py +++ b/hw_component.py @@ -60,6 +60,29 @@ def getTileDims(self): return self.getPower2TileDims() # return self.getArbitraryTileDims(self): + def calc_waves_per_sm(self, M, K, N, m, k, n): + bytes_required = self.precision * (m * k + k * n) + grid_size = math.ceil(M / m) * math.ceil(N / n) + + compute_time = 2 * (m * n * k) / 312e12 + load_time = (self.precision * (m * k + k * n)) / (7050 * 1024 ** 3) + + stages = 3 if load_time < compute_time else 5 + smem_required = bytes_required * stages + if smem_required > self.size_per_bundle: + return -1 + + reg_accum = m * n + reg_input = 16 * (m + n) + if reg_accum + reg_input > 65536: + return -1 + + smem_blocks = self.size_per_bundle // smem_required + reg_blocks = 65536 // (reg_accum + reg_input) + + return grid_size / min(smem_blocks, reg_blocks) / 108 + + def getGEMMBasedTileDims(self, M, K, N): m_dims = [ M >> i for i in range(M.bit_length()) if (M >> i) >= 1 ] k_dims = [ K >> i for i in range(K.bit_length()) if (K >> i) >= 1 ] diff --git a/time_calculation.py b/time_calculation.py index 4890e986..54a6e9fa 100644 --- a/time_calculation.py +++ b/time_calculation.py @@ -734,7 +734,7 @@ def printSysConfig(self, exp_hw_config, exp_model_config, output_file): self.memLayer[i].printStats(f) self.network.printStats(f) - def roofline(self, flop, mem_access_, name="", info=False, mem_level=None, flashattn_enable=False): + def roofline(self, flop, mem_access_, name="", util=1, info=False, mem_level=None, flashattn_enable=False): # print("Roofline: entered {}".format(name)) # Parse mem_access_ into consistent format @@ -749,6 +749,8 @@ def roofline(self, flop, mem_access_, name="", info=False, mem_level=None, flash print("mem_access_ should be integer or list, wrong input", flush=True) sys.exit(0) + throughput = self.th * util + # Determine which levels to compute if mem_level is not None: # Single level mode @@ -774,13 +776,13 @@ def roofline(self, flop, mem_access_, name="", info=False, mem_level=None, flash for level_idx, num_mem in levels_to_compute: mem_bw = self.memLayer[level_idx].getThroughput() mem_latency = self.memLayer[level_idx].getLatency() - inflection_point = float("inf") if mem_bw == 0 else self.th / mem_bw + inflection_point = float("inf") if mem_bw == 0 else throughput / mem_bw # print(f"Level {level_idx}: mem_bw={mem_bw}, mem_latency={mem_latency}, inflection_point={inflection_point}", flush=True) comp_int = float("inf") if num_mem == 0 else flop / num_mem if comp_int < inflection_point: # mem-bound level_time = (float("inf") if (mem_bw == 0 or num_mem == 0) else (num_mem / mem_bw)) + mem_latency else: # compute-bound - level_time = float("inf") if (self.th == 0) else (flop / self.th) + level_time = float("inf") if (throughput == 0) else (flop / throughput) times.append(level_time) @@ -826,7 +828,14 @@ def getGEMMTime(self, dim1, dim2, dim3, name="", if eff_sm > 0: mem_access_per_sm[1] = mem_access_per_sm[1] / eff_sm - GEMM_time = self.roofline(GEMM_flop, mem_access_per_sm, name, flashattn_enable=flashattn_enable) + sm_util = 1 + gemm_waves = self.memLayer[1].calc_waves_per_sm(dim1, dim2, dim3, gemm.l2_M, gemm.l2_K, gemm.l2_N) + if gemm_waves != -1: + full_waves = math.floor(gemm_waves) + partial_wave = gemm_waves - full_waves + sm_util = (full_waves + partial_wave * partial_wave) / gemm_waves + + GEMM_time = self.roofline(GEMM_flop, mem_access_per_sm, name, util=sm_util, flashattn_enable=flashattn_enable) if flashattn_enable or disable_overhead: pass else: From 432c9e948a1826ca765445b185f545aa7a4daaa8 Mon Sep 17 00:00:00 2001 From: NanoMembers Date: Tue, 18 Nov 2025 13:46:25 -0800 Subject: [PATCH 46/77] Adjust inference run checks in TimeCalculation --- time_calculation.py | 24 ++++++++++++++---------- 1 file changed, 14 insertions(+), 10 deletions(-) diff --git a/time_calculation.py b/time_calculation.py index 54a6e9fa..6ba2306f 100644 --- a/time_calculation.py +++ b/time_calculation.py @@ -142,8 +142,10 @@ def _analytical_all_reduce(self, size_bytes, participants, ib, ll, label): data_transfer = ((per_rank / ib) + mem_access + ll) * 2 * (participants - 1) prep_comp = per_rank prep_mem = int(math.ceil(3 * size_bytes / participants)) + + # NOTE: removed overhead because it made it impossible to eliminate network effects even with inf bandwidth and zero latency data_prep = ( - self._roofline(prep_comp, prep_mem, name=f"{label}-prep") + self.O + self._roofline(prep_comp, prep_mem, name=f"{label}-prep") # + self.O ) * (participants - 1) return data_transfer + data_prep @@ -276,15 +278,17 @@ def __init__(self, hw_config, model_config, mode, *, astra_policy_override: Opti else: self.kp1 = None self.kp2 = None - run_type = model_config.model_config.run_type - if run_type == "inference": - if self.cp > 1: - raise ValueError( - "Context parallelism (cp) is not supported for LLM inference. " - "Please set parallelism.cp to 1 for inference runs." - ) - if self.mb > 1: - print(f"[WARNING]: LLM inference configured with mb={self.mb} (>1). \n Pipeline micro-batching is ill-defined for autoregressive decode and should be avoided.") + + if self.mode != "GEMM": + run_type = model_config.model_config.run_type + if run_type == "inference": + if self.cp > 1: + raise ValueError( + "Context parallelism (cp) is not supported for LLM inference. " + "Please set parallelism.cp to 1 for inference runs." + ) + if self.mb > 1: + print(f"[WARNING]: LLM inference configured with mb={self.mb} (>1). \n Pipeline micro-batching is ill-defined for autoregressive decode and should be avoided.") # Statistics Param From afdb7315f2df300f758a4afcb1ef63b02c17f3a0 Mon Sep 17 00:00:00 2001 From: NanoMembers Date: Tue, 18 Nov 2025 15:42:44 -0800 Subject: [PATCH 47/77] Split validation tests into A100 and H100 devices and update data handling --- ...dation.py => test_IMEC_A100_no_network.py} | 13 +++++++--- tests/test_IMEC_H100_no_network.py | 25 +++++++++++++++++++ validation_scripts/imec_data/A100_inf.csv | 12 +++++++++ validation_scripts/imec_data/H100_inf.csv | 12 +++++++++ validation_scripts/nvidia_inf.csv | 23 ----------------- validation_scripts/nvidia_inf.py | 14 ++++------- 6 files changed, 63 insertions(+), 36 deletions(-) rename tests/{test_nvidia_validation.py => test_IMEC_A100_no_network.py} (53%) create mode 100644 tests/test_IMEC_H100_no_network.py create mode 100644 validation_scripts/imec_data/A100_inf.csv create mode 100644 validation_scripts/imec_data/H100_inf.csv delete mode 100644 validation_scripts/nvidia_inf.csv diff --git a/tests/test_nvidia_validation.py b/tests/test_IMEC_A100_no_network.py similarity index 53% rename from tests/test_nvidia_validation.py rename to tests/test_IMEC_A100_no_network.py index a2ef5095..699d33ab 100644 --- a/tests/test_nvidia_validation.py +++ b/tests/test_IMEC_A100_no_network.py @@ -9,9 +9,14 @@ from validation_scripts import nvidia_inf def test_avg_abs_error_below_threshold_network_ignored(): - results = nvidia_inf.run(enable_plot=True, network_ignored=True) + results = nvidia_inf.run(enable_plot=False, network_ignored=True, device="A100") avg_abs_error = results["avg_abs_error"] - assert not math.isnan(avg_abs_error), "NVIDIA validation produced no valid measurements" - # Historically the average absolute error bottoms out at roughly 5.97%, so we allow + assert not math.isnan(avg_abs_error), "IMEC validation produced no valid measurements" + # Historically the average absolute error bottoms out at roughly 6.5%, so we allow # a little slack (6.5%) to absorb minor DeepFlow changes while still flagging regressions. - assert avg_abs_error <= 17.0, f"Average absolute error {avg_abs_error:.2f}% exceeds 17.0%" + assert avg_abs_error <= 6.5, f"Average absolute error {avg_abs_error:.2f}% exceeds 6.5%" + + # TODO: split these into tests per device/model + +if __name__ == "__main__": + test_avg_abs_error_below_threshold_network_ignored() diff --git a/tests/test_IMEC_H100_no_network.py b/tests/test_IMEC_H100_no_network.py new file mode 100644 index 00000000..f6ee5b58 --- /dev/null +++ b/tests/test_IMEC_H100_no_network.py @@ -0,0 +1,25 @@ +import math +import sys +from pathlib import Path + +PROJECT_ROOT = Path(__file__).resolve().parents[1] +if str(PROJECT_ROOT) not in sys.path: + sys.path.insert(0, str(PROJECT_ROOT)) + +from validation_scripts import nvidia_inf + +# Test for H100 device configuration needs more refinement +# NOT WORKING YET + +def test_avg_abs_error_below_threshold_network_ignored(): + results = nvidia_inf.run(enable_plot=False, network_ignored=True, device="H100") + avg_abs_error = results["avg_abs_error"] + assert not math.isnan(avg_abs_error), "IMEC validation produced no valid measurements" + # Historically the average absolute error bottoms out at roughly 6.5%, so we allow + # a little slack (6.5%) to absorb minor DeepFlow changes while still flagging regressions. + assert avg_abs_error <= 6.5, f"Average absolute error {avg_abs_error:.2f}% exceeds 6.5%" + + # TODO: split these into tests per device/model + +if __name__ == "__main__": + test_avg_abs_error_below_threshold_network_ignored() diff --git a/validation_scripts/imec_data/A100_inf.csv b/validation_scripts/imec_data/A100_inf.csv new file mode 100644 index 00000000..99b20f95 --- /dev/null +++ b/validation_scripts/imec_data/A100_inf.csv @@ -0,0 +1,12 @@ +device,model,TP,actual +A100,Llama 2-7B,1,2.19 +A100,Llama 2-7B,2,1.544 +A100,Llama 2-7B,4,1.28 +A100,Llama 2-7B,8,1.187 +A100,Llama 2-13B,1,3.884 +A100,Llama 2-13B,2,2.499 +A100,Llama 2-13B,4,1.894 +A100,Llama 2-13B,8,1.693 +A100,Llama 2-70B,2,10.5 +A100,Llama 2-70B,4,6.403 +A100,Llama 2-70B,8,4.735 \ No newline at end of file diff --git a/validation_scripts/imec_data/H100_inf.csv b/validation_scripts/imec_data/H100_inf.csv new file mode 100644 index 00000000..9db0e2c0 --- /dev/null +++ b/validation_scripts/imec_data/H100_inf.csv @@ -0,0 +1,12 @@ +device,model,TP,actual +H100,Llama 2-7B,1,1.44 +H100,Llama 2-7B,2,1.143 +H100,Llama 2-7B,4,0.924 +H100,Llama 2-7B,8,0.828 +H100,Llama 2-13B,1,2.396 +H100,Llama 2-13B,2,1.717 +H100,Llama 2-13B,4,1.431 +H100,Llama 2-13B,8,1.201 +H100,Llama 2-70B,2,6.267 +H100,Llama 2-70B,4,4.1226 +H100,Llama 2-70B,8,3.202 \ No newline at end of file diff --git a/validation_scripts/nvidia_inf.csv b/validation_scripts/nvidia_inf.csv deleted file mode 100644 index 7dcf5fca..00000000 --- a/validation_scripts/nvidia_inf.csv +++ /dev/null @@ -1,23 +0,0 @@ -device,model,TP,actual,seconds -A100,Llama 2-7B,1,2.19,2.36 -A100,Llama 2-7B,2,1.544,1.82 -A100,Llama 2-7B,4,1.28,1.77 -A100,Llama 2-7B,8,1.187,2.16 -A100,Llama 2-13B,1,3.884,3.84 -A100,Llama 2-13B,2,2.499,2.75 -A100,Llama 2-13B,4,1.894,2.46 -A100,Llama 2-13B,8,1.693,2.85 -A100,Llama 2-70B,2,10.5,9.55 -A100,Llama 2-70B,4,6.403,6.98 -A100,Llama 2-70B,8,4.735,6.74 -H100,Llama 2-7B,1,1.44,1.82 -H100,Llama 2-7B,2,1.143,1.55 -H100,Llama 2-7B,4,0.924,1.62 -H100,Llama 2-7B,8,0.828,2.06 -H100,Llama 2-13B,1,2.396,2.75 -H100,Llama 2-13B,2,1.717,2.2 -H100,Llama 2-13B,4,1.431,2.17 -H100,Llama 2-13B,8,1.201,2.67 -H100,Llama 2-70B,2,6.267,6.77 -H100,Llama 2-70B,4,4.1226,5.56 -H100,Llama 2-70B,8,3.202,6 diff --git a/validation_scripts/nvidia_inf.py b/validation_scripts/nvidia_inf.py index 9983da22..0a8dc201 100644 --- a/validation_scripts/nvidia_inf.py +++ b/validation_scripts/nvidia_inf.py @@ -35,8 +35,7 @@ def _load_data(csv_path: Path) -> pd.DataFrame: # Normalize dtypes df["TP"] = df["TP"].astype(int) df["actual"] = pd.to_numeric(df["actual"], errors="coerce") - df["seconds"] = pd.to_numeric(df["seconds"], errors="coerce") - return df.dropna(subset=["device", "model", "TP", "actual", "seconds"]) + return df.dropna(subset=["device", "model", "TP", "actual"]) def _iter_tests(df: pd.DataFrame): @@ -75,13 +74,10 @@ def _run_test( "dimensions": [ { "id": "dim0", - "topology": {"bandwidth": "100000 GB", "latency": "0"}, + "topology": {"bandwidth": "100000 GB", "latency": "1e-9" }, } ] }, - "execution_backend": { - "model": "analytical", - }, } hw_overrides.update(override_bw) @@ -221,9 +217,9 @@ def plot_device(df: pd.DataFrame, device: str, outdir: Path) -> Path | None: return outpath -def run(enable_plot: bool = True, network_ignored: bool = True): +def run(enable_plot: bool = True, network_ignored: bool = True, device: str = "A100"): pct_errors = [] - data = _load_data(Path(__file__).parent / "nvidia_inf.csv") + data = _load_data(Path(__file__).parent / f"imec_data/{device}_inf.csv") for idx, (device, model, tp) in enumerate(_iter_tests(data)): seconds = _run_test(device, model, tp, idx, network_ignored) @@ -277,7 +273,7 @@ def run(enable_plot: bool = True, network_ignored: bool = True): avg_abs_error = sum(e for e in pct_errors if not math.isnan(e)) / len( [e for e in pct_errors if not math.isnan(e)] ) - # print("Average absolute percent error across all tests: {:.2f}%".format(avg_abs_error)) + print("Average absolute percent error across all tests: {:.2f}%".format(avg_abs_error)) return {"avg_abs_error": avg_abs_error} From 6e6578751768ded7b4ea5e3fe8d01dcc01793d0d Mon Sep 17 00:00:00 2001 From: NanoMembers Date: Tue, 18 Nov 2025 15:50:44 -0800 Subject: [PATCH 48/77] Mark IMEC H100 test for failure --- tests/test_IMEC_H100_no_network.py | 2 ++ 1 file changed, 2 insertions(+) diff --git a/tests/test_IMEC_H100_no_network.py b/tests/test_IMEC_H100_no_network.py index f6ee5b58..cd3b0571 100644 --- a/tests/test_IMEC_H100_no_network.py +++ b/tests/test_IMEC_H100_no_network.py @@ -1,6 +1,7 @@ import math import sys from pathlib import Path +import pytest PROJECT_ROOT = Path(__file__).resolve().parents[1] if str(PROJECT_ROOT) not in sys.path: @@ -11,6 +12,7 @@ # Test for H100 device configuration needs more refinement # NOT WORKING YET +@pytest.mark.xfail(reason="Pending model refinements", strict=False) def test_avg_abs_error_below_threshold_network_ignored(): results = nvidia_inf.run(enable_plot=False, network_ignored=True, device="H100") avg_abs_error = results["avg_abs_error"] From 7ce961aadd5d8b0f9c3ab1bd5f6836a21480714c Mon Sep 17 00:00:00 2001 From: gkarfakis19 Date: Fri, 21 Nov 2025 13:09:05 -0800 Subject: [PATCH 49/77] multithreaded and cleaner per device-model tests --- tests/conftest.py | 10 + tests/test_IMEC_A100_no_network.py | 22 -- tests/test_IMEC_H100_no_network.py | 27 --- tests/test_IMEC_inf_no_network.py | 73 +++++++ tests/test_koyeb_validation.py | 3 +- validation_scripts/nvidia_inf.py | 252 +++++++++++++++++------ validation_scripts/validation_helpers.py | 36 +++- 7 files changed, 311 insertions(+), 112 deletions(-) create mode 100644 tests/conftest.py delete mode 100644 tests/test_IMEC_A100_no_network.py delete mode 100644 tests/test_IMEC_H100_no_network.py create mode 100644 tests/test_IMEC_inf_no_network.py diff --git a/tests/conftest.py b/tests/conftest.py new file mode 100644 index 00000000..0607f997 --- /dev/null +++ b/tests/conftest.py @@ -0,0 +1,10 @@ +import sys +from pathlib import Path + +# This is magic file that pytest loads. We need it unfortunately to hook the validation_helpers module into pytest. + +PROJECT_ROOT = Path(__file__).resolve().parents[1] +if str(PROJECT_ROOT) not in sys.path: + sys.path.insert(0, str(PROJECT_ROOT)) + +pytest_plugins = ["validation_scripts.validation_helpers"] diff --git a/tests/test_IMEC_A100_no_network.py b/tests/test_IMEC_A100_no_network.py deleted file mode 100644 index 699d33ab..00000000 --- a/tests/test_IMEC_A100_no_network.py +++ /dev/null @@ -1,22 +0,0 @@ -import math -import sys -from pathlib import Path - -PROJECT_ROOT = Path(__file__).resolve().parents[1] -if str(PROJECT_ROOT) not in sys.path: - sys.path.insert(0, str(PROJECT_ROOT)) - -from validation_scripts import nvidia_inf - -def test_avg_abs_error_below_threshold_network_ignored(): - results = nvidia_inf.run(enable_plot=False, network_ignored=True, device="A100") - avg_abs_error = results["avg_abs_error"] - assert not math.isnan(avg_abs_error), "IMEC validation produced no valid measurements" - # Historically the average absolute error bottoms out at roughly 6.5%, so we allow - # a little slack (6.5%) to absorb minor DeepFlow changes while still flagging regressions. - assert avg_abs_error <= 6.5, f"Average absolute error {avg_abs_error:.2f}% exceeds 6.5%" - - # TODO: split these into tests per device/model - -if __name__ == "__main__": - test_avg_abs_error_below_threshold_network_ignored() diff --git a/tests/test_IMEC_H100_no_network.py b/tests/test_IMEC_H100_no_network.py deleted file mode 100644 index cd3b0571..00000000 --- a/tests/test_IMEC_H100_no_network.py +++ /dev/null @@ -1,27 +0,0 @@ -import math -import sys -from pathlib import Path -import pytest - -PROJECT_ROOT = Path(__file__).resolve().parents[1] -if str(PROJECT_ROOT) not in sys.path: - sys.path.insert(0, str(PROJECT_ROOT)) - -from validation_scripts import nvidia_inf - -# Test for H100 device configuration needs more refinement -# NOT WORKING YET - -@pytest.mark.xfail(reason="Pending model refinements", strict=False) -def test_avg_abs_error_below_threshold_network_ignored(): - results = nvidia_inf.run(enable_plot=False, network_ignored=True, device="H100") - avg_abs_error = results["avg_abs_error"] - assert not math.isnan(avg_abs_error), "IMEC validation produced no valid measurements" - # Historically the average absolute error bottoms out at roughly 6.5%, so we allow - # a little slack (6.5%) to absorb minor DeepFlow changes while still flagging regressions. - assert avg_abs_error <= 6.5, f"Average absolute error {avg_abs_error:.2f}% exceeds 6.5%" - - # TODO: split these into tests per device/model - -if __name__ == "__main__": - test_avg_abs_error_below_threshold_network_ignored() diff --git a/tests/test_IMEC_inf_no_network.py b/tests/test_IMEC_inf_no_network.py new file mode 100644 index 00000000..caf3e7d4 --- /dev/null +++ b/tests/test_IMEC_inf_no_network.py @@ -0,0 +1,73 @@ +import math +from pathlib import Path + +import pandas as pd +import pytest + +from validation_scripts import nvidia_inf + +PROJECT_ROOT = Path(__file__).resolve().parents[1] +DATA_DIR = PROJECT_ROOT / "validation_scripts/imec_data" +DEFAULT_THRESHOLD = 100.0 +# Override per (device, model) if needed. Falls back to DEFAULT_THRESHOLD. +# Actual + 0.5% margin of error, rounded up to nearest 0.5%. +GROUP_THRESHOLDS = { + ("A100", "Llama 2-7B"): 5.5, # 4.62 + 0.5 = 5.12 -> 5.5 + ("A100", "Llama 2-13B"): 3.0, # 2.44 + 0.5 = 2.94 -> 3.0 + ("A100", "Llama 2-70B"): 13.0, # 12.09 + 0.5 = 12.59 -> 13.0 +} + + +def _load_params(): + params = [] + for device in ("A100", "H100"): + csv_path = DATA_DIR / f"{device}_inf.csv" + if not csv_path.exists(): + continue + df = pd.read_csv(csv_path, comment="/") + df["TP"] = df["TP"].astype(int) + df["actual"] = pd.to_numeric(df["actual"], errors="coerce") + df = df.dropna(subset=["device", "model", "TP", "actual"]) + for model, group in df.groupby("model"): + tp_actuals = [(int(tp), float(act)) for tp, act in zip(group["TP"], group["actual"])] + marks = () + if device == "H100": + marks = (pytest.mark.xfail(reason="Pending model refinements", strict=False),) + params.append( + pytest.param( + device, + model, + tp_actuals, + marks=marks, + id=f"{device}-{model}", + ) + ) + return params + + +PARAMS = _load_params() + + +@pytest.mark.parametrize("device,model,tp_actuals", PARAMS) +def test_avg_abs_error_below_threshold_network_ignored(device, model, tp_actuals, record_validation): + label = f"{device} | {model}" + threshold = GROUP_THRESHOLDS.get((device, model), DEFAULT_THRESHOLD) + result = nvidia_inf.run( + enable_plot=False, + network_ignored=True, + device=device, + models=[model], + emit_logs=False, + ) + rows = result["rows"] + assert rows, f"{label} produced no validation rows" + + pct_errors = [] + for row in rows: + pct_error = row["pct_error"] + assert not math.isnan(pct_error), f"{label}-TP{row['tp']} produced no valid measurements" + pct_errors.append(pct_error) + val_label = f"{label} avg_abs_error" + avg_error = sum(pct_errors) / len(pct_errors) + record_validation(val_label, avg_error, expected_pct=threshold) + assert avg_error <= threshold, f"{label} avg error {avg_error:.2f}% exceeds {threshold}%" diff --git a/tests/test_koyeb_validation.py b/tests/test_koyeb_validation.py index e9c526a0..02d1a8fd 100644 --- a/tests/test_koyeb_validation.py +++ b/tests/test_koyeb_validation.py @@ -9,9 +9,10 @@ from validation_scripts import koyeb -def test_koyeb_average_absolute_error_below_threshold(): +def test_koyeb_average_absolute_error_below_threshold(record_validation): results = koyeb.run_all_and_plot(enable_plot=False, verbose=False) avg_abs_error = results["avg_abs_error"] + record_validation("avg_abs_error", avg_abs_error, expected_pct=6.5) assert not math.isnan(avg_abs_error), "Koyeb validation produced no valid measurements" # Historically the average absolute error bottoms out at roughly 5.97%, so we allow # a little slack (6.5%) to absorb minor DeepFlow changes while still flagging regressions. diff --git a/validation_scripts/nvidia_inf.py b/validation_scripts/nvidia_inf.py index 0a8dc201..06a83e3f 100644 --- a/validation_scripts/nvidia_inf.py +++ b/validation_scripts/nvidia_inf.py @@ -1,16 +1,26 @@ import argparse import math import os +from functools import lru_cache from pathlib import Path -import pandas as pd -from typing import List +from typing import Dict, Iterable, List, Optional, Sequence, Tuple +import sys import matplotlib.pyplot as plt +import pandas as pd -from .validation_helpers import ( - ValidationSpec, - run_validation_suite, - parse_inference_time, -) +try: + from .validation_helpers import ( + ValidationSpec, + run_validation_suite, + parse_inference_time, + ) +except ImportError: + sys.path.append(str(Path(__file__).resolve().parents[1])) + from validation_scripts.validation_helpers import ( # type: ignore + ValidationSpec, + run_validation_suite, + parse_inference_time, + ) HW_CONFIGS = { "A100": "a100_80GB.yaml", @@ -27,6 +37,7 @@ RUN_PERF = os.path.join(PROJECT_ROOT, "run_perf.py") HARDWARE_CONFIG_PATH = os.path.join(PROJECT_ROOT, "configs/hardware-config") MODEL_CONFIG_PATH = os.path.join(PROJECT_ROOT, "configs/model-config") +DATA_DIR = Path(__file__).parent / "imec_data" def _load_data(csv_path: Path) -> pd.DataFrame: @@ -38,16 +49,20 @@ def _load_data(csv_path: Path) -> pd.DataFrame: return df.dropna(subset=["device", "model", "TP", "actual"]) +@lru_cache(maxsize=None) +def _load_device_data(device: str) -> pd.DataFrame: + csv_path = DATA_DIR / f"{device}_inf.csv" + if not csv_path.exists(): + raise FileNotFoundError(f"Validation CSV not found for device {device}: {csv_path}") + return _load_data(csv_path) + + def _iter_tests(df: pd.DataFrame): for _, row in df.iterrows(): yield (row["device"], row["model"], row["TP"]) -def _run_test( - device: str, model: str, tp: int, idx: int, network_ignored: bool = True -) -> float: - print(f"Running test {idx + 1}: Device={device}, Model={model}, TP={tp}") - specs: List[ValidationSpec] = [] +def _build_spec(device: str, model: str, tp: int, idx: int, network_ignored: bool) -> Tuple[ValidationSpec, str, str]: label = f"{device} {model} TP={tp}" model_overrides = { @@ -58,7 +73,7 @@ def _run_test( } } - hw_overrides = { + hw_overrides: Dict[str, Dict[str, object]] = { "parallelism": { "dp": 1, "tp": int(tp), @@ -81,20 +96,43 @@ def _run_test( } hw_overrides.update(override_bw) - specs.append( - ValidationSpec( - label=label, - model_overrides=model_overrides, - hardware_overrides=hw_overrides, - order=idx, - ) + spec = ValidationSpec( + label=label, + model_overrides=model_overrides, + hardware_overrides=hw_overrides, + model_config_path=os.path.join(MODEL_CONFIG_PATH, MODEL_CONFIGS.get(model)), + hardware_config_path=os.path.join(HARDWARE_CONFIG_PATH, HW_CONFIGS.get(device)), + metadata={"device": device, "model": model, "tp": int(tp)}, + order=idx, ) + hw_config_path = spec.hardware_config_path # type: ignore[arg-type] + model_config_path = spec.model_config_path # type: ignore[arg-type] + return spec, hw_config_path, model_config_path - hw_config_path = os.path.join(HARDWARE_CONFIG_PATH, HW_CONFIGS.get(device)) - model_config_path = os.path.join(MODEL_CONFIG_PATH, MODEL_CONFIGS.get(model)) +def _lookup_actual(device: str, model: str, tp: int) -> float: + data = _load_device_data(device) + matches = data.loc[ + (data["device"] == device) & (data["model"] == model) & (data["TP"] == int(tp)), + "actual", + ] + if matches.empty: + raise ValueError(f"No reference inference time for {device} {model} TP={tp}") + return float(matches.values[0]) + + +def run_single( + device: str, + model: str, + tp: int, + *, + network_ignored: bool = True, + actual_inference_time_s: Optional[float] = None, + emit_logs: bool = False, +) -> Dict[str, object]: + spec, hw_config_path, model_config_path = _build_spec(device, model, tp, idx=0, network_ignored=network_ignored) validation_results = run_validation_suite( - specs, + [spec], base_model_config_path=model_config_path, base_hardware_config_path=hw_config_path, result_parser=parse_inference_time, @@ -102,16 +140,97 @@ def _run_test( ) result = validation_results[0] - if result.success: - inf_time_s = float(result.metrics.get("inference_time_s", float("nan"))) + inference_time_s = float(result.metrics.get("inference_time_s", float("nan"))) if result.success else float("nan") + err_detail = None if result.success else (result.error or "DeepFlow run failed") + + if actual_inference_time_s is None: + try: + actual_inference_time_s = _lookup_actual(device, model, tp) + except Exception: + actual_inference_time_s = float("nan") + + if math.isnan(inference_time_s) or actual_inference_time_s == 0 or math.isnan(actual_inference_time_s): + pct_error = float("nan") else: - inf_time_s = float("nan") - err_detail = result.error or "DeepFlow run failed" - print(f" ERROR: {err_detail}") - if result.raw_output: - print(result.raw_output.strip()) + pct_error = abs(inference_time_s - actual_inference_time_s) / actual_inference_time_s * 100.0 + + if emit_logs: + block_lines = [ + f"\n=== Result (device={device}, model={model}, TP={tp}) ===", + ] + if not math.isnan(inference_time_s): + block_lines.append(f" DeepFlow Inference Time: {inference_time_s:.2f}s") + block_lines.append(f" Actual Inference Time: {actual_inference_time_s:.2f}s") + block_lines.append(f" Percent Error: {pct_error:.2f}%") + else: + block_lines.append(f" DeepFlow run failed. {err_detail or ''}".rstrip()) + if result.raw_output: + block_lines.append(result.raw_output.strip()) + print("\n".join(block_lines)) + + return { + "success": result.success, + "inference_time_s": inference_time_s, + "actual_inference_time_s": actual_inference_time_s, + "pct_error": pct_error, + "error": err_detail, + "raw_output": result.raw_output, + } - return inf_time_s + +def build_specs_for_device( + device: str, + *, + network_ignored: bool = True, + models: Optional[Iterable[str]] = None, +) -> Tuple[List[ValidationSpec], Dict[Tuple[str, int], float], str, str]: + data = _load_device_data(device) + model_filter = set(models) if models is not None else None + specs: List[ValidationSpec] = [] + actual_lookup: Dict[Tuple[str, int], float] = {} + base_model_path: Optional[str] = None + hw_config_path: Optional[str] = None + idx = 0 + for device_val, model, tp in _iter_tests(data): + if model_filter and model not in model_filter: + continue + spec, hw_path, model_path = _build_spec(device_val, model, tp, idx, network_ignored) + specs.append(spec) + actual_lookup[(model, int(tp))] = _lookup_actual(device_val, model, int(tp)) + base_model_path = base_model_path or model_path + hw_config_path = hw_config_path or hw_path + idx += 1 + + if not specs: + raise ValueError(f"No validation specs generated for device={device} (models={models}).") + return specs, actual_lookup, base_model_path, hw_config_path + + +def compute_pct_errors(results, actual_lookup: Dict[Tuple[str, int], float]): + rows: List[Dict[str, object]] = [] + for res in results: + metadata = res.spec.metadata or {} + model = metadata.get("model") + tp = int(metadata.get("tp")) if "tp" in metadata else None + inf_time = float(res.metrics.get("inference_time_s", float("nan"))) if res.success else float("nan") + actual = actual_lookup.get((model, tp)) if model is not None and tp is not None else float("nan") + if math.isnan(inf_time) or actual is None or actual == 0 or math.isnan(actual): + pct_error = float("nan") + else: + pct_error = abs(inf_time - actual) / actual * 100.0 + rows.append( + { + "device": metadata.get("device"), + "model": model, + "tp": tp, + "inference_time_s": inf_time, + "actual_inference_time_s": actual, + "pct_error": pct_error, + "success": res.success, + "error": res.error, + } + ) + return rows def plot_device(df: pd.DataFrame, device: str, outdir: Path) -> Path | None: @@ -217,43 +336,54 @@ def plot_device(df: pd.DataFrame, device: str, outdir: Path) -> Path | None: return outpath -def run(enable_plot: bool = True, network_ignored: bool = True, device: str = "A100"): +def run( + enable_plot: bool = True, + network_ignored: bool = True, + device: str = "A100", + models: Optional[Sequence[str]] = None, + emit_logs: bool = True, +): pct_errors = [] - data = _load_data(Path(__file__).parent / f"imec_data/{device}_inf.csv") + data = _load_device_data(device) - for idx, (device, model, tp) in enumerate(_iter_tests(data)): - seconds = _run_test(device, model, tp, idx, network_ignored) + specs, actual_lookup, base_model_path, hw_config_path = build_specs_for_device( + device, network_ignored=network_ignored, models=models + ) - actual = data.loc[ - (data["device"] == device) & (data["model"] == model) & (data["TP"] == tp), - "actual", - ].values[0] + validation_results = run_validation_suite( + specs, + base_model_config_path=base_model_path, + base_hardware_config_path=hw_config_path, + result_parser=parse_inference_time, + run_perf_path=RUN_PERF, + ) - pct_error = ( - abs(seconds - actual) / actual * 100.0 if actual != 0 else float("nan") - ) - pct_errors.append(pct_error) + rows = compute_pct_errors(validation_results, actual_lookup) + for row in rows: + pct_error = float(row["pct_error"]) + pct_errors.append(pct_error) data.loc[ - (data["device"] == device) & (data["model"] == model) & (data["TP"] == tp), + (data["device"] == row["device"]) & (data["model"] == row["model"]) & (data["TP"] == row["tp"]), "seconds", - ] = seconds + ] = row["inference_time_s"] data.loc[ - (data["device"] == device) & (data["model"] == model) & (data["TP"] == tp), + (data["device"] == row["device"]) & (data["model"] == row["model"]) & (data["TP"] == row["tp"]), "pct_error", ] = pct_error - block_lines = [ - f"\n=== Result (device={device}, model={model}, TP={tp}) ===", - ] - if not math.isnan(seconds): - block_lines.append(f" DeepFlow Inference Time: {seconds:.2f}s") - block_lines.append(f" Actual Inference Time: {actual:.2f}s") - block_lines.append(f" Percent Error: {pct_error:.2f}%") - else: - block_lines.append(" DeepFlow run failed.") - print("\n".join(block_lines)) + if emit_logs: + block_lines = [ + f"\n=== Result (device={row['device']}, model={row['model']}, TP={row['tp']}) ===", + ] + if not math.isnan(row["inference_time_s"]): + block_lines.append(f" DeepFlow Inference Time: {float(row['inference_time_s']):.2f}s") + block_lines.append(f" Actual Inference Time: {float(row['actual_inference_time_s']):.2f}s") + block_lines.append(f" Percent Error: {pct_error:.2f}%") + else: + block_lines.append(f" DeepFlow run failed. {(row.get('error') or '')}".rstrip()) + print("\n".join(block_lines)) if enable_plot: out_dir = os.path.join(PROJECT_ROOT, "output", "validation") @@ -270,11 +400,11 @@ def run(enable_plot: bool = True, network_ignored: bool = True, device: str = "A else: for p in outputs: print(f"Saved: {p}") - avg_abs_error = sum(e for e in pct_errors if not math.isnan(e)) / len( - [e for e in pct_errors if not math.isnan(e)] - ) - print("Average absolute percent error across all tests: {:.2f}%".format(avg_abs_error)) - return {"avg_abs_error": avg_abs_error} + valid_errors = [e for e in pct_errors if not math.isnan(e)] + avg_abs_error = sum(valid_errors) / len(valid_errors) if valid_errors else float("nan") + if emit_logs: + print("Average absolute percent error across all tests: {:.2f}%".format(avg_abs_error)) + return {"avg_abs_error": avg_abs_error, "rows": rows} if __name__ == "__main__": diff --git a/validation_scripts/validation_helpers.py b/validation_scripts/validation_helpers.py index 3e8a063e..6c4cdb0d 100644 --- a/validation_scripts/validation_helpers.py +++ b/validation_scripts/validation_helpers.py @@ -12,6 +12,7 @@ from dataclasses import dataclass, field, replace from typing import Any, Callable, Dict, List, Mapping, Optional, Sequence, Tuple +import pytest import yaml PROJECT_ROOT = os.path.dirname(os.path.dirname(os.path.abspath(__file__))) @@ -369,10 +370,43 @@ def parse_inference_time(output: str, spec: ValidationSpec) -> Dict[str, Any]: raise ValueError(f"Failed to parse inference time for experiment '{spec.label}'.") return {"inference_time_s": float(match.group(1))} +# ---- Pytest integration for quiet validation reporting ---- +_VALIDATION_METRICS: List[Tuple[str, str, float, Optional[float]]] = [] + +# magic function, needed for pytest to work +def pytest_configure(config): + _VALIDATION_METRICS.clear() + + +@pytest.fixture +def record_validation(request): + def _record(label: str, achieved_pct: float, expected_pct: Optional[float] = None) -> None: + #node id is tests/test_IMEC_A100_no_network.py::test_avg_abs_error_below_threshold_network_ignored + # I want just the filename under tests/, not the function name + filename = request.node.nodeid.split("::")[0] + filename = filename.split("/")[-1] + _VALIDATION_METRICS.append( + (filename, label, achieved_pct, expected_pct) + ) + return _record + + +def pytest_terminal_summary(terminalreporter, exitstatus): + if not _VALIDATION_METRICS: + return + tr = terminalreporter + tr.write_sep("=", "Test results") + for nodeid, label, achieved, expected in _VALIDATION_METRICS: + line = f"{nodeid}: {label} achieved {achieved:.2f}%" + if expected is not None: + line += f" (expected <= {expected:.2f}%)" + tr.write_line(line) + __all__ = [ "ValidationSpec", "ValidationResult", "run_validation_suite", "parse_decode_time", - "parse_inference_time" + "parse_inference_time", + "record_validation", ] From 0796c9b201980b7b2bea95ebe4331a7259c480f8 Mon Sep 17 00:00:00 2001 From: gkarfakis19 Date: Fri, 21 Nov 2025 13:21:29 -0800 Subject: [PATCH 50/77] doubled num tests --- tests/test_IMEC_inf_no_network.py | 92 ++++++++++++++++++++++--------- 1 file changed, 67 insertions(+), 25 deletions(-) diff --git a/tests/test_IMEC_inf_no_network.py b/tests/test_IMEC_inf_no_network.py index caf3e7d4..7462dfe2 100644 --- a/tests/test_IMEC_inf_no_network.py +++ b/tests/test_IMEC_inf_no_network.py @@ -8,17 +8,26 @@ PROJECT_ROOT = Path(__file__).resolve().parents[1] DATA_DIR = PROJECT_ROOT / "validation_scripts/imec_data" -DEFAULT_THRESHOLD = 100.0 -# Override per (device, model) if needed. Falls back to DEFAULT_THRESHOLD. -# Actual + 0.5% margin of error, rounded up to nearest 0.5%. -GROUP_THRESHOLDS = { - ("A100", "Llama 2-7B"): 5.5, # 4.62 + 0.5 = 5.12 -> 5.5 - ("A100", "Llama 2-13B"): 3.0, # 2.44 + 0.5 = 2.94 -> 3.0 - ("A100", "Llama 2-70B"): 13.0, # 12.09 + 0.5 = 12.59 -> 13.0 + + +DEFAULT_THRESHOLD = 10.0 +# Override per (device, model) average pct error. +# Error is actual + 0.5% margin of error, rounded up to nearest 0.5%. +MODEL_THRESHOLDS = { + ("A100", "Llama 2-7B"): 5.5, # 4.62 + 0.5 -> ceil to 5.5 + ("A100", "Llama 2-13B"): 3.0, # 2.44 + 0.5 -> ceil to 3.0 + ("A100", "Llama 2-70B"): 13.0, # 12.09 + 0.5 -> ceil to 13.0 +} +# Override per (device, TP) average pct error across all models at that TP. +TP_THRESHOLDS = { + ("A100", 1): 4.5, # 3.98 + 0.5 -> ceil to 4.5 + ("A100", 2): 8.5, # 7.60 + 0.5 -> ceil to 8.5 + ("A100", 4): 5.5, # 4.91 + 0.5 -> ceil to 5.5 + ("A100", 8): 7.0, # 6.34 + 0.5 -> ceil to 7.0 } -def _load_params(): +def _load_params(by: str): params = [] for device in ("A100", "H100"): csv_path = DATA_DIR / f"{device}_inf.csv" @@ -28,38 +37,51 @@ def _load_params(): df["TP"] = df["TP"].astype(int) df["actual"] = pd.to_numeric(df["actual"], errors="coerce") df = df.dropna(subset=["device", "model", "TP", "actual"]) - for model, group in df.groupby("model"): - tp_actuals = [(int(tp), float(act)) for tp, act in zip(group["TP"], group["actual"])] + if by == "model": + grouped = df.groupby("model") + else: + grouped = df.groupby("TP") + for key, group in grouped: marks = () if device == "H100": marks = (pytest.mark.xfail(reason="Pending model refinements", strict=False),) params.append( pytest.param( device, - model, - tp_actuals, + key, marks=marks, - id=f"{device}-{model}", + id=f"{device}-{by}-{key}", ) ) return params -PARAMS = _load_params() +MODEL_PARAMS = _load_params("model") +TP_PARAMS = _load_params("TP") -@pytest.mark.parametrize("device,model,tp_actuals", PARAMS) -def test_avg_abs_error_below_threshold_network_ignored(device, model, tp_actuals, record_validation): +@pytest.fixture(scope="module") +def cached_rows(): + cache = {} + for device in ("A100", "H100"): + try: + result = nvidia_inf.run( + enable_plot=False, + network_ignored=True, + device=device, + emit_logs=False, + ) + except FileNotFoundError: + continue + cache[device] = result["rows"] + return cache + + +@pytest.mark.parametrize("device,model", MODEL_PARAMS) +def test_avg_abs_error_by_model(device, model, cached_rows, record_validation): label = f"{device} | {model}" - threshold = GROUP_THRESHOLDS.get((device, model), DEFAULT_THRESHOLD) - result = nvidia_inf.run( - enable_plot=False, - network_ignored=True, - device=device, - models=[model], - emit_logs=False, - ) - rows = result["rows"] + threshold = MODEL_THRESHOLDS.get((device, model), DEFAULT_THRESHOLD) + rows = [r for r in cached_rows.get(device, []) if r["model"] == model] assert rows, f"{label} produced no validation rows" pct_errors = [] @@ -67,7 +89,27 @@ def test_avg_abs_error_below_threshold_network_ignored(device, model, tp_actuals pct_error = row["pct_error"] assert not math.isnan(pct_error), f"{label}-TP{row['tp']} produced no valid measurements" pct_errors.append(pct_error) + + avg_error = sum(pct_errors) / len(pct_errors) val_label = f"{label} avg_abs_error" + record_validation(val_label, avg_error, expected_pct=threshold) + assert avg_error <= threshold, f"{label} avg error {avg_error:.2f}% exceeds {threshold}%" + + +@pytest.mark.parametrize("device,tp", TP_PARAMS) +def test_avg_abs_error_by_tp(device, tp, cached_rows, record_validation): + label = f"{device} | TP{tp}" + threshold = TP_THRESHOLDS.get((device, int(tp)), DEFAULT_THRESHOLD) + rows = [r for r in cached_rows.get(device, []) if int(r["tp"]) == int(tp)] + assert rows, f"{label} produced no validation rows" + + pct_errors = [] + for row in rows: + pct_error = row["pct_error"] + assert not math.isnan(pct_error), f"{label}-{row['model']} produced no valid measurements" + pct_errors.append(pct_error) + avg_error = sum(pct_errors) / len(pct_errors) + val_label = f"{label} avg_abs_error" record_validation(val_label, avg_error, expected_pct=threshold) assert avg_error <= threshold, f"{label} avg error {avg_error:.2f}% exceeds {threshold}%" From 5fbb468fb8472fc81c694aa33c28ea9aa5d38e44 Mon Sep 17 00:00:00 2001 From: gkarfakis19 Date: Fri, 21 Nov 2025 13:28:35 -0800 Subject: [PATCH 51/77] separated val configs --- validation_scripts/koyeb.py | 5 +- validation_scripts/nvidia_inf.py | 5 +- .../hardware-config}/H100_SXM5_80GB.yaml | 119 +++++++++--------- .../hardware-config}/a100_80GB.yaml | 103 +++++++-------- .../a100_80GB_no_parallelism.yaml | 119 ++++++++---------- .../model-config/Llama2-13B_inf.yaml | 21 ++++ .../model-config/Llama2-70B_inf.yaml | 21 ++++ .../model-config/Llama2-7B_inf.yaml | 21 ++++ .../model-config/Llama3.1-8B_inf.yaml | 21 ++++ 9 files changed, 254 insertions(+), 181 deletions(-) rename {configs/old_lstm_hw_configs => validation_scripts/validation_configs/hardware-config}/H100_SXM5_80GB.yaml (63%) rename {configs/old_lstm_hw_configs => validation_scripts/validation_configs/hardware-config}/a100_80GB.yaml (66%) rename configs/old_lstm_hw_configs/v100.yaml => validation_scripts/validation_configs/hardware-config/a100_80GB_no_parallelism.yaml (57%) create mode 100644 validation_scripts/validation_configs/model-config/Llama2-13B_inf.yaml create mode 100644 validation_scripts/validation_configs/model-config/Llama2-70B_inf.yaml create mode 100644 validation_scripts/validation_configs/model-config/Llama2-7B_inf.yaml create mode 100644 validation_scripts/validation_configs/model-config/Llama3.1-8B_inf.yaml diff --git a/validation_scripts/koyeb.py b/validation_scripts/koyeb.py index b883fdb7..d1f3a000 100644 --- a/validation_scripts/koyeb.py +++ b/validation_scripts/koyeb.py @@ -71,8 +71,9 @@ PROJECT_ROOT = os.path.dirname(os.path.dirname(os.path.abspath(__file__))) RUN_PERF = os.path.join(PROJECT_ROOT, "run_perf.py") -HARDWARE_CONFIG = os.path.join(PROJECT_ROOT, "configs/hardware-config/a100_80GB_no_parallelism.yaml") -BASE_MODEL_CONFIG = os.path.join(PROJECT_ROOT, "configs/model-config/Llama3.1-8B_inf.yaml") +VALIDATION_CONFIG_ROOT = os.path.join(PROJECT_ROOT, "validation_scripts", "validation_configs") +HARDWARE_CONFIG = os.path.join(VALIDATION_CONFIG_ROOT, "hardware-config/a100_80GB_no_parallelism.yaml") +BASE_MODEL_CONFIG = os.path.join(VALIDATION_CONFIG_ROOT, "model-config/Llama3.1-8B_inf.yaml") def _collect_points_in_order(data: Dict[str, Dict[Tuple[int, int, int], Tuple[float, float]]]) -> List[Tuple[str, Tuple[int, int, int], Tuple[float, float]]]: diff --git a/validation_scripts/nvidia_inf.py b/validation_scripts/nvidia_inf.py index 06a83e3f..df7c45b6 100644 --- a/validation_scripts/nvidia_inf.py +++ b/validation_scripts/nvidia_inf.py @@ -35,8 +35,9 @@ PROJECT_ROOT = os.path.dirname(os.path.dirname(os.path.abspath(__file__))) RUN_PERF = os.path.join(PROJECT_ROOT, "run_perf.py") -HARDWARE_CONFIG_PATH = os.path.join(PROJECT_ROOT, "configs/hardware-config") -MODEL_CONFIG_PATH = os.path.join(PROJECT_ROOT, "configs/model-config") +VALIDATION_CONFIG_ROOT = os.path.join(PROJECT_ROOT, "validation_scripts", "validation_configs") +HARDWARE_CONFIG_PATH = os.path.join(VALIDATION_CONFIG_ROOT, "hardware-config") +MODEL_CONFIG_PATH = os.path.join(VALIDATION_CONFIG_ROOT, "model-config") DATA_DIR = Path(__file__).parent / "imec_data" diff --git a/configs/old_lstm_hw_configs/H100_SXM5_80GB.yaml b/validation_scripts/validation_configs/hardware-config/H100_SXM5_80GB.yaml similarity index 63% rename from configs/old_lstm_hw_configs/H100_SXM5_80GB.yaml rename to validation_scripts/validation_configs/hardware-config/H100_SXM5_80GB.yaml index 183a9f33..be3fbdfc 100644 --- a/configs/old_lstm_hw_configs/H100_SXM5_80GB.yaml +++ b/validation_scripts/validation_configs/hardware-config/H100_SXM5_80GB.yaml @@ -8,7 +8,7 @@ sw_param: kernel_launch_overhead: 9e-6 # (s) # dp_zero_stage selects ZeRO-style data-parallel sharding: 0=DDP, 1=ZeRO-1, 2=ZeRO-2, 3=ZeRO-3 # only used if data parallelism (dp) > 1 - dp_zero_stage: 0 + dp_zero_stage: 1 precision: # tensor_format accepts numeric byte counts or dtype strings (fp16, bf16, fp32). # Take care that the number of flops defined below has been carefully tuned for FP16/BF16 formats. @@ -19,7 +19,7 @@ sw_param: # -> false forces all activations/weights/gradients/optimizer states/stats to use tensor_format bytes. # -> true enables AMP-style FP32 accumulations/optimizer states depending on param_storage_mode. # if true, optimizer states and stats are always FP32. - mixed_precision: false + mixed_precision: true # param_storage_mode: as_tensor_format | tensor_plus_fp32_master | fp32_params. # -> as_tensor_format makes weights/gradients use tensor_format bytes. # -> tensor_plus_fp32_master makes weights/gradients use tensor_format bytes and @@ -28,8 +28,7 @@ sw_param: param_storage_mode: as_tensor_format # kv_cache precision may be set as bytes, dtype string, or "as_tensor_format". kv_cache: as_tensor_format - h2d_bandwidth: 13297674240 # Bytes/s; set -1 to disable H2D timing - + h2d_bandwidth: -1 # Bytes/s; set -1 to disable H2D timing tech_param: core: operating_frequency: 1.98e9 # (Hz) GPU boost clock frequency [TePU] @@ -37,13 +36,13 @@ tech_param: nominal_flop_rate_per_mcu: 1024 # F16 FLOP / cycle / SM is 4096 [Semi] nominal_energy_per_flop: 1.8e-13 num_mcu_per_bundle: 4 # number of tensorcores per SM [TePU] - FMA_d1: 8 # tensor core can compute the results for an d1 x d2 x d1 mixed-precision matrix multiplication per clock - FMA_d2: 4 - dataflow: best #{wst. ast, ost, best, none}: wst: weight stationary, ast: activation stationary, ost: output stationary - util: 0.75 #util should be 0.75 (12/16) for tensorcore + FMA_d1: 16 # tensor core can compute the results for an d1 x d2 x d1 mixed-precision matrix multiplication per clock + FMA_d2: 8 + dataflow: best + util: 1.0 DRAM: size: 80 GB # [TePU] - bandwidth: 3440 GB # (/s) [TePU] + bandwidth: 3430 GB # (/s) [TePU] dynamic_energy_per_bit: 4.5e-12 # (J) [https://passlab.github.io/mchpc/mchpc2019/presentations/MCHPC_Pawlowski_keynote.pdf, slide 36] latency: 100e-9 # (s) util: 1 @@ -61,33 +60,10 @@ tech_param: util: 1 SRAM-R: size: 33 MB # register file memory per SM * number of SMs [NVIDIA] - bandwidth: 244 TB # traffic for 2 reads and 1 write for GEMMs computed per clock + bandwidth: 122 TB # traffic for 2 reads and 1 write for GEMMs computed per clock dynamic_energy_per_bit: 110e-15 latency: 0 util: 1 - network: - intra_node: - latency: 5e-6 - util: 1 - bandwidth: 400 Gb - inter_node: - latency: 5e-6 - util: 0.96 - bandwidth: 400 Gb - -system_hierarchy: - num_devices_per_node: 1 - num_nodes: 1 - kp1_inter: false - kp2_inter: false - dp_inter: false - lp_inter: false - tp_inter: false - -network_topology: - # Allowed values: ring | fc | switch - inter_node: ring - intra_node: ring memory_hierarchy: l0: #Register Memory @@ -103,48 +79,71 @@ memory_hierarchy: type: DRAM scope: global -scheduling_param: +parallelism: auto: false ### LLM PARAMS ### dp: 1 # data parallelism. Acts as replica count during inference + tp: 8 # tensor parallelism + tp_sp: false # tensor-sequence parallelism + cp: 1 # context parallelism. Not supported for inference. lp: 1 # pipeline parallelism. Number of pipeline stages - mb: 1 # micro-batches (for best results set to >= lp) - tp: 1 # tensor parallelism - tp_sp: true # tensor-sequence parallelism - cp : 1 # context parallelism - ### LSTM/GEMM LEGACY PARAMS ### - t: CR - kp1: 1 - kp2: 1 - kp_hidden_dim1: 1 - kp_hidden_dim2: 1 - kp_softmax_dim1: 1 - kp_softmax_dim2: 1 - kp_embedding_dim1: 1 - kp_embedding_dim2: 1 - kp_projection_dim1: 1 - kp_projection_dim2: 1 - kp_hidden_type: -1 - kp_softmax_type: -1 - kp_embedding_type: -1 - kp_projection_type: -1 + # mb is not a true parallelism parameter, does not affect num of GPUs. + mb: 1 # micro-batches (for best results set to >= lp). + +network: + dimensions: + - id: dim0 + label: torus2d_tp + size: auto + topology: + type: Ring + bandwidth: 100 GB # We assume 8 25GB/s NVLINKs in the torus. Torus2D is hardcoded bidirectional, so it's actually 4 links, 25GB/s each (because they get replicated). So 100GB/s total. + latency: 5e-6 + energy_per_bit: 8e-12 # (J/bit) + util: 1.0 + optimize_2dmap: false + collective_override: {} # overrides execution_backend.collectives when provided + parallelisms: [tp, cp] + - id: dim1 + label: inter_torus_infiniband + size: auto + topology: + type: Ring + bandwidth: 25 GB + latency: 5e-6 + energy_per_bit: 8e-12 # (J/bit) + util: 1.0 + collective_override: {} # overrides execution_backend.collectives when provided + parallelisms: [lp,dp] + - id: dim2 + label: inter_torus_infiniband_dp + size: auto + topology: + type: Ring + bandwidth: 25 GB + latency: 5e-6 + energy_per_bit: 8e-12 # (J/bit) + util: 1.0 + parallelisms: [] + collective_override: {} # overrides execution_backend.collectives when provided execution_backend: - model: astra # analytical | astra + model: analytical # analytical | astra astra: backend: astra - mode: hybrid # hybrid | full_astrasim_hierarchical | full_astrasim_flattened - collective_override: - # Allowed values per op: auto | ring | direct | halvingDoubling | doubleBinaryTree + mode: full_astrasim_flattened # hybrid | full_astrasim_hierarchical | full_astrasim_flattened + collectives: + # Default collective algorithms. Per-dimension collective_override entries can override these. + # Allowed values per op: auto | ring | direct | halvingDoubling | doubleBinaryTree | mesh | mesh2d | hypercube | torus2d # Note: leave to auto unless you have a specific reason to change it. all_gather: auto all_reduce: auto reduce_scatter: auto all_to_all: auto sys_options: - endpoint_delay: 10 + endpoint_delay: 5 active_chunks_per_dimension: 32 # Higher value means significantly slower astrasim runtime, but more accurate especially for tensor parallelism - preferred_dataset_splits: 32 # Should be same as active_chunks_per_dimension + preferred_dataset_splits: 4 # Should be same as active_chunks_per_dimension inference: kvcache_type: hbm_only diff --git a/configs/old_lstm_hw_configs/a100_80GB.yaml b/validation_scripts/validation_configs/hardware-config/a100_80GB.yaml similarity index 66% rename from configs/old_lstm_hw_configs/a100_80GB.yaml rename to validation_scripts/validation_configs/hardware-config/a100_80GB.yaml index cfbebc6c..7e27d5a5 100644 --- a/configs/old_lstm_hw_configs/a100_80GB.yaml +++ b/validation_scripts/validation_configs/hardware-config/a100_80GB.yaml @@ -2,7 +2,7 @@ sw_param: kernel_launch_overhead: 9e-6 # (s) # dp_zero_stage selects ZeRO-style data-parallel sharding: 0=DDP, 1=ZeRO-1, 2=ZeRO-2, 3=ZeRO-3 # only used if data parallelism (dp) > 1 - dp_zero_stage: 2 + dp_zero_stage: 1 precision: # tensor_format accepts numeric byte counts or dtype strings (fp16, bf16, fp32). # Take care that the number of flops defined below has been carefully tuned for FP16/BF16 formats. @@ -13,7 +13,7 @@ sw_param: # -> false forces all activations/weights/gradients/optimizer states/stats to use tensor_format bytes. # -> true enables AMP-style FP32 accumulations/optimizer states depending on param_storage_mode. # if true, optimizer states and stats are always FP32. - mixed_precision: false + mixed_precision: true # param_storage_mode: as_tensor_format | tensor_plus_fp32_master | fp32_params. # -> as_tensor_format makes weights/gradients use tensor_format bytes. # -> tensor_plus_fp32_master makes weights/gradients use tensor_format bytes and @@ -33,10 +33,10 @@ tech_param: FMA_d1: 8 # tensor core can compute the results for an d1 x d2 x d1 mixed-precision matrix multiplication per clock FMA_d2: 4 dataflow: best # {wst | ast | ost | best}: wst: weight stationary, ast: activation stationary, ost: output stationary, best: best of all 3 - util: 100.0 + util: 1.0 DRAM: size: 80 GB - bandwidth: 198600 GB # (/s) + bandwidth: 1986 GB # (/s) dynamic_energy_per_bit: 4.5e-12 # (J) latency: 100e-9 # (s) util: 1 @@ -58,30 +58,7 @@ tech_param: dynamic_energy_per_bit: 110e-15 # (J) latency: 0 util: 1 - network: - intra_node: - latency: 5e-6 - util: 1 - bandwidth: 200 Gb - inter_node: - latency: 5e-6 - util: 0.96 - bandwidth: 200 Gb -system_hierarchy: - num_devices_per_node: 8 # currently only support astrasim integration with all devices in one node. - num_nodes: 1 - ## LLM PARAMS ## - dp_inter: false - lp_inter: false - tp_inter: false - cp_inter: false - ## LEGACY LSTM PARAMS ## - kp1_inter: false - kp2_inter: false -network_topology: - # Allowed values: ring | fc | switch | HyperCube | Mesh | Mesh2D | Torus2D - inter_node: ring - intra_node: ring + memory_hierarchy: l0: #Register Memory type: SRAM-R @@ -95,39 +72,63 @@ memory_hierarchy: l3: #Global Memory type: DRAM scope: global -scheduling_param: - auto: false # Not supported for LLM mode. - ### LLM PARAMS ### + +parallelism: dp: 2 # data parallelism. Acts as replica count during inference - lp: 2 # pipeline parallelism. Number of pipeline stages - mb: 2 # micro-batches (for best results set to >= lp) tp: 2 # tensor parallelism tp_sp: true # tensor-sequence parallelism - cp : 1 # context parallelism - ### LSTM/GEMM LEGACY PARAMS ### - t: "CR" - kp1: 1 - kp2: 1 - kp_hidden_dim1: 1 - kp_hidden_dim2: 1 - kp_softmax_dim1: 1 - kp_softmax_dim2: 1 - kp_embedding_dim1: 1 - kp_embedding_dim2: 1 - kp_projection_dim1: 1 - kp_projection_dim2: 1 - kp_hidden_type: -1 - kp_softmax_type: -1 - kp_embedding_type: -1 - kp_projection_type: -1 + cp: 2 # context parallelism. Not supported for inference. + lp: 4 # pipeline parallelism. Number of pipeline stages + # mb is not a true parallelism parameter, does not affect num of GPUs. + mb: 4 # micro-batches (for best results set to >= lp). +network: + # Global faulty links apply to the active AstraSim dimension: [[src, dst, weight]] + # faulty_links: [[0,1,0.6],[16,17,0.1]] + dimensions: + # currently only support astrasim integration with all devices in one node. + - id: dim0 + label: torus2d_tp + size: auto + topology: + type: Ring + bandwidth: 100 GB # We assume 8 25GB/s NVLINKs in the torus. Torus2D is hardcoded bidirectional, so it's actually 4 links, 25GB/s each (because they get replicated). So 100GB/s total. + latency: 5e-6 + energy_per_bit: 8e-12 # (J/bit) + util: 1.0 + optimize_2dmap: false + collective_override: {} # overrides execution_backend.collectives when provided + parallelisms: [tp, cp] + - id: dim1 + label: inter_torus_infiniband + size: auto + topology: + type: Ring + bandwidth: 25 GB + latency: 5e-6 + energy_per_bit: 8e-12 # (J/bit) + util: 1.0 + collective_override: {} # overrides execution_backend.collectives when provided + parallelisms: [lp,dp] + - id: dim2 + label: inter_torus_infiniband_dp + size: auto + topology: + type: Ring + bandwidth: 25 GB + latency: 5e-6 + energy_per_bit: 8e-12 # (J/bit) + util: 1.0 + parallelisms: [] + collective_override: {} # overrides execution_backend.collectives when provided execution_backend: model: astra # analytical | astra astra: backend: astra mode: full_astrasim_flattened # hybrid | full_astrasim_hierarchical | full_astrasim_flattened - collective_override: + collectives: + # Default collective algorithms. Per-dimension collective_override entries can override these. # Allowed values per op: auto | ring | direct | halvingDoubling | doubleBinaryTree | mesh | mesh2d | hypercube | torus2d # Note: leave to auto unless you have a specific reason to change it. all_gather: auto diff --git a/configs/old_lstm_hw_configs/v100.yaml b/validation_scripts/validation_configs/hardware-config/a100_80GB_no_parallelism.yaml similarity index 57% rename from configs/old_lstm_hw_configs/v100.yaml rename to validation_scripts/validation_configs/hardware-config/a100_80GB_no_parallelism.yaml index b930572b..798101a2 100644 --- a/configs/old_lstm_hw_configs/v100.yaml +++ b/validation_scripts/validation_configs/hardware-config/a100_80GB_no_parallelism.yaml @@ -1,9 +1,11 @@ +# EXAMPLE FILE USED FOR TESTING/README ONLY. +# DO NOT MODIFY. sw_param: - kernel_launch_overhead: 9e-6 #9us + kernel_launch_overhead: 9e-6 # dp_zero_stage selects ZeRO-style data-parallel sharding: 0=DDP, 1=ZeRO-1, 2=ZeRO-2, 3=ZeRO-3 # only used if data parallelism (dp) > 1 - dp_zero_stage: 0 + dp_zero_stage: 2 precision: # tensor_format accepts numeric byte counts or dtype strings (fp16, bf16, fp32). # Take care that the number of flops defined below has been carefully tuned for FP16/BF16 formats. @@ -23,40 +25,39 @@ sw_param: param_storage_mode: as_tensor_format # kv_cache precision may be set as bytes, dtype string, or "as_tensor_format". kv_cache: as_tensor_format - h2d_bandwidth: 13297674240 # Bytes/s; set -1 to disable H2D timing - + h2d_bandwidth: -1 # Bytes/s; set -1 to disable H2D timing tech_param: core: - operating_frequency: 1.33e9 #Hz - num_bundles: 80 # V100 has 80 SMs - nominal_flop_rate_per_mcu: 128 + operating_frequency: 1.41e9 + num_bundles: 108 + nominal_flop_rate_per_mcu: 512 nominal_energy_per_flop: 1.8e-13 - num_mcu_per_bundle: 8 #a bundle can be thought of as an SM, number of tensorcores per SM - FMA_d1: 4 # tensor core can compute the results for an d1 x d2 x d1 mixed-precision matrix multiplication per clock + num_mcu_per_bundle: 4 + FMA_d1: 8 FMA_d2: 4 - dataflow: best #{wst. ast, ost, best, none}: wst: weight stationary, ast: activation stationary, ost: output stationary - util: 0.85 #util should be 0.75 (12/16) for tensorcore + dataflow: best + util: 1.0 DRAM: - size: 16 GB # V100 has 16GB HBM2 - bandwidth: 900 GB # V100 has ~900 GB/s bandwidth - dynamic_energy_per_bit: 4.2e-12 #Joules + size: 80 GB + bandwidth: 1986 GB + dynamic_energy_per_bit: 4.5e-12 latency: 100e-9 util: 1 SRAM-L2: - size: 6 MB # V100 L2 cache - bandwidth: 3072 GB # Estimated L2 bandwidth + size: 40 MB + bandwidth: 7050 GB dynamic_energy_per_bit: 130e-15 latency: 0 - util: 0.45 + util: 1 SRAM-L1: - size: 6400 KB # 80 KB per SM * 80 SMs - bandwidth: 15 TB # Estimated L1 bandwidth + size: 20736 KB + bandwidth: 18 TB dynamic_energy_per_bit: 130e-15 latency: 0 util: 1 SRAM-R: - size: 20 MB # Register file - bandwidth: 80 TB # Estimated register bandwidth + size: 27 MB + bandwidth: 122 TB dynamic_energy_per_bit: 110e-15 latency: 0 util: 1 @@ -64,73 +65,59 @@ tech_param: intra_node: latency: 5e-6 util: 1 - bandwidth: 300 Gb + bandwidth: 400 Gb inter_node: latency: 5e-6 util: 0.96 - bandwidth: 300 Gb - -system_hierarchy: - num_devices_per_node: 1 - num_nodes: 1 - kp1_inter: false - kp2_inter: false - dp_inter: false - lp_inter: false - tp_inter: false - -network_topology: - # Allowed values: ring | fc | switch - inter_node: ring - intra_node: ring - + bandwidth: 400 Gb memory_hierarchy: - l0: #Register Memory + l0: type: SRAM-R scope: mcu - l1: #Shared Memory + l1: type: SRAM-L1 scope: mcu-bundle - l2: #L2 + l2: type: SRAM-L2 scope: global - l3: #Global Memory + l3: type: DRAM scope: global -scheduling_param: - auto: false + +parallelism: ### LLM PARAMS ### dp: 1 # acts as replica count during inference lp: 1 # pipeline stages mb: 1 # number of micro-batches tp: 1 # tensor parallelism - tp_sp: true # tensor-sequence parallelism + tp_sp: false # tensor-sequence parallelism cp : 1 # context parallelism - ### LSTM/GEMM LEGACY PARAMS ### - t: CR - kp1: 1 - kp2: 1 - kp_hidden_dim1: 1 - kp_hidden_dim2: 1 - kp_softmax_dim1: 1 - kp_softmax_dim2: 1 - kp_embedding_dim1: 1 - kp_embedding_dim2: 1 - kp_projection_dim1: 1 - kp_projection_dim2: 1 - kp_hidden_type: -1 - kp_softmax_type: -1 - kp_embedding_type: -1 - kp_projection_type: -1 + +network: + dimensions: + - id: dim0 + label: device_fabric + size: 1 + topology: + type: Ring + bandwidth: 400 Gb + latency: 5e-6 + energy_per_bit: 8e-12 # (J/bit) + util: 1.0 + collective_override: {} # overrides execution_backend.collectives when provided + parallelisms: [tp, cp, lp, dp] + + execution_backend: - model: astra # analytical | astra + model: analytical # analytical | astra astra: backend: astra - mode: hybrid # hybrid | full_astrasim_hierarchical | full_astrasim_flattened - collective_override: - # Allowed values per op: auto | ring | direct | halvingDoubling | doubleBinaryTree + mode: full_astrasim_flattened # hybrid | full_astrasim_hierarchical | full_astrasim_flattened + collectives: + # Default collective algorithms. Per-dimension collective_override entries can override these. + # Allowed values per op: auto | ring | direct | halvingDoubling | doubleBinaryTree | mesh | mesh2d | hypercube | torus2d # Note: leave to auto unless you have a specific reason to change it. all_gather: auto all_reduce: auto @@ -142,4 +129,4 @@ execution_backend: preferred_dataset_splits: 32 # Should be same as active_chunks_per_dimension inference: - kvcache_type: hbm_only + kvcache_type: hbm_only # Only supported mode for now. diff --git a/validation_scripts/validation_configs/model-config/Llama2-13B_inf.yaml b/validation_scripts/validation_configs/model-config/Llama2-13B_inf.yaml new file mode 100644 index 00000000..a5105fa0 --- /dev/null +++ b/validation_scripts/validation_configs/model-config/Llama2-13B_inf.yaml @@ -0,0 +1,21 @@ +model_param: + mode: "LLM" # Mode of operation + run_type: "inference" # Type of run: "training" or "inference" + tied_embeddings: true # Share input and output embedding weights (WIP: affects memory estimation only) + model_type: "llama" # gpt assumes dense-style architecture with GELU; llama assumes llama-style architecture with SwiGLU + global_batch_size: 16 # Sequences per global batch + gradient_accumulation_steps: 1 # Number of gradient accumulation micro-steps per optimizer update + seq_len: 2048 # Total context length (prefill + decode_len for inference) + decode_len: 1024 # Number of tokens to decode during inference + hidden_dim: 5120 # Transformer hidden dimension + attention: # Attention-related parameters + attention_type: "mha" # Supported attention types: "mha" or "gqa" ("mla" not implemented) + num_heads: 40 # Total number of attention heads + kv_heads: 8 # Number of key/value heads (used only when attention_type = "gqa") + use_flashattention: true # Enable FlashAttention optimization + attention_tile_size: 256 # Tile size used in attention computation + num_experts: 1 # Number of experts in MoE layers. Set to 1 to disable MoE. + top_k: 1 # Number of active experts per token in MoE + intermediate_size: 13824 # Output dimension of the first FFN layer (typically 4×hidden_dim) + vocab_size: 32000 # Vocabulary size + num_layers: 40 # Number of transformer layers diff --git a/validation_scripts/validation_configs/model-config/Llama2-70B_inf.yaml b/validation_scripts/validation_configs/model-config/Llama2-70B_inf.yaml new file mode 100644 index 00000000..24facac3 --- /dev/null +++ b/validation_scripts/validation_configs/model-config/Llama2-70B_inf.yaml @@ -0,0 +1,21 @@ +model_param: + mode: "LLM" # Mode of operation + run_type: "inference" # Type of run: "training" or "inference" + tied_embeddings: true # Share input and output embedding weights (WIP: affects memory estimation only) + model_type: "llama" # gpt assumes dense-style architecture with GELU; llama assumes llama-style architecture with SwiGLU + global_batch_size: 16 # Sequences per global batch + gradient_accumulation_steps: 1 # Number of gradient accumulation micro-steps per optimizer update + seq_len: 2048 # Total context length (prefill + decode_len for inference) + decode_len: 1024 # Number of tokens to decode during inference + hidden_dim: 8192 # Transformer hidden dimension + attention: # Attention-related parameters + attention_type: "gqa" # Supported attention types: "mha" or "gqa" ("mla" not implemented) + num_heads: 64 # Total number of attention heads + kv_heads: 8 # Number of key/value heads (used only when attention_type = "gqa") + use_flashattention: true # Enable FlashAttention optimization + attention_tile_size: 256 # Tile size used in attention computation + num_experts: 1 # Number of experts in MoE layers. Set to 1 to disable MoE. + top_k: 1 # Number of active experts per token in MoE + intermediate_size: 28672 # Output dimension of the first FFN layer (typically 4×hidden_dim) + vocab_size: 32000 # Vocabulary size + num_layers: 80 # Number of transformer layers diff --git a/validation_scripts/validation_configs/model-config/Llama2-7B_inf.yaml b/validation_scripts/validation_configs/model-config/Llama2-7B_inf.yaml new file mode 100644 index 00000000..0bed4a63 --- /dev/null +++ b/validation_scripts/validation_configs/model-config/Llama2-7B_inf.yaml @@ -0,0 +1,21 @@ +model_param: + mode: "LLM" # Mode of operation + run_type: "inference" # Type of run: "training" or "inference" + tied_embeddings: true # Share input and output embedding weights (WIP: affects memory estimation only) + model_type: "llama" # gpt assumes dense-style architecture with GELU; llama assumes llama-style architecture with SwiGLU + global_batch_size: 16 # Sequences per global batch + gradient_accumulation_steps: 1 # Number of gradient accumulation micro-steps per optimizer update + seq_len: 2048 # Total context length (prefill + decode_len for inference) + decode_len: 1024 # Number of tokens to decode during inference + hidden_dim: 4096 # Transformer hidden dimension + attention: # Attention-related parameters + attention_type: "mha" # Supported attention types: "mha" or "gqa" ("mla" not implemented) + num_heads: 32 # Total number of attention heads + kv_heads: 8 # Number of key/value heads (used only when attention_type = "gqa") + use_flashattention: true # Enable FlashAttention optimization + attention_tile_size: 256 # Tile size used in attention computation + num_experts: 1 # Number of experts in MoE layers. Set to 1 to disable MoE. + top_k: 1 # Number of active experts per token in MoE + intermediate_size: 11008 # Output dimension of the first FFN layer (typically 4×hidden_dim) + vocab_size: 32000 # Vocabulary size + num_layers: 32 # Number of transformer layers diff --git a/validation_scripts/validation_configs/model-config/Llama3.1-8B_inf.yaml b/validation_scripts/validation_configs/model-config/Llama3.1-8B_inf.yaml new file mode 100644 index 00000000..439048b9 --- /dev/null +++ b/validation_scripts/validation_configs/model-config/Llama3.1-8B_inf.yaml @@ -0,0 +1,21 @@ +model_param: + mode: "LLM" # Mode of operation + run_type: "inference" # Type of run: "training" or "inference" + tied_embeddings: true # Share input and output embedding weights (WIP: affects memory estimation only) + model_type: "llama" # gpt assumes dense-style architecture with GELU; llama assumes llama-style architecture with SwiGLU + global_batch_size: 1 # Sequences per global batch + gradient_accumulation_steps: 1 # Number of gradient accumulation micro-steps per optimizer update + seq_len: 2048 # Total context length (prefill + decode_len for inference) + decode_len: 1024 # Number of tokens to decode during inference + hidden_dim: 4096 # Transformer hidden dimension + attention: # Attention-related parameters + attention_type: "gqa" # Supported attention types: "mha" or "gqa" ("mla" not implemented) + num_heads: 32 # Total number of attention heads + kv_heads: 8 # Number of key/value heads (used only when attention_type = "gqa") + use_flashattention: true # Enable FlashAttention optimization + attention_tile_size: 128 # Tile size used in attention computation + num_experts: 1 # Number of experts in MoE layers. Set to 1 to disable MoE. + top_k: 1 # Number of active experts per token in MoE + intermediate_size: 14336 # Output dimension of the first FFN layer (typically 4×hidden_dim) + vocab_size: 128256 # Vocabulary size + num_layers: 32 # Number of transformer layers From 0e88891f5e47d27565bb5e9798c31dec8585084e Mon Sep 17 00:00:00 2001 From: gkarfakis19 Date: Fri, 21 Nov 2025 15:32:48 -0800 Subject: [PATCH 52/77] WIP train test --- ...dation.py => test_koyeb_inf_single_gpu.py} | 0 validation_scripts/imec_data/A100_train.csv | 3 + validation_scripts/nvidia_inf.py | 3 + validation_scripts/nvidia_train.py | 211 ++++++++++++++++++ 4 files changed, 217 insertions(+) rename tests/{test_koyeb_validation.py => test_koyeb_inf_single_gpu.py} (100%) create mode 100644 validation_scripts/imec_data/A100_train.csv create mode 100644 validation_scripts/nvidia_train.py diff --git a/tests/test_koyeb_validation.py b/tests/test_koyeb_inf_single_gpu.py similarity index 100% rename from tests/test_koyeb_validation.py rename to tests/test_koyeb_inf_single_gpu.py diff --git a/validation_scripts/imec_data/A100_train.csv b/validation_scripts/imec_data/A100_train.csv new file mode 100644 index 00000000..db137c48 --- /dev/null +++ b/validation_scripts/imec_data/A100_train.csv @@ -0,0 +1,3 @@ +// placeholder training validation data +device,model,batch,dp,tp,pp,cp,actual +A100,Llama 2-7B,32,2,2,1,1,123.45 diff --git a/validation_scripts/nvidia_inf.py b/validation_scripts/nvidia_inf.py index df7c45b6..1d8ff2a7 100644 --- a/validation_scripts/nvidia_inf.py +++ b/validation_scripts/nvidia_inf.py @@ -22,6 +22,9 @@ parse_inference_time, ) +# Data used to be in https://docs.nvidia.com/nemo-framework/user-guide/latest/llms/llama/performance.html in 2022. +# no longer available now, found via paper. "Performance Modeling and Workload Analysis of Distributed Large Language Model Training and Inference" by J. Kundu et al. + HW_CONFIGS = { "A100": "a100_80GB.yaml", "H100": "H100_SXM5_80GB.yaml", diff --git a/validation_scripts/nvidia_train.py b/validation_scripts/nvidia_train.py new file mode 100644 index 00000000..33876204 --- /dev/null +++ b/validation_scripts/nvidia_train.py @@ -0,0 +1,211 @@ +import math +import os +from functools import lru_cache +from pathlib import Path +from typing import Dict, Iterable, List, Optional, Sequence, Tuple + +from .validation_helpers import ( + ValidationSpec, + run_validation_suite, + parse_inference_time, +) + +PROJECT_ROOT = os.path.dirname(os.path.dirname(os.path.abspath(__file__))) +RUN_PERF = os.path.join(PROJECT_ROOT, "run_perf.py") +VALIDATION_CONFIG_ROOT = os.path.join(PROJECT_ROOT, "validation_scripts", "validation_configs") +HARDWARE_CONFIG_PATH = os.path.join(VALIDATION_CONFIG_ROOT, "hardware-config") +MODEL_CONFIG_PATH = os.path.join(VALIDATION_CONFIG_ROOT, "model-config") +DATA_DIR = Path(__file__).parent / "imec_data" + +HW_CONFIGS = { + "A100": "a100_80GB.yaml", +} + +MODEL_CONFIGS = { + "Llama 2-7B": "Llama2-7B_inf.yaml", +} + + +def _load_data(csv_path: Path): + from validation_scripts.nvidia_inf import _load_data as _load_data_inf # reuse parsing behavior + return _load_data_inf(csv_path) + + +@lru_cache(maxsize=None) +def _load_device_data(device: str): + csv_path = DATA_DIR / f"{device}_train.csv" + if not csv_path.exists(): + raise FileNotFoundError(f"Validation CSV not found for device {device}: {csv_path}") + return _load_data(csv_path) + + +def _iter_tests(df): + for _, row in df.iterrows(): + yield (row["device"], row["model"], row["batch"], row["dp"], row["tp"], row["pp"], row["cp"]) + + +def _build_spec(device: str, model: str, batch: int, dp: int, tp: int, pp: int, cp: int, idx: int) -> Tuple[ValidationSpec, str, str]: + label = f"{device} {model} bs={batch} dp={dp} tp={tp} pp={pp} cp={cp}" + + model_overrides = { + "model_param": { + "global_batch_size": int(batch), + } + } + + hw_overrides = { + "parallelism": { + "dp": int(dp), + "tp": int(tp), + "tp_sp": False, + "cp": int(cp), + "lp": int(pp), + "mb": max(1, int(pp)), # micro-batches align with pipeline by default + } + } + + spec = ValidationSpec( + label=label, + model_overrides=model_overrides, + hardware_overrides=hw_overrides, + model_config_path=os.path.join(MODEL_CONFIG_PATH, MODEL_CONFIGS.get(model)), + hardware_config_path=os.path.join(HARDWARE_CONFIG_PATH, HW_CONFIGS.get(device)), + metadata={ + "device": device, + "model": model, + "batch": int(batch), + "dp": int(dp), + "tp": int(tp), + "pp": int(pp), + "cp": int(cp), + }, + order=idx, + ) + return spec, spec.model_config_path, spec.hardware_config_path # type: ignore[return-value] + + +def _lookup_actual(device: str, model: str, batch: int, dp: int, tp: int, pp: int, cp: int) -> float: + data = _load_device_data(device) + matches = data.loc[ + (data["device"] == device) + & (data["model"] == model) + & (data["batch"] == int(batch)) + & (data["dp"] == int(dp)) + & (data["tp"] == int(tp)) + & (data["pp"] == int(pp)) + & (data["cp"] == int(cp)), + "actual", + ] + if matches.empty: + raise ValueError(f"No reference training time for {device} {model} bs={batch} dp={dp} tp={tp} pp={pp} cp={cp}") + return float(matches.values[0]) + + +def compute_pct_errors(results, actual_lookup: Dict[Tuple[str, int, int, int, int], float]): + rows: List[Dict[str, object]] = [] + for res in results: + meta = res.spec.metadata or {} + model = meta.get("model") + batch = int(meta.get("batch")) if "batch" in meta else None + dp = int(meta.get("dp")) if "dp" in meta else None + tp = int(meta.get("tp")) if "tp" in meta else None + pp = int(meta.get("pp")) if "pp" in meta else None + cp = int(meta.get("cp")) if "cp" in meta else None + key = (model, batch, dp, tp, pp, cp) + inf_time = float(res.metrics.get("inference_time_s", float("nan"))) if res.success else float("nan") + actual = actual_lookup.get(key, float("nan")) + if math.isnan(inf_time) or actual == 0 or math.isnan(actual): + pct_error = float("nan") + else: + pct_error = abs(inf_time - actual) / actual * 100.0 + rows.append( + { + "device": meta.get("device"), + "model": model, + "batch": batch, + "dp": dp, + "tp": tp, + "pp": pp, + "cp": cp, + "inference_time_s": inf_time, + "actual_inference_time_s": actual, + "pct_error": pct_error, + "success": res.success, + "error": res.error, + } + ) + return rows + + +def build_specs_for_device( + device: str, + *, + models: Optional[Iterable[str]] = None, +) -> Tuple[List[ValidationSpec], Dict[Tuple[str, int, int, int, int, int], float], str, str]: + data = _load_device_data(device) + model_filter = set(models) if models is not None else None + specs: List[ValidationSpec] = [] + actual_lookup: Dict[Tuple[str, int, int, int, int, int], float] = {} + base_model_path: Optional[str] = None + hw_config_path: Optional[str] = None + idx = 0 + for device_val, model, batch, dp, tp, pp, cp in _iter_tests(data): + if model_filter and model not in model_filter: + continue + spec, model_path, hw_path = _build_spec(device_val, model, batch, dp, tp, pp, cp, idx) + specs.append(spec) + actual_lookup[(model, int(batch), int(dp), int(tp), int(pp), int(cp))] = _lookup_actual( + device_val, model, batch, dp, tp, pp, cp + ) + base_model_path = base_model_path or model_path + hw_config_path = hw_config_path or hw_path + idx += 1 + if not specs: + raise ValueError(f"No validation specs generated for device={device} (models={models}).") + return specs, actual_lookup, base_model_path, hw_config_path + + +def run( + device: str = "A100", + models: Optional[Sequence[str]] = None, + emit_logs: bool = True, +): + specs, actual_lookup, base_model_path, hw_config_path = build_specs_for_device( + device, models=models + ) + + validation_results = run_validation_suite( + specs, + base_model_config_path=base_model_path, + base_hardware_config_path=hw_config_path, + result_parser=parse_inference_time, + run_perf_path=RUN_PERF, + ) + + rows = compute_pct_errors(validation_results, actual_lookup) + pct_errors = [r["pct_error"] for r in rows if not math.isnan(r["pct_error"])] + + if emit_logs: + for row in rows: + block = [ + f"\n=== Result (device={row['device']}, model={row['model']}, bs={row['batch']}, dp={row['dp']}, tp={row['tp']}, pp={row['pp']}, cp={row['cp']}) ===" + ] + if row["success"] and not math.isnan(row["pct_error"]): + block.append(f" DeepFlow time: {float(row['inference_time_s']):.2f}s") + block.append(f" Actual time: {float(row['actual_inference_time_s']):.2f}s") + block.append(f" Percent Error: {float(row['pct_error']):.2f}%") + else: + block.append(f" DeepFlow run failed. {(row.get('error') or '')}".rstrip()) + print("\n".join(block)) + + avg_abs_error = sum(pct_errors) / len(pct_errors) if pct_errors else float("nan") + if emit_logs: + print("Average absolute percent error across all training tests: {:.2f}%".format(avg_abs_error)) + return {"avg_abs_error": avg_abs_error, "rows": rows} + + +if __name__ == "__main__": + if __package__ is None: + import sys + sys.path.append(str(Path(__file__).resolve().parents[1])) + run(device="A100", emit_logs=True) From 22bef8fa617ac03c2acaaa736e7fa4fcbd089d67 Mon Sep 17 00:00:00 2001 From: gkarfakis19 Date: Fri, 21 Nov 2025 17:53:26 -0800 Subject: [PATCH 53/77] added full train validation script based on IMEC data. Test is commented out. --- tests/test_IMEC_train.py | 43 ++++ validation_scripts/imec_data/A100_train.csv | 8 +- validation_scripts/nvidia_train.py | 220 ++++++++++++++---- .../a100_80GB_korthikanti.yaml | 144 ++++++++++++ .../hardware-config/a100_80GB_selene_sc.yaml | 144 ++++++++++++ .../model-config/GPT_175_B.yaml | 21 ++ .../model-config/GPT_1T.yaml | 21 ++ .../model-config/GPT_22_B.yaml | 21 ++ .../model-config/GPT_310_B.yaml | 21 ++ .../model-config/GPT_530_B.yaml | 21 ++ validation_scripts/validation_helpers.py | 24 ++ 11 files changed, 643 insertions(+), 45 deletions(-) create mode 100644 tests/test_IMEC_train.py create mode 100644 validation_scripts/validation_configs/hardware-config/a100_80GB_korthikanti.yaml create mode 100644 validation_scripts/validation_configs/hardware-config/a100_80GB_selene_sc.yaml create mode 100644 validation_scripts/validation_configs/model-config/GPT_175_B.yaml create mode 100644 validation_scripts/validation_configs/model-config/GPT_1T.yaml create mode 100644 validation_scripts/validation_configs/model-config/GPT_22_B.yaml create mode 100644 validation_scripts/validation_configs/model-config/GPT_310_B.yaml create mode 100644 validation_scripts/validation_configs/model-config/GPT_530_B.yaml diff --git a/tests/test_IMEC_train.py b/tests/test_IMEC_train.py new file mode 100644 index 00000000..24466d57 --- /dev/null +++ b/tests/test_IMEC_train.py @@ -0,0 +1,43 @@ +import math + +import pytest + +from validation_scripts import nvidia_train + +DEFAULT_THRESHOLD = 20.0 +# Optional overrides per (device, tp_sp_enabled) +GROUP_THRESHOLDS = { + # ("A100_korthi", True): 10.0, +} + + +def _run_and_check(device: str, tp_sp: bool, record_validation): + result = nvidia_train.run( + device=device, + tp_sp=tp_sp, + enable_plot=False, + emit_logs=False, + ) + rows = result["rows"] + assert rows, f"{device} tp_sp={tp_sp} produced no validation rows" + pct_errors = [] + for row in rows: + pct = row["pct_error"] + assert not math.isnan(pct), f"{device} tp_sp={tp_sp} produced NaN pct_error for {row['model']}" + pct_errors.append(pct) + avg_error = sum(pct_errors) / len(pct_errors) + threshold = GROUP_THRESHOLDS.get((device, tp_sp), DEFAULT_THRESHOLD) + record_validation(f"{device} tp_sp={tp_sp}", avg_error, expected_pct=threshold) + assert avg_error <= threshold, f"{device} tp_sp={tp_sp} avg error {avg_error:.2f}% exceeds {threshold}%" + + +# def test_train_a100_korthi_tp_sp_disabled(record_validation): +# _run_and_check("A100_korthi", False, record_validation) + + +# def test_train_a100_korthi_tp_sp_enabled(record_validation): +# _run_and_check("A100_korthi", True, record_validation) + + +# def test_train_a100_selene_tp_sp_disabled(record_validation): +# _run_and_check("A100_selene", False, record_validation) diff --git a/validation_scripts/imec_data/A100_train.csv b/validation_scripts/imec_data/A100_train.csv index db137c48..d580d834 100644 --- a/validation_scripts/imec_data/A100_train.csv +++ b/validation_scripts/imec_data/A100_train.csv @@ -1,3 +1,5 @@ -// placeholder training validation data -device,model,batch,dp,tp,pp,cp,actual -A100,Llama 2-7B,32,2,2,1,1,123.45 +device,model,batch,mb,dp,tp,pp,cp,tp_sp,recomputation,actual +A100_korthi,GPT 22B,4,1,1,8,1,1,True,selective,1.1 +A100_korthi,GPT 175B,64,64,1,8,8,1,True,selective,12.9 +A100_korthi,GPT 530B,280,280,1,8,35,1,True,selective,35.5 +A100_korthi,GPT 1T,512,512,1,8,64,1,True,selective,69.1 diff --git a/validation_scripts/nvidia_train.py b/validation_scripts/nvidia_train.py index 33876204..c87ee18a 100644 --- a/validation_scripts/nvidia_train.py +++ b/validation_scripts/nvidia_train.py @@ -3,12 +3,25 @@ from functools import lru_cache from pathlib import Path from typing import Dict, Iterable, List, Optional, Sequence, Tuple +import sys +import matplotlib.pyplot as plt +import pandas as pd + + +try: + from .validation_helpers import ( + ValidationSpec, + run_validation_suite, + parse_training_time, + ) +except ImportError: + sys.path.append(str(Path(__file__).resolve().parents[1])) + from validation_scripts.validation_helpers import ( # type: ignore + ValidationSpec, + run_validation_suite, + parse_training_time, + ) -from .validation_helpers import ( - ValidationSpec, - run_validation_suite, - parse_inference_time, -) PROJECT_ROOT = os.path.dirname(os.path.dirname(os.path.abspath(__file__))) RUN_PERF = os.path.join(PROJECT_ROOT, "run_perf.py") @@ -19,37 +32,91 @@ HW_CONFIGS = { "A100": "a100_80GB.yaml", + "A100_korthi": "a100_80GB_korthikanti.yaml", + "A100_selene": "a100_80GB_selene_sc.yaml", } MODEL_CONFIGS = { "Llama 2-7B": "Llama2-7B_inf.yaml", + "GPT 22B": "GPT_22_B.yaml", + "GPT 175B": "GPT_175_B.yaml", + "GPT 310B": "GPT_310_B.yaml", + "GPT 530B": "GPT_530_B.yaml", + "GPT 1T": "GPT_1T.yaml", } def _load_data(csv_path: Path): - from validation_scripts.nvidia_inf import _load_data as _load_data_inf # reuse parsing behavior - return _load_data_inf(csv_path) + # Load training validation CSVs; treat leading '//' as comments for path hints. + df = pd.read_csv(csv_path, comment="/") + if "mb" not in df.columns: + raise ValueError(f"Training validation CSV missing required 'mb' column: {csv_path}") + for col in ("batch", "mb", "dp", "tp", "pp", "cp"): + if col in df.columns: + df[col] = df[col].astype(int) + if "tp_sp" in df.columns: + df["tp_sp"] = df["tp_sp"].astype(bool) + if "recomputation" in df.columns: + df["recomputation"] = df["recomputation"].astype(str) + return df @lru_cache(maxsize=None) def _load_device_data(device: str): - csv_path = DATA_DIR / f"{device}_train.csv" - if not csv_path.exists(): - raise FileNotFoundError(f"Validation CSV not found for device {device}: {csv_path}") - return _load_data(csv_path) + candidates = [DATA_DIR / f"{device}_train.csv"] + base = device.split("_")[0] + if base and base != device: + candidates.append(DATA_DIR / f"{base}_train.csv") + candidates.append(DATA_DIR / "A100_train.csv") + + for path in candidates: + if path.exists(): + df = _load_data(path) + if "device" in df.columns: + sub = df[df["device"] == device] + if not sub.empty: + return sub + else: + continue + return df + raise FileNotFoundError(f"Validation CSV not found for device {device}: tried {', '.join(str(p) for p in candidates)}") def _iter_tests(df): for _, row in df.iterrows(): - yield (row["device"], row["model"], row["batch"], row["dp"], row["tp"], row["pp"], row["cp"]) + yield ( + row["device"], + row["model"], + row["batch"], + row["mb"], + row["dp"], + row["tp"], + row["pp"], + row["cp"], + row.get("tp_sp", False), + row.get("recomputation", "partial"), + ) -def _build_spec(device: str, model: str, batch: int, dp: int, tp: int, pp: int, cp: int, idx: int) -> Tuple[ValidationSpec, str, str]: - label = f"{device} {model} bs={batch} dp={dp} tp={tp} pp={pp} cp={cp}" +def _build_spec( + device: str, + model: str, + batch: int, + mb: int, + dp: int, + tp: int, + pp: int, + cp: int, + tp_sp: bool, + recomputation: str, + idx: int, +) -> Tuple[ValidationSpec, str, str]: + label = f"{device} {model} bs={batch} mb={mb} dp={dp} tp={tp} pp={pp} cp={cp} tp_sp={tp_sp}" model_overrides = { "model_param": { "global_batch_size": int(batch), + "recomputation": str(recomputation), } } @@ -57,10 +124,10 @@ def _build_spec(device: str, model: str, batch: int, dp: int, tp: int, pp: int, "parallelism": { "dp": int(dp), "tp": int(tp), - "tp_sp": False, + "tp_sp": bool(tp_sp), "cp": int(cp), "lp": int(pp), - "mb": max(1, int(pp)), # micro-batches align with pipeline by default + "mb": max(1, int(mb)), } } @@ -74,64 +141,82 @@ def _build_spec(device: str, model: str, batch: int, dp: int, tp: int, pp: int, "device": device, "model": model, "batch": int(batch), + "mb": int(mb), "dp": int(dp), "tp": int(tp), "pp": int(pp), "cp": int(cp), + "tp_sp": bool(tp_sp), + "recomputation": str(recomputation), }, order=idx, ) return spec, spec.model_config_path, spec.hardware_config_path # type: ignore[return-value] -def _lookup_actual(device: str, model: str, batch: int, dp: int, tp: int, pp: int, cp: int) -> float: +def _lookup_actual(device: str, model: str, batch: int, mb: int, dp: int, tp: int, pp: int, cp: int, tp_sp: bool, recomputation: str) -> float: data = _load_device_data(device) matches = data.loc[ (data["device"] == device) & (data["model"] == model) & (data["batch"] == int(batch)) + & (data["mb"] == int(mb)) & (data["dp"] == int(dp)) & (data["tp"] == int(tp)) & (data["pp"] == int(pp)) - & (data["cp"] == int(cp)), + & (data["cp"] == int(cp)) + & (data.get("tp_sp", False) == bool(tp_sp)) + & (data.get("recomputation", "").astype(str) == str(recomputation)), "actual", ] if matches.empty: - raise ValueError(f"No reference training time for {device} {model} bs={batch} dp={dp} tp={tp} pp={pp} cp={cp}") + raise ValueError( + f"No reference training time for {device} {model} bs={batch} mb={mb} dp={dp} tp={tp} pp={pp} cp={cp} tp_sp={tp_sp} recomputation={recomputation}" + ) return float(matches.values[0]) -def compute_pct_errors(results, actual_lookup: Dict[Tuple[str, int, int, int, int], float]): +def compute_pct_errors(results, actual_lookup: Dict[Tuple[str, int, int, int, int, int, int, bool, str], float]): rows: List[Dict[str, object]] = [] for res in results: meta = res.spec.metadata or {} model = meta.get("model") batch = int(meta.get("batch")) if "batch" in meta else None + mb = int(meta.get("mb")) if "mb" in meta else None dp = int(meta.get("dp")) if "dp" in meta else None tp = int(meta.get("tp")) if "tp" in meta else None pp = int(meta.get("pp")) if "pp" in meta else None cp = int(meta.get("cp")) if "cp" in meta else None - key = (model, batch, dp, tp, pp, cp) - inf_time = float(res.metrics.get("inference_time_s", float("nan"))) if res.success else float("nan") + tp_sp = bool(meta.get("tp_sp")) if "tp_sp" in meta else False + recomputation = str(meta.get("recomputation")) if "recomputation" in meta else "partial" + key = (model, batch, mb, dp, tp, pp, cp, tp_sp, recomputation) + train_time = float(res.metrics.get("training_time_s", float("nan"))) if res.success else float("nan") actual = actual_lookup.get(key, float("nan")) - if math.isnan(inf_time) or actual == 0 or math.isnan(actual): + if math.isnan(train_time) or actual == 0 or math.isnan(actual): + signed_pct_error = float("nan") pct_error = float("nan") else: - pct_error = abs(inf_time - actual) / actual * 100.0 + signed_pct_error = (train_time - actual) / actual * 100.0 + pct_error = abs(signed_pct_error) rows.append( { "device": meta.get("device"), "model": model, "batch": batch, + "mb": mb, "dp": dp, "tp": tp, "pp": pp, "cp": cp, - "inference_time_s": inf_time, + "tp_sp": tp_sp, + "recomputation": recomputation, + "training_time_s": train_time, "actual_inference_time_s": actual, + "signed_pct_error": signed_pct_error, "pct_error": pct_error, "success": res.success, "error": res.error, + "raw_output": res.raw_output, } ) return rows @@ -141,21 +226,26 @@ def build_specs_for_device( device: str, *, models: Optional[Iterable[str]] = None, -) -> Tuple[List[ValidationSpec], Dict[Tuple[str, int, int, int, int, int], float], str, str]: + tp_sp_only: Optional[bool] = None, +) -> Tuple[List[ValidationSpec], Dict[Tuple[str, int, int, int, int, int, int, bool, str], float], str, str]: data = _load_device_data(device) model_filter = set(models) if models is not None else None specs: List[ValidationSpec] = [] - actual_lookup: Dict[Tuple[str, int, int, int, int, int], float] = {} + actual_lookup: Dict[Tuple[str, int, int, int, int, int, int, bool, str], float] = {} base_model_path: Optional[str] = None hw_config_path: Optional[str] = None idx = 0 - for device_val, model, batch, dp, tp, pp, cp in _iter_tests(data): + for device_val, model, batch, mb, dp, tp, pp, cp, tp_sp, recomputation in _iter_tests(data): if model_filter and model not in model_filter: continue - spec, model_path, hw_path = _build_spec(device_val, model, batch, dp, tp, pp, cp, idx) + if tp_sp_only is not None and bool(tp_sp) != bool(tp_sp_only): + continue + spec, model_path, hw_path = _build_spec( + device_val, model, batch, mb, dp, tp, pp, cp, tp_sp, recomputation, idx + ) specs.append(spec) - actual_lookup[(model, int(batch), int(dp), int(tp), int(pp), int(cp))] = _lookup_actual( - device_val, model, batch, dp, tp, pp, cp + actual_lookup[(model, int(batch), int(mb), int(dp), int(tp), int(pp), int(cp), bool(tp_sp), str(recomputation))] = _lookup_actual( + device_val, model, batch, mb, dp, tp, pp, cp, tp_sp, recomputation ) base_model_path = base_model_path or model_path hw_config_path = hw_config_path or hw_path @@ -168,18 +258,22 @@ def build_specs_for_device( def run( device: str = "A100", models: Optional[Sequence[str]] = None, + tp_sp: Optional[bool] = None, + enable_plot: bool = True, + show_progress: bool = False, emit_logs: bool = True, ): specs, actual_lookup, base_model_path, hw_config_path = build_specs_for_device( - device, models=models + device, models=models, tp_sp_only=tp_sp ) validation_results = run_validation_suite( specs, base_model_config_path=base_model_path, base_hardware_config_path=hw_config_path, - result_parser=parse_inference_time, + result_parser=parse_training_time, run_perf_path=RUN_PERF, + show_progress=show_progress, ) rows = compute_pct_errors(validation_results, actual_lookup) @@ -188,24 +282,66 @@ def run( if emit_logs: for row in rows: block = [ - f"\n=== Result (device={row['device']}, model={row['model']}, bs={row['batch']}, dp={row['dp']}, tp={row['tp']}, pp={row['pp']}, cp={row['cp']}) ===" + f"\n=== Result (device={row['device']}, model={row['model']}, bs={row['batch']}, mb={row['mb']}, dp={row['dp']}, tp={row['tp']}, pp={row['pp']}, cp={row['cp']}, tp_sp={row['tp_sp']}, recomputation={row['recomputation']}) ===" ] if row["success"] and not math.isnan(row["pct_error"]): - block.append(f" DeepFlow time: {float(row['inference_time_s']):.2f}s") - block.append(f" Actual time: {float(row['actual_inference_time_s']):.2f}s") - block.append(f" Percent Error: {float(row['pct_error']):.2f}%") + block.append(f" DeepFlow train time: {float(row['training_time_s']):.2f}s") + block.append(f" Actual train time: {float(row['actual_inference_time_s']):.2f}s") + block.append(f" Percent Error: {float(row['signed_pct_error']):+.2f}%") else: block.append(f" DeepFlow run failed. {(row.get('error') or '')}".rstrip()) + if row.get("raw_output"): + block.append(" --- Raw output ---") + block.append(str(row["raw_output"]).strip()) print("\n".join(block)) + plot_path = None + if enable_plot and rows: + out_dir = Path(PROJECT_ROOT) / "output" / "validation" / "train" + out_dir.mkdir(parents=True, exist_ok=True) + labels = [] + errors = [] + colors = [] + color_map = {"full": "#ff8c00", "partial": "#1f77b4", "selective": "#1f77b4"} + for row in rows: + gpus = int(row["dp"]) * int(row["tp"]) * int(row["pp"]) * int(row["cp"]) + labels.append(f"{row['model']} x{gpus}") + errors.append(row["signed_pct_error"]) + colors.append(color_map.get(str(row.get("recomputation")).lower(), "#1f77b4")) + fig_w = max(6.0, 0.6 * len(labels)) + fig, ax = plt.subplots(figsize=(fig_w, 4)) + bars = ax.bar(range(len(errors)), errors, color=colors) + ax.set_xticks(range(len(labels))) + ax.set_xticklabels(labels, rotation=20, ha="right", fontsize=9) + ax.set_ylabel("Percent Error") + ax.set_title(f"Training validation ({device})") + # Legend + handles = [ + plt.Rectangle((0, 0), 1, 1, color="#ff8c00"), + plt.Rectangle((0, 0), 1, 1, color="#1f77b4"), + ] + ax.legend(handles, ["recompute=full", "recompute=partial/selective"]) + for rect, err in zip(bars, errors): + if math.isnan(err): + continue + height = rect.get_height() + ax.text(rect.get_x() + rect.get_width() / 2, height, f"{err:.1f}%", ha="center", va="bottom", fontsize=8) + ax.grid(axis="y", linestyle="--", alpha=0.3) + plot_path = out_dir / f"train_{device}.png" + fig.tight_layout() + fig.savefig(plot_path, dpi=200) + plt.close(fig) + avg_abs_error = sum(pct_errors) / len(pct_errors) if pct_errors else float("nan") if emit_logs: print("Average absolute percent error across all training tests: {:.2f}%".format(avg_abs_error)) - return {"avg_abs_error": avg_abs_error, "rows": rows} + if plot_path: + print(f"Saved plot: {plot_path}") + return {"avg_abs_error": avg_abs_error, "rows": rows, "plot": plot_path} if __name__ == "__main__": - if __package__ is None: - import sys - sys.path.append(str(Path(__file__).resolve().parents[1])) - run(device="A100", emit_logs=True) + print("=== Running A100_korthi training validation ===") + run(device="A100_korthi", emit_logs=True, show_progress=True) + # print("\n=== Running A100_selene training validation ===") + # run(device="A100_selene", emit_logs=True, show_progress=True) diff --git a/validation_scripts/validation_configs/hardware-config/a100_80GB_korthikanti.yaml b/validation_scripts/validation_configs/hardware-config/a100_80GB_korthikanti.yaml new file mode 100644 index 00000000..61500997 --- /dev/null +++ b/validation_scripts/validation_configs/hardware-config/a100_80GB_korthikanti.yaml @@ -0,0 +1,144 @@ +sw_param: + kernel_launch_overhead: 9e-6 # (s) + # dp_zero_stage selects ZeRO-style data-parallel sharding: 0=DDP, 1=ZeRO-1, 2=ZeRO-2, 3=ZeRO-3 + # only used if data parallelism (dp) > 1 + dp_zero_stage: 1 + precision: + # tensor_format accepts numeric byte counts or dtype strings (fp16, bf16, fp32). + # Take care that the number of flops defined below has been carefully tuned for FP16/BF16 formats. + # For FP32 or other formats they need to be manually adjusted. + # The following are equivalent: (tensor_format: 2.0, tensor_format: fp16, tensor_format: bf16). + tensor_format: bf16 # baseline storage width for activations/gradients. + # mixed_precision: true | false + # -> false forces all activations/weights/gradients/optimizer states/stats to use tensor_format bytes. + # -> true enables AMP-style FP32 accumulations/optimizer states depending on param_storage_mode. + # if true, optimizer states and stats are always FP32. + mixed_precision: true + # param_storage_mode: as_tensor_format | tensor_plus_fp32_master | fp32_params. + # -> as_tensor_format makes weights/gradients use tensor_format bytes. + # -> tensor_plus_fp32_master makes weights/gradients use tensor_format bytes and + # keeps a separate FP32 master copy for weights. + # -> fp32_params makes weights *and* gradients use FP32 bytes without a master copy. + param_storage_mode: tensor_plus_fp32_master + # kv_cache precision may be set as bytes, dtype string, or "as_tensor_format". + kv_cache: as_tensor_format + h2d_bandwidth: -1 # Bytes/s; set -1 to disable H2D timing +tech_param: + core: + operating_frequency: 1.41e9 # (Hz) + num_bundles: 108 + nominal_flop_rate_per_mcu: 512 + nominal_energy_per_flop: 1.8e-13 # (J) + num_mcu_per_bundle: 4 # number of tensorcores per SM + FMA_d1: 8 # tensor core can compute the results for an d1 x d2 x d1 mixed-precision matrix multiplication per clock + FMA_d2: 4 + dataflow: best # {wst | ast | ost | best}: wst: weight stationary, ast: activation stationary, ost: output stationary, best: best of all 3 + util: 1.0 + DRAM: + size: 80 GB + bandwidth: 1986 GB # (/s) + dynamic_energy_per_bit: 4.5e-12 # (J) + latency: 100e-9 # (s) + util: 1.0 + SRAM-L2: + size: 40 MB + bandwidth: 7050 GB + dynamic_energy_per_bit: 130e-15 # (J) + latency: 0 + util: 1 + SRAM-L1: + size: 20736 KB + bandwidth: 18 TB + dynamic_energy_per_bit: 130e-15 # (J) + latency: 0 + util: 1 + SRAM-R: + size: 27 MB + bandwidth: 122 TB + dynamic_energy_per_bit: 110e-15 # (J) + latency: 0 + util: 1 + +memory_hierarchy: + l0: #Register Memory + type: SRAM-R + scope: mcu + l1: #Shared Memory + type: SRAM-L1 + scope: mcu-bundle + l2: #L2 + type: SRAM-L2 + scope: global + l3: #Global Memory + type: DRAM + scope: global + +parallelism: + dp: 2 # data parallelism. Acts as replica count during inference + tp: 2 # tensor parallelism + tp_sp: true # tensor-sequence parallelism + cp: 2 # context parallelism. Not supported for inference. + lp: 4 # pipeline parallelism. Number of pipeline stages + # mb is not a true parallelism parameter, does not affect num of GPUs. + mb: 4 # micro-batches (for best results set to >= lp). + +network: + # Global faulty links apply to the active AstraSim dimension: [[src, dst, weight]] + # faulty_links: [[0,1,0.6],[16,17,0.1]] + dimensions: + # currently only support astrasim integration with all devices in one node. + - id: dim0 + label: torus2d_tp + size: auto + topology: + type: FullyConnected + bandwidth: 300 GB # based on num nvlinks per GPU * 25 GB/s per link + energy_per_bit: 8e-12 # (J/bit) + latency: 1e-6 + util: 1.0 + optimize_2dmap: false + collective_override: {} # overrides execution_backend.collectives when provided + parallelisms: [tp, cp] + - id: dim1 + label: inter_torus_infiniband + size: auto + topology: + type: FullyConnected + bandwidth: 25 GB # 25GB/s per GPU, N GPUs per node. NOT AGGREGATE. + latency: 5e-6 + energy_per_bit: 8e-12 # (J/bit) + util: 1.0 + collective_override: {} # overrides execution_backend.collectives when provided + parallelisms: [lp,dp] + - id: dim2 + label: inter_torus_infiniband_dp + size: auto + topology: + type: Ring + bandwidth: 25 GB + latency: 5e-6 + energy_per_bit: 8e-12 # (J/bit) + util: 1.0 + parallelisms: [] + collective_override: {} # overrides execution_backend.collectives when provided + +execution_backend: + model: astra # analytical | astra + astra: + backend: astra + mode: full_astrasim_hierarchical # hybrid | full_astrasim_hierarchical | full_astrasim_flattened + collectives: + # Default collective algorithms. Per-dimension collective_override entries can override these. + # Allowed values per op: auto | ring | direct | halvingDoubling | doubleBinaryTree | mesh | mesh2d | hypercube | torus2d + # Note: leave to auto unless you have a specific reason to change it. + all_gather: auto + all_reduce: auto + reduce_scatter: auto + all_to_all: auto + sys_options: + endpoint_delay: 5 + active_chunks_per_dimension: 32 # Higher value means significantly slower astrasim runtime, but more accurate especially for tensor parallelism + preferred_dataset_splits: 4 # Should be same as active_chunks_per_dimension + +inference: + kvcache_type: hbm_only # Only supported mode for now. diff --git a/validation_scripts/validation_configs/hardware-config/a100_80GB_selene_sc.yaml b/validation_scripts/validation_configs/hardware-config/a100_80GB_selene_sc.yaml new file mode 100644 index 00000000..8aad58ad --- /dev/null +++ b/validation_scripts/validation_configs/hardware-config/a100_80GB_selene_sc.yaml @@ -0,0 +1,144 @@ +sw_param: + kernel_launch_overhead: 9e-6 # (s) + # dp_zero_stage selects ZeRO-style data-parallel sharding: 0=DDP, 1=ZeRO-1, 2=ZeRO-2, 3=ZeRO-3 + # only used if data parallelism (dp) > 1 + dp_zero_stage: 1 + precision: + # tensor_format accepts numeric byte counts or dtype strings (fp16, bf16, fp32). + # Take care that the number of flops defined below has been carefully tuned for FP16/BF16 formats. + # For FP32 or other formats they need to be manually adjusted. + # The following are equivalent: (tensor_format: 2.0, tensor_format: fp16, tensor_format: bf16). + tensor_format: bf16 # baseline storage width for activations/gradients. + # mixed_precision: true | false + # -> false forces all activations/weights/gradients/optimizer states/stats to use tensor_format bytes. + # -> true enables AMP-style FP32 accumulations/optimizer states depending on param_storage_mode. + # if true, optimizer states and stats are always FP32. + mixed_precision: true + # param_storage_mode: as_tensor_format | tensor_plus_fp32_master | fp32_params. + # -> as_tensor_format makes weights/gradients use tensor_format bytes. + # -> tensor_plus_fp32_master makes weights/gradients use tensor_format bytes and + # keeps a separate FP32 master copy for weights. + # -> fp32_params makes weights *and* gradients use FP32 bytes without a master copy. + param_storage_mode: as_tensor_format + # kv_cache precision may be set as bytes, dtype string, or "as_tensor_format". + kv_cache: as_tensor_format + h2d_bandwidth: -1 # Bytes/s; set -1 to disable H2D timing +tech_param: + core: + operating_frequency: 1.41e9 # (Hz) + num_bundles: 108 + nominal_flop_rate_per_mcu: 512 + nominal_energy_per_flop: 1.8e-13 # (J) + num_mcu_per_bundle: 4 # number of tensorcores per SM + FMA_d1: 8 # tensor core can compute the results for an d1 x d2 x d1 mixed-precision matrix multiplication per clock + FMA_d2: 4 + dataflow: best # {wst | ast | ost | best}: wst: weight stationary, ast: activation stationary, ost: output stationary, best: best of all 3 + util: 1.0 + DRAM: + size: 80 GB + bandwidth: 1986 GB # (/s) + dynamic_energy_per_bit: 4.5e-12 # (J) + latency: 100e-9 # (s) + util: 1 + SRAM-L2: + size: 40 MB + bandwidth: 7050 GB + dynamic_energy_per_bit: 130e-15 # (J) + latency: 0 + util: 1 + SRAM-L1: + size: 20736 KB + bandwidth: 18 TB + dynamic_energy_per_bit: 130e-15 # (J) + latency: 0 + util: 1 + SRAM-R: + size: 27 MB + bandwidth: 122 TB + dynamic_energy_per_bit: 110e-15 # (J) + latency: 0 + util: 1 + +memory_hierarchy: + l0: #Register Memory + type: SRAM-R + scope: mcu + l1: #Shared Memory + type: SRAM-L1 + scope: mcu-bundle + l2: #L2 + type: SRAM-L2 + scope: global + l3: #Global Memory + type: DRAM + scope: global + +parallelism: + dp: 2 # data parallelism. Acts as replica count during inference + tp: 2 # tensor parallelism + tp_sp: true # tensor-sequence parallelism + cp: 2 # context parallelism. Not supported for inference. + lp: 4 # pipeline parallelism. Number of pipeline stages + # mb is not a true parallelism parameter, does not affect num of GPUs. + mb: 4 # micro-batches (for best results set to >= lp). + +network: + # Global faulty links apply to the active AstraSim dimension: [[src, dst, weight]] + # faulty_links: [[0,1,0.6],[16,17,0.1]] + dimensions: + # currently only support astrasim integration with all devices in one node. + - id: dim0 + label: torus2d_tp + size: auto + topology: + type: FullyConnected + bandwidth: 300 GB # based on num nvlinks per GPU * 25 GB/s per link + latency: 1e-6 + energy_per_bit: 8e-12 # (J/bit) + util: 1.0 + optimize_2dmap: false + collective_override: {} # overrides execution_backend.collectives when provided + parallelisms: [tp, cp] + - id: dim1 + label: inter_torus_infiniband + size: auto + topology: + type: FullyConnected + bandwidth: 25 GB # 25GB/s per GPU, N GPUs per node. NOT AGGREGATE. + latency: 5e-6 + energy_per_bit: 8e-12 # (J/bit) + util: 1.0 + collective_override: {} # overrides execution_backend.collectives when provided + parallelisms: [lp,dp] + - id: dim2 + label: inter_torus_infiniband_dp + size: auto + topology: + type: Ring + bandwidth: 25 GB + latency: 5e-6 + energy_per_bit: 8e-12 # (J/bit) + util: 1.0 + parallelisms: [] + collective_override: {} # overrides execution_backend.collectives when provided + +execution_backend: + model: astra # analytical | astra + astra: + backend: astra + mode: full_astrasim_hierarchical # hybrid | full_astrasim_hierarchical | full_astrasim_flattened + collectives: + # Default collective algorithms. Per-dimension collective_override entries can override these. + # Allowed values per op: auto | ring | direct | halvingDoubling | doubleBinaryTree | mesh | mesh2d | hypercube | torus2d + # Note: leave to auto unless you have a specific reason to change it. + all_gather: auto + all_reduce: auto + reduce_scatter: auto + all_to_all: auto + sys_options: + endpoint_delay: 5 + active_chunks_per_dimension: 32 # Higher value means significantly slower astrasim runtime, but more accurate especially for tensor parallelism + preferred_dataset_splits: 4 # Should be same as active_chunks_per_dimension + +inference: + kvcache_type: hbm_only # Only supported mode for now. diff --git a/validation_scripts/validation_configs/model-config/GPT_175_B.yaml b/validation_scripts/validation_configs/model-config/GPT_175_B.yaml new file mode 100644 index 00000000..ba591aee --- /dev/null +++ b/validation_scripts/validation_configs/model-config/GPT_175_B.yaml @@ -0,0 +1,21 @@ +model_param: + mode: "LLM" + run_type: "training" + tied_embeddings: false + model_type: "gpt" + global_batch_size: 8 + gradient_accumulation_steps: 1 + recomputation: partial + seq_len: 2048 + hidden_dim: 12288 + attention: + attention_type: "mha" + num_heads: 96 + kv_heads: 96 + use_flashattention: false + attention_tile_size: 128 + num_experts: 1 + top_k: 1 + intermediate_size: 49152 + vocab_size: 51200 + num_layers: 96 diff --git a/validation_scripts/validation_configs/model-config/GPT_1T.yaml b/validation_scripts/validation_configs/model-config/GPT_1T.yaml new file mode 100644 index 00000000..5f631bd9 --- /dev/null +++ b/validation_scripts/validation_configs/model-config/GPT_1T.yaml @@ -0,0 +1,21 @@ +model_param: + mode: "LLM" + run_type: "training" + tied_embeddings: false + model_type: "gpt" + global_batch_size: 8 + gradient_accumulation_steps: 1 + recomputation: partial + seq_len: 2048 + hidden_dim: 25600 + attention: + attention_type: "mha" + num_heads: 160 + kv_heads: 160 + use_flashattention: false + attention_tile_size: 128 + num_experts: 1 + top_k: 1 + intermediate_size: 102400 + vocab_size: 51200 + num_layers: 128 diff --git a/validation_scripts/validation_configs/model-config/GPT_22_B.yaml b/validation_scripts/validation_configs/model-config/GPT_22_B.yaml new file mode 100644 index 00000000..67598d95 --- /dev/null +++ b/validation_scripts/validation_configs/model-config/GPT_22_B.yaml @@ -0,0 +1,21 @@ +model_param: + mode: "LLM" + run_type: "training" + tied_embeddings: false + model_type: "gpt" + global_batch_size: 8 + gradient_accumulation_steps: 1 + recomputation: partial + seq_len: 2048 + hidden_dim: 6144 + attention: + attention_type: "mha" + num_heads: 64 + kv_heads: 64 + use_flashattention: false + attention_tile_size: 128 + num_experts: 1 + top_k: 1 + intermediate_size: 24576 + vocab_size: 51200 + num_layers: 48 diff --git a/validation_scripts/validation_configs/model-config/GPT_310_B.yaml b/validation_scripts/validation_configs/model-config/GPT_310_B.yaml new file mode 100644 index 00000000..3988181a --- /dev/null +++ b/validation_scripts/validation_configs/model-config/GPT_310_B.yaml @@ -0,0 +1,21 @@ +model_param: + mode: "LLM" + run_type: "training" + tied_embeddings: false + model_type: "gpt" + global_batch_size: 8 + gradient_accumulation_steps: 1 + recomputation: partial + seq_len: 2048 + hidden_dim: 16384 + attention: + attention_type: "mha" + num_heads: 128 + kv_heads: 128 + use_flashattention: false + attention_tile_size: 128 + num_experts: 1 + top_k: 1 + intermediate_size: 65536 + vocab_size: 51200 + num_layers: 96 diff --git a/validation_scripts/validation_configs/model-config/GPT_530_B.yaml b/validation_scripts/validation_configs/model-config/GPT_530_B.yaml new file mode 100644 index 00000000..25fcbb10 --- /dev/null +++ b/validation_scripts/validation_configs/model-config/GPT_530_B.yaml @@ -0,0 +1,21 @@ +model_param: + mode: "LLM" + run_type: "training" + tied_embeddings: false + model_type: "gpt" + global_batch_size: 8 + gradient_accumulation_steps: 1 + recomputation: partial + seq_len: 2048 + hidden_dim: 20480 + attention: + attention_type: "mha" + num_heads: 128 + kv_heads: 128 + use_flashattention: false + attention_tile_size: 128 + num_experts: 1 + top_k: 1 + intermediate_size: 81920 + vocab_size: 51200 + num_layers: 105 diff --git a/validation_scripts/validation_helpers.py b/validation_scripts/validation_helpers.py index 6c4cdb0d..a361cd77 100644 --- a/validation_scripts/validation_helpers.py +++ b/validation_scripts/validation_helpers.py @@ -291,6 +291,7 @@ def run_validation_suite( env_overrides: Optional[Dict[str, str]] = None, cache_mode: str = "NO_CACHE", extra_run_perf_args: Optional[Sequence[str]] = None, + show_progress: bool = False, ) -> List[ValidationResult]: if not specs: return [] @@ -327,6 +328,14 @@ def run_validation_suite( extra_run_perf_args=extra_args, ) results: List[ValidationResult] = [] + progress = None + if show_progress: + try: + from tqdm import tqdm # type: ignore + progress = tqdm(total=len(assigned_specs), desc="validation", leave=True) + except Exception: + progress = None + with ProcessPoolExecutor(max_workers=worker_count, initializer=_worker_init, initargs=(context,)) as executor: future_map = {executor.submit(_execute_spec, spec): spec for spec in assigned_specs} for future in as_completed(future_map): @@ -346,6 +355,10 @@ def run_validation_suite( hardware_config_used=None, ) results.append(result) + if progress is not None: + progress.update(1) + if progress is not None: + progress.close() results.sort(key=lambda item: item.spec.order) return results @@ -357,6 +370,9 @@ def run_validation_suite( _INF_TIME_REGEX = re.compile( r"LLM inference time:\s*([0-9]+(?:\.[0-9]+)?)s" ) +_TRAIN_TIME_REGEX = re.compile( + r"Training time for batch:\s*([0-9]+(?:\.[0-9]+)?)s" +) def parse_decode_time(output: str, spec: ValidationSpec) -> Dict[str, Any]: match = _DECODE_TIME_REGEX.search(output) @@ -370,6 +386,13 @@ def parse_inference_time(output: str, spec: ValidationSpec) -> Dict[str, Any]: raise ValueError(f"Failed to parse inference time for experiment '{spec.label}'.") return {"inference_time_s": float(match.group(1))} + +def parse_training_time(output: str, spec: ValidationSpec) -> Dict[str, Any]: + match = _TRAIN_TIME_REGEX.search(output) + if not match: + raise ValueError(f"Failed to parse training time for experiment '{spec.label}'.") + return {"training_time_s": float(match.group(1))} + # ---- Pytest integration for quiet validation reporting ---- _VALIDATION_METRICS: List[Tuple[str, str, float, Optional[float]]] = [] @@ -408,5 +431,6 @@ def pytest_terminal_summary(terminalreporter, exitstatus): "run_validation_suite", "parse_decode_time", "parse_inference_time", + "parse_training_time", "record_validation", ] From b87b2ece304fb4946b79e9c90fb24ca6b3e53bce Mon Sep 17 00:00:00 2001 From: gkarfakis19 Date: Sat, 22 Nov 2025 18:18:29 -0800 Subject: [PATCH 54/77] added optimizer FLOPS to bwd pass, big cleanup of timing and comm rules and apply grad --- LLM_execution.py | 60 ++ examples/korthi.sh | 11 + simulate_LLM.py | 33 +- time_calculation.py | 104 +--- time_calculation_LLM.py | 534 ++++++------------ timing_model.py | 17 +- .../a100_80GB_korthikanti.yaml | 14 +- .../model-config/GPT_22_B.yaml | 4 +- 8 files changed, 317 insertions(+), 460 deletions(-) create mode 100755 examples/korthi.sh diff --git a/LLM_execution.py b/LLM_execution.py index 9c9a655d..e68af4ec 100644 --- a/LLM_execution.py +++ b/LLM_execution.py @@ -216,6 +216,35 @@ def _clone(self, obj: Any) -> Any: obj.duration, fwd=obj.fwd, ) + elif "optimizer" in obj.name: + # Optimizer nodes need to be expanded per TP rank + cloned_nodes = [] + for tp_rank in range(self._par_degree): + hw_id = self._hw_id_for_rank(obj.hw_id, tp_rank) + cloned_node = simulate_LLM.Node( + name=f"{obj.name}_rank{tp_rank}", + op_id=self._next_op_id(), + hw_id=hw_id, + duration=obj.duration, + fwd=obj.fwd, + ) + cloned_nodes.append(cloned_node) + + cloned_tuple = tuple(cloned_nodes) + self._clone_cache[obj_id] = cloned_tuple + + # We also need to copy metadata if any + for cloned_node in cloned_nodes: + self._copy_metadata(obj, cloned_node) + + # Attach children (if any) to all clones? + # Let's follow the pattern: + for child in getattr(obj, "children", []): + child_clone = self._clone(child) + if child_clone is not None: + self._attach(cloned_tuple, child_clone) + + return cloned_tuple else: cloned = simulate_LLM.Node( obj.name, @@ -225,6 +254,7 @@ def _clone(self, obj: Any) -> Any: fwd=obj.fwd, ) + # following _expand_transformer_node logic, we need to find siblings that are zero3 tp_shard=True. zero3_attachments: List[simulate_LLM.Edge] = [] for parent in getattr(obj, "parents", []): @@ -513,6 +543,14 @@ def _expand_transformer_node(self, node: simulate_LLM.Node) -> Tuple[Any, ...]: child_clone = self._clone(child) if child_clone is None: continue + + # Special handling for optimizer nodes: 1-to-1 connection + if isinstance(child, simulate_LLM.Node) and "optimizer" in child.name and isinstance(child_clone, (list, tuple)): + if len(child_clone) == len(rank_tails): + for r in range(len(rank_tails)): + rank_tails[r].add_child(child_clone[r]) + continue + self._attach(downstream_parents, child_clone) for zero3_edge in zero3_attachments: @@ -1069,6 +1107,28 @@ def _run_full_astrasim_hierarchical(self) -> ExecutionResult: if transformer_time is not None: self._apply_transformer_time(transformer_time) + if _env_flag("DEEPFLOW_VISUALIZE_GRAPHS") and self.transformer_graph: + transformer_timed_forward_root = self.transformer_graph.convert_comm_sizes_to_times( + self.transformer_forward_root, + self.time_calc.network_model, + self.interconnect_params, + ) + transformer_timed_backward_root = self.transformer_graph.convert_comm_sizes_to_times( + self.transformer_backward_root, + self.time_calc.network_model, + self.interconnect_params, + ) + self.transformer_graph.save_graph( + transformer_timed_forward_root, + self.time_calc.output_dir, + "/hierarchical_graph_transformer_forward", + ) + self.transformer_graph.save_graph( + transformer_timed_backward_root, + self.time_calc.output_dir, + "/hierarchical_graph_transformer_backward", + ) + dp_count = getattr(self.time_calc, "dp", 1) or 1 if not self.pipeline_root: raise RuntimeError("Pipeline graph root is not available for AstraSim execution") diff --git a/examples/korthi.sh b/examples/korthi.sh new file mode 100755 index 00000000..d2232137 --- /dev/null +++ b/examples/korthi.sh @@ -0,0 +1,11 @@ +#!/bin/bash + +# Runs the first Korthikanti training case (GPT-22B, bs=4, mb=1, dp=1, tp=8, pp=1, cp=1, tp_sp=true) +# with all debug/visualization flags enabled. + +DEEPFLOW_PERSIST_ASTRASIM_ARTIFACTS=1 \ +DEEPFLOW_VISUALIZE_GRAPHS=1 \ +DEEPFLOW_PERSIST_ARTIFACT_VIZ=1 \ +python3 run_perf.py \ + --hardware_config validation_scripts/validation_configs/hardware-config/a100_80GB_korthikanti.yaml \ + --model_config validation_scripts/validation_configs/model-config/GPT_22_B.yaml diff --git a/simulate_LLM.py b/simulate_LLM.py index 81cdf694..08c345cc 100644 --- a/simulate_LLM.py +++ b/simulate_LLM.py @@ -578,7 +578,7 @@ def traverse_and_convert(node, visited=None): traverse_and_convert(roots) return roots - def construct_fwd_bwd_graph(self, include_backward: bool = True): + def construct_fwd_bwd_graph(self, include_backward: bool = True, include_optimizer: bool = False): embedding_node = [] data_batch_node = [] softmax_node = [] @@ -1025,6 +1025,37 @@ def attach_parallel_edge(target, gather_edge, skip_non_comm_children=None, skip_ op_id += 1 attach_parallel_edge(host, gather_edge) + if include_backward and include_optimizer: + optimizer_time = self._time("optimizer") + if optimizer_time > 0: + # Find the last node for each stage to attach the optimizer + # For stage 0, it's the embedding backward of the last microbatch + # For other stages, it's the first transformer layer (lowest index) of that stage for the last microbatch + + # Helper to find min layer for each stage + stage_min_layer = {} + for l in range(self.num_layer): + stage = self._stage_for_layer(l) + if stage not in stage_min_layer: + stage_min_layer[stage] = l + else: + stage_min_layer[stage] = min(stage_min_layer[stage], l) + + # bwd flips mbs, so we attach to first one, not last. + for stage in range(self.lp): + last_node = None + if stage == 0: + last_node = embedding_node_b[0] + else: + min_l = stage_min_layer.get(stage) + if min_l is not None: + last_node = transformer_nodes_b[0][min_l] + + if last_node: + opt_node = Node(f"optimizer_stage{stage}", op_id, stage, optimizer_time, fwd=False) + op_id += 1 + last_node.add_child(opt_node) + return root_forward_entry def construct_transformer_graph(self, direction: str = "both"): diff --git a/time_calculation.py b/time_calculation.py index 6ba2306f..5c22d99b 100644 --- a/time_calculation.py +++ b/time_calculation.py @@ -1470,86 +1470,38 @@ def getCb(self): # Reduction and all-gather time estimation - def gradClipping(self, Dim0=None, Dim1=None, name=None): - if Dim0 == None: - Dim0 = 2 * self.D - if Dim1 == None: - Dim1 = self.G * self.D - if name == None: - name = "Hidden" - # t_list[i] * clip_norm / max(global_norm, clip_norm) - # where: - # global_norm = sqrt(sum([l2norm(t)**2 for t in t_list])) - - norm_comp = Dim0 * Dim1 * 2 - # 1: power 2 - # 1: summ - norm_mem = (Dim0 * Dim1 * 1) * self.precision_bytes - # 1: one read per element and power it by 2 in local registers anfd - # summing to local acc - - clip_comp = Dim0 * Dim1 * 2 - # 1: pointwise mul - # 1: pointwise div - - clip_mem = (Dim0 * Dim1 * 2) * self.precision_bytes - # 1: one read for pointwise mul - # 1: one write for pointwise div - - gradclip_mem = norm_mem + clip_mem - gradclip_comp = norm_comp + clip_comp - - gradclip_time = self.roofline( - gradclip_comp, gradclip_mem, name="pointwise-grad-clipping" - ) - - if self.debug: - print( - "({}) gradclip_flop: {:,}, gradclip_mem: {:,}".format( - name, gradclip_comp, gradclip_mem - ) - ) - print("({}) gradclip_time: {:,}\n".format(name, gradclip_time)) + def grad_clipping(self, num_params: int) -> float: + """Gradient clipping (L2 norm + scale) cost.""" + norm_comp = num_params * 2 # square + sum + clip_comp = num_params * 2 # scale + divide + clip_mem = num_params * 2 * self.precision.gradients # read + write - return gradclip_time + gradclip_mem = clip_mem + gradclip_comp = norm_comp + clip_comp - def apply_grad(self, Dim0=None, Dim1=None, name=None): - if Dim0 == None: - Dim0 = 2 * self.D - if Dim1 == None: - Dim1 = self.G * self.D - if name == None: - name = "Hidden" + return self.roofline(gradclip_comp, gradclip_mem, name="pointwise-grad-clipping") - applyGrad_comp = Dim0 * Dim1 * 3 - # 3: one pointwise division by scalar after reducing all the gradients, - # one final addition of gradients to the weights - # one multiply by learning rate - applyGrad_mem = ( - (1 * Dim0 * Dim1 * self.precision_bytes) - + (2 * Dim0 * Dim1 * self.precision_bytes) - + (1 * Dim0 * Dim1 * self.precision_bytes) + def apply_grad(self, num_params: int) -> float: + """Approximate optimizer update cost (Adam/AdamW-style) per parameter tensor.""" + OPT_FLOPS_PER_PARAM = 14.0 # typical for Adam + bytes_per_param_read = ( + self.precision.parameters + + self.precision.gradients + + 2 * self.precision.optimizer_states + + self.precision.master_parameters ) - # 1: read for pointiwse div - # 2: 1 reads and one write for pointwise add - # 1: one write for multiplication by lr - applyGrad_time = self.roofline( - applyGrad_comp, applyGrad_mem, name="pointwise-applyGrad" + bytes_per_param_write = ( + self.precision.parameters + + 2 * self.precision.optimizer_states + + self.precision.master_parameters ) - clip_time = self.gradClipping(Dim0, Dim1, name) - - grad_time = applyGrad_time + clip_time - - if self.debug: - print( - "({}) applyGrad_flop: {:,}, applyGrad_mem: {:,}".format( - name, applyGrad_comp, applyGrad_mem - ) - ) - print("({}) applyGrad_time: {:,}\n".format(name, applyGrad_time)) + apply_grad_comp = num_params * OPT_FLOPS_PER_PARAM + apply_grad_mem = num_params * (bytes_per_param_read + bytes_per_param_write) - return grad_time + apply_grad_time = self.roofline(apply_grad_comp, apply_grad_mem, name="pointwise-applyGrad") + clip_time = self.grad_clipping(num_params) + return apply_grad_time + clip_time @@ -1731,7 +1683,7 @@ def getDataParallelReduction(self, k, n, dim1, dim2, name): local_bytes=0.0, debug_label=name or "comm", ) - apply_grad_time = self.apply_grad(Dim0=k / dim1, Dim1=n, name=name) + apply_grad_time = self.apply_grad(int(math.ceil((k / dim1) * n))) elif self.kp_hidden_type == 2: # RC total_bytes = math.ceil(self.precision_bytes * (k / dim1) * (n / dim2)) @@ -1757,7 +1709,7 @@ def getDataParallelReduction(self, k, n, dim1, dim2, name): local_bytes=3 * total_bytes, debug_label=name or "comm", ) - apply_grad_time = self.apply_grad(Dim0=k, Dim1=n / dim2, name=name) + apply_grad_time = self.apply_grad(int(math.ceil(k * (n / dim2)))) else: reduction_time_wt_kp = 0 total_bytes = math.ceil(self.precision_bytes * k * n) @@ -1770,7 +1722,7 @@ def getDataParallelReduction(self, k, n, dim1, dim2, name): local_bytes=0.0, debug_label=name or "comm", ) - apply_grad_time = self.apply_grad(Dim0=k, Dim1=n, name=name) + apply_grad_time = self.apply_grad(int(math.ceil(k * n))) if self.debug: print(f"reduction_time_wt_kp: {reduction_time_wt_kp}") print(f"reduction_time_wt_dp: {reduction_time_wt_dp}") diff --git a/time_calculation_LLM.py b/time_calculation_LLM.py index e76783ea..1f9ec742 100644 --- a/time_calculation_LLM.py +++ b/time_calculation_LLM.py @@ -12,7 +12,7 @@ from time_calculation import TimeCalculation from itertools import zip_longest # for element-wise aggregation of memory access lists from timing_model import CommSpec, DirectionTiming, OperationTiming, OperationGroup - +import yaml def _env_flag(name: str) -> bool: value = os.environ.get(name) if value is None: @@ -44,67 +44,67 @@ class ParallelismMode(Enum): # Map each parallelism mode to operation-level collective specs used across the # metadata pipeline. Each spec records the collective kind, the participant -# scope (tp/cp/seq/etc.), the interconnect label, and an identifying suffix. +# scope (tp/cp/seq/etc.), and the interconnect label. COMM_RULE_DEFAULT_KEY = "__default__" COMMUNICATION_RULES: Dict[ ParallelismMode, Dict[str, Dict[str, Optional[Dict[str, str]]]] ] = { ParallelismMode.TENSOR: { COMM_RULE_DEFAULT_KEY: { - 'forward': {'kind': 'all_reduce', 'participants': 'tp', 'interconnect': 'tp', 'suffix': 'all_reduce'}, - 'backward': {'kind': 'all_reduce', 'participants': 'tp', 'interconnect': 'tp', 'suffix': 'all_reduce'}, + 'forward': {'kind': 'all_reduce', 'participants': 'tp', 'interconnect': 'tp'}, + 'backward': {'kind': 'all_reduce', 'participants': 'tp', 'interconnect': 'tp'}, }, }, ParallelismMode.TENSOR_SEQUENCE: { - COMM_RULE_DEFAULT_KEY: {'forward': None, 'backward': None}, + COMM_RULE_DEFAULT_KEY: {'forward': None, 'backward': None}, # <- dangerous 'layernorm1': { - 'forward': {'kind': 'all_gather', 'participants': 'tp', 'interconnect': 'tp', 'suffix': 'all_gather'}, - 'backward': {'kind': 'all_gather', 'participants': 'tp', 'interconnect': 'tp', 'suffix': 'all_gather'}, + 'forward': {'kind': 'all_gather', 'participants': 'tp', 'interconnect': 'tp'}, + 'backward': {'kind': 'all_gather', 'participants': 'tp', 'interconnect': 'tp'}, }, 'layernorm2': { - 'forward': {'kind': 'all_gather', 'participants': 'tp', 'interconnect': 'tp', 'suffix': 'all_gather'}, - 'backward': {'kind': 'all_gather', 'participants': 'tp', 'interconnect': 'tp', 'suffix': 'all_gather'}, + 'forward': {'kind': 'all_gather', 'participants': 'tp', 'interconnect': 'tp'}, + 'backward': {'kind': 'all_gather', 'participants': 'tp', 'interconnect': 'tp'}, }, 'MHA': { - 'forward': {'kind': 'reduce_scatter', 'participants': 'tp', 'interconnect': 'tp', 'suffix': 'reduce_scatter'}, - 'backward': {'kind': 'reduce_scatter', 'participants': 'tp', 'interconnect': 'tp', 'suffix': 'reduce_scatter'}, + 'forward': {'kind': 'reduce_scatter', 'participants': 'tp', 'interconnect': 'tp'}, + 'backward': {'kind': 'reduce_scatter', 'participants': 'tp', 'interconnect': 'tp'}, }, 'MLP': { - 'forward': {'kind': 'reduce_scatter', 'participants': 'tp', 'interconnect': 'tp', 'suffix': 'reduce_scatter'}, - 'backward': {'kind': 'reduce_scatter', 'participants': 'tp', 'interconnect': 'tp', 'suffix': 'reduce_scatter'}, + 'forward': {'kind': 'reduce_scatter', 'participants': 'tp', 'interconnect': 'tp'}, + 'backward': {'kind': 'reduce_scatter', 'participants': 'tp', 'interconnect': 'tp'}, }, }, ParallelismMode.CONTEXT: { COMM_RULE_DEFAULT_KEY: { - 'forward': {'kind': 'all_gather', 'participants': 'cp', 'interconnect': 'cp', 'suffix': 'all_gather'}, - 'backward': {'kind': 'reduce_scatter', 'participants': 'cp', 'interconnect': 'cp', 'suffix': 'reduce_scatter'}, + 'forward': {'kind': 'all_gather', 'participants': 'cp', 'interconnect': 'cp'}, + 'backward': {'kind': 'reduce_scatter', 'participants': 'cp', 'interconnect': 'cp'}, }, - 'attention': {'forward': None, 'backward': {'kind': 'reduce_scatter', 'participants': 'cp', 'interconnect': 'cp', 'suffix': 'reduce_scatter'}}, - 'output_proj': {'forward': None, 'backward': {'kind': 'all_gather', 'participants': 'cp', 'interconnect': 'cp', 'suffix': 'all_gather'}}, - 'qkv_proj': {'forward': {'kind': 'all_gather', 'participants': 'cp', 'interconnect': 'cp', 'suffix': 'all_gather'}, 'backward': None}, + 'attention': {'forward': None, 'backward': {'kind': 'reduce_scatter', 'participants': 'cp', 'interconnect': 'cp'}}, + 'output_proj': {'forward': None, 'backward': {'kind': 'all_gather', 'participants': 'cp', 'interconnect': 'cp'}}, + 'qkv_proj': {'forward': {'kind': 'all_gather', 'participants': 'cp', 'interconnect': 'cp'}, 'backward': None}, }, ParallelismMode.TENSOR_CONTEXT_HYBRID: { COMM_RULE_DEFAULT_KEY: {'forward': None, 'backward': None}, 'layernorm1': { - 'forward': {'kind': 'all_gather', 'participants': 'tp', 'interconnect': 'tp', 'suffix': 'all_gather'}, - 'backward': {'kind': 'all_gather', 'participants': 'tp', 'interconnect': 'tp', 'suffix': 'all_gather'}, + 'forward': {'kind': 'all_gather', 'participants': 'tp', 'interconnect': 'tp'}, + 'backward': {'kind': 'all_gather', 'participants': 'tp', 'interconnect': 'tp'}, }, 'layernorm2': { - 'forward': {'kind': 'all_gather', 'participants': 'tp', 'interconnect': 'tp', 'suffix': 'all_gather'}, - 'backward': {'kind': 'all_gather', 'participants': 'tp', 'interconnect': 'tp', 'suffix': 'all_gather'}, + 'forward': {'kind': 'all_gather', 'participants': 'tp', 'interconnect': 'tp'}, + 'backward': {'kind': 'all_gather', 'participants': 'tp', 'interconnect': 'tp'}, }, 'MLP': { - 'forward': {'kind': 'reduce_scatter', 'participants': 'tp', 'interconnect': 'tp', 'suffix': 'reduce_scatter'}, - 'backward': {'kind': 'reduce_scatter', 'participants': 'tp', 'interconnect': 'tp', 'suffix': 'reduce_scatter'}, + 'forward': {'kind': 'reduce_scatter', 'participants': 'tp', 'interconnect': 'tp'}, + 'backward': {'kind': 'reduce_scatter', 'participants': 'tp', 'interconnect': 'tp'}, }, - 'attention': {'forward': None, 'backward': {'kind': 'reduce_scatter', 'participants': 'cp', 'interconnect': 'cp', 'suffix': 'reduce_scatter'}}, + 'attention': {'forward': None, 'backward': {'kind': 'reduce_scatter', 'participants': 'cp', 'interconnect': 'cp'}}, 'output_proj': { - 'forward': {'kind': 'reduce_scatter', 'participants': 'tp', 'interconnect': 'cp', 'suffix': 'reduce_scatter'}, - 'backward': {'kind': 'all_gather', 'participants': 'cp', 'interconnect': 'cp', 'suffix': 'all_gather'}, + 'forward': {'kind': 'reduce_scatter', 'participants': 'tp', 'interconnect': 'cp'}, + 'backward': {'kind': 'all_gather', 'participants': 'cp', 'interconnect': 'cp'}, }, 'qkv_proj': { - 'forward': {'kind': 'all_gather', 'participants': 'cp', 'interconnect': 'tp', 'suffix': 'all_gather'}, - 'backward': {'kind': 'reduce_scatter', 'participants': 'tp', 'interconnect': 'tp', 'suffix': 'reduce_scatter'}, + 'forward': {'kind': 'all_gather', 'participants': 'cp', 'interconnect': 'tp'}, + 'backward': {'kind': 'reduce_scatter', 'participants': 'tp', 'interconnect': 'tp'}, }, }, ParallelismMode.SINGLE: { @@ -117,13 +117,11 @@ class ParallelismMode(Enum): 'kind': 'all_to_all', 'participants': 'moe', 'interconnect': 'tp', - 'suffix': 'all_to_all', }, 'backward': { 'kind': 'reduce_scatter', 'participants': 'moe', 'interconnect': 'tp', - 'suffix': 'reduce_scatter', }, } _MOE_MLP_RULE = { @@ -131,13 +129,11 @@ class ParallelismMode(Enum): 'kind': 'all_to_all', 'participants': 'moe', 'interconnect': 'tp', - 'suffix': 'all_to_all', }, 'backward': { 'kind': 'all_to_all', 'participants': 'moe', 'interconnect': 'tp', - 'suffix': 'all_to_all', }, } _MOE_LAYER_NORM1_RULE = { @@ -145,7 +141,6 @@ class ParallelismMode(Enum): 'kind': 'all_to_all', 'participants': 'moe', 'interconnect': 'tp', - 'suffix': 'all_to_all', }, } MOE_COMMUNICATION_RULES: Dict[ @@ -1431,12 +1426,7 @@ def get_data_parallel_reduction_sizes(self, d, intermediate_size): """Calculate communication sizes for data parallel reductions (no timing).""" if not getattr(self, "dp", 1) or self.dp <= 1: # No communication needed for dp=1 - return { - 'qkv_size': 0, - 'output_size': 0, - 'ffn_size': 0, - 'total_size': 0 - } + return 0 # Calculate sizes only qkv_size = math.ceil(self.precision.grad_communication * d * 3 * d) @@ -1446,154 +1436,21 @@ def get_data_parallel_reduction_sizes(self, d, intermediate_size): ffn2_size = math.ceil(self.precision.grad_communication * intermediate_size * d) total_size = qkv_size + output_size + ffn1_size + ffn2_size - return { - 'qkv_size': qkv_size, - 'output_size': output_size, - 'ffn_size': ffn1_size + ffn2_size, - 'total_size': total_size - } + return total_size - def get_data_parallel_local_computation(self, d, intermediate_size): - """Calculate local computation times for apply_grad operations.""" - qkv_local = self.apply_grad(Dim0=d, Dim1=3*d, name="qkv_proj reduction") - output_local = self.apply_grad(Dim0=d, Dim1=d, name="output_proj reduction") + def get_data_parallel_reduction_llm(self, d, intermediate_size): + """Return apply_grad compute time per rank (no communication), honoring ZeRO sharding.""" + apply_grad_time = self.apply_grad(int(d * 3 * d)) # QKV + apply_grad_time += self.apply_grad(int(d * d)) # Output ffn1_dim = self._ffn1_output_dim(intermediate_size) - ffn_local = ( - self.apply_grad(Dim0=ffn1_dim, Dim1=d, name="ffn1 reduction") - + self.apply_grad(Dim0=intermediate_size, Dim1=d, name="ffn2 reduction") - ) + apply_grad_time += self.apply_grad(int(ffn1_dim * d)) # FFN1 + apply_grad_time += self.apply_grad(int(intermediate_size * d)) # FFN2 - return { - 'qkv_local': qkv_local, - 'output_local': output_local, - 'ffn_local': ffn_local, - 'total_local': qkv_local + output_local + ffn_local - } - - def get_data_parallel_reduction_llm(self, d, intermediate_size): - # If no data parallelism, still apply gradients locally but no cross-device reduction - if self.dp <= 1: - apply_grad_time = 0.0 - apply_grad_time += self.apply_grad(Dim0=d, Dim1=3*d, name="qkv_proj reduction") - apply_grad_time += self.apply_grad(Dim0=d, Dim1=d, name="output_proj reduction") - ffn1_dim = self._ffn1_output_dim(intermediate_size) - apply_grad_time += self.apply_grad(Dim0=ffn1_dim, Dim1=d, name="ffn1 reduction") - apply_grad_time += self.apply_grad(Dim0=intermediate_size, Dim1=d, name="ffn2 reduction") - if self.debug: - print(f"(dp=1) apply_grad_time: {apply_grad_time}") - return apply_grad_time - - reduction_sizes = self.get_data_parallel_reduction_sizes(d, intermediate_size) - grad_bytes = int(reduction_sizes['total_size']) - ( - total_params_per_rank, - _, - _, - _, - _, - ) = self._param_stats_per_rank(d, intermediate_size, self.vocab_size) - param_bytes = int(math.ceil(total_params_per_rank * self.precision.parameters)) - grad_shard = self.dp if self.zero_stage >= 2 else 1 - - if grad_shard > 1: # ZeRO stages >= 2: gradients are sharded via reduce-scatter and parameters are gathered. - apply_grad_time = 0.0 - apply_grad_time += self.apply_grad(Dim0=d, Dim1=3 * d, name="qkv_proj reduction") - apply_grad_time += self.apply_grad(Dim0=d, Dim1=d, name="output_proj reduction") - ffn1_dim = self._ffn1_output_dim(intermediate_size) - apply_grad_time += self.apply_grad(Dim0=ffn1_dim, Dim1=d, name="ffn1 reduction") - apply_grad_time += self.apply_grad(Dim0=intermediate_size, Dim1=d, name="ffn2 reduction") - # divide by dp to get the time per device + grad_shard = self.dp if (self.zero_stage >= 2 and self.dp > 1) else 1 + if grad_shard > 1: apply_grad_time /= grad_shard - reduction_time = self.network_model.collective( - kind="reduce_scatter", - size_bytes=grad_bytes, - participants=int(self.dp), - ib=self.IBD, - ll=self.LLD, - local_bytes=0.0, - debug_label="dp_zero_gradients", - axis="dp", - ) - - gather_time = 0.0 - if self.zero_stage == 2: - gather_bytes = int(math.ceil(param_bytes)) - gather_time += self.network_model.collective( - kind="all_gather", - size_bytes=gather_bytes, - participants=int(self.dp), - ib=self.IBD, - ll=self.LLD, - local_bytes=0.0, - debug_label="dp_zero_param_gather_zero2", - axis="dp", - ) - if self.debug: - print(f"apply_grad_time (ZeRO-{self.zero_stage}): {apply_grad_time}") - # TODO(ZeRO): revisit optimizer FLOPs once ZeRO-3 materialization is finalized. - return reduction_time + gather_time + apply_grad_time - else: # legacy path for zero stages <= 2. - reduction_time = 0.0 - apply_grad_time = 0.0 - - total_bytes = math.ceil(self.precision.grad_communication * d * 3 * d) - reduction_time += self.network_model.collective( - kind="all_reduce", - size_bytes=total_bytes, - participants=int(self.dp), - ib=self.IBD, - ll=self.LLD, - local_bytes=0.0, - debug_label="qkv_proj reduction", - axis="dp", - ) - apply_grad_time += self.apply_grad(Dim0=d, Dim1=3*d, name="qkv_proj reduction") - - total_bytes = math.ceil(self.precision.grad_communication * d * d) - reduction_time += self.network_model.collective( - kind="all_reduce", - size_bytes=total_bytes, - participants=int(self.dp), - ib=self.IBD, - ll=self.LLD, - local_bytes=0.0, - debug_label="output_proj reduction", - axis="dp", - ) - apply_grad_time += self.apply_grad(Dim0=d, Dim1=d, name="output_proj reduction") - - ffn1_dim = self._ffn1_output_dim(intermediate_size) - total_bytes = math.ceil(self.precision.grad_communication * ffn1_dim * d) - reduction_time += self.network_model.collective( - kind="all_reduce", - size_bytes=total_bytes, - participants=int(self.dp), - ib=self.IBD, - ll=self.LLD, - local_bytes=0.0, - debug_label="ffn1 reduction", - axis="dp", - ) - apply_grad_time += self.apply_grad(Dim0=ffn1_dim, Dim1=d, name="ffn1 reduction") - - total_bytes = math.ceil(self.precision.grad_communication * intermediate_size * d) - reduction_time += self.network_model.collective( - kind="all_reduce", - size_bytes=total_bytes, - participants=int(self.dp), - ib=self.IBD, - ll=self.LLD, - local_bytes=0.0, - debug_label="ffn2 reduction", - axis="dp", - ) - apply_grad_time += self.apply_grad(Dim0=intermediate_size, Dim1=d, name="ffn2 reduction") - - if self.debug: - print(f"apply_grad_time: {apply_grad_time}") - - return reduction_time + apply_grad_time + return apply_grad_time def _combine_mem(self, *args): combined = {} @@ -1618,7 +1475,18 @@ def _mem_levels(self, arr): # TODO TODO: # we need a significant refactor here. The comm sizes are ingested in a weird way and never used. Instead we use old precomputed sizes. # FIX at some point! - def compute_all_gemm_and_node_times(self, batch_size, vocab_size, hidden_dim, seq_len, num_heads, kv_heads, intermediate_size, num_SMs): + def compute_all_gemm_and_node_times( + self, + batch_size, + vocab_size, + hidden_dim, + seq_len, + num_heads, + kv_heads, + intermediate_size, + num_SMs, + include_optimizer: bool = False, + ): """Compute latency for all GEMM operations and node breakdown times.""" gemm_shapes = LLM_util.process_gemm_shapes( @@ -1631,21 +1499,6 @@ def compute_all_gemm_and_node_times(self, batch_size, vocab_size, hidden_dim, se intermediate_size, vocab_size, ) - if self.debug: - print( - "generating gemm shapes for transformer batch size:", - batch_size, - "seq_len:", - seq_len, - "hidden_dim:", - hidden_dim, - "num_heads:", - num_heads, - "kv_heads:", - kv_heads, - "intermediate_size:", - intermediate_size, - ) gemm_qkv_proj = gemm_shapes["qkv_proj"] gemm_attention_score = gemm_shapes["attention_score"] @@ -1658,26 +1511,6 @@ def compute_all_gemm_and_node_times(self, batch_size, vocab_size, hidden_dim, se transformer_timings: Dict[str, OperationTiming] = {} - def _make_direction( - op_name: str, - direction: str, - compute_time: float, - comm_time: float, - comm_bytes: float, - *, - flops: float = 0.0, - memory: Optional[Mapping[str, float]] = None, - ) -> DirectionTiming: - mem_map = dict(memory) if memory else {} - bytes_int = int(math.ceil(float(comm_bytes or 0.0))) - return DirectionTiming( - compute_time=compute_time, - comm_time=comm_time, - comm_bytes=bytes_int, - flops=flops, - memory_accesses=mem_map, - ) - def _extract_forward(ret: Sequence[Any]) -> Tuple[float, float, float, float, Any]: if len(ret) == 5: time, reduction, size, flops, mem = ret @@ -1694,21 +1527,17 @@ def _extract_forward(ret: Sequence[Any]) -> Tuple[float, float, float, float, An qkv_proj_gemm_b, qkv_proj_reduction_b, qkv_proj_size_b = self.parallelism_gemm_backward( gemm_qkv_proj, "qkv_projection_b", gemm_type=GemmType.QKV ) - qkv_forward = _make_direction( - "qkv_proj", - "forward", + qkv_forward = DirectionTiming( compute_time=qkv_proj_gemm_f, comm_time=qkv_proj_reduction_f, - comm_bytes=qkv_proj_size_f, + comm_bytes=int(qkv_proj_size_f or 0), flops=qkv_proj_flops_f, - memory=self._mem_levels(qkv_proj_mem_f), + memory_accesses=dict(self._mem_levels(qkv_proj_mem_f)), ) - qkv_backward = _make_direction( - "qkv_proj", - "backward", + qkv_backward = DirectionTiming( compute_time=qkv_proj_gemm_b, comm_time=qkv_proj_reduction_b, - comm_bytes=qkv_proj_size_b, + comm_bytes=int(qkv_proj_size_b or 0), ) transformer_timings["qkv_proj"] = OperationTiming("qkv_proj", forward=qkv_forward, backward=qkv_backward) @@ -1719,21 +1548,17 @@ def _extract_forward(ret: Sequence[Any]) -> Tuple[float, float, float, float, An attn_score_gemm_b, attn_score_reduction_b, attn_score_size_b = self.parallelism_gemm_backward( gemm_attention_score, "attention_score_b", gemm_type=GemmType.ATTENTION_SCORE ) - attn_score_forward = _make_direction( - "attention_score", - "forward", + attn_score_forward = DirectionTiming( compute_time=attn_score_gemm_f, comm_time=attn_score_reduction_f, - comm_bytes=attn_score_size_f, + comm_bytes=int(attn_score_size_f or 0), flops=attn_score_flops_f, - memory=self._mem_levels(attn_score_mem_f), + memory_accesses=dict(self._mem_levels(attn_score_mem_f)), ) - attn_score_backward = _make_direction( - "attention_score", - "backward", + attn_score_backward = DirectionTiming( compute_time=attn_score_gemm_b, comm_time=attn_score_reduction_b, - comm_bytes=attn_score_size_b, + comm_bytes=int(attn_score_size_b or 0), ) transformer_timings["attention_score"] = OperationTiming( "attention_score", @@ -1747,19 +1572,15 @@ def _extract_forward(ret: Sequence[Any]) -> Tuple[float, float, float, float, An transformer_timings["attention_scale_softmax"] = OperationTiming( name="attention_scale_softmax", - forward=_make_direction( - "attention_scale_softmax", - "forward", + forward=DirectionTiming( compute_time=attention_scale_softmax_f, comm_time=0.0, - comm_bytes=0.0, + comm_bytes=0, ), - backward=_make_direction( - "attention_scale_softmax", - "backward", + backward=DirectionTiming( compute_time=attention_scale_softmax_b, comm_time=0.0, - comm_bytes=0.0, + comm_bytes=0, ), ) @@ -1769,21 +1590,17 @@ def _extract_forward(ret: Sequence[Any]) -> Tuple[float, float, float, float, An attn_out_gemm_b, attn_out_reduction_b, attn_out_size_b = self.parallelism_gemm_backward( gemm_attention_output, "attention_output_b", gemm_type=GemmType.ATTENTION_OUTPUT ) - attn_out_forward = _make_direction( - "attention_output", - "forward", + attn_out_forward = DirectionTiming( compute_time=attn_out_gemm_f, comm_time=attn_out_reduction_f, - comm_bytes=attn_out_size_f, + comm_bytes=int(attn_out_size_f or 0), flops=attn_out_flops_f, - memory=self._mem_levels(attn_out_mem_f), + memory_accesses=dict(self._mem_levels(attn_out_mem_f)), ) - attn_out_backward = _make_direction( - "attention_output", - "backward", + attn_out_backward = DirectionTiming( compute_time=attn_out_gemm_b, comm_time=attn_out_reduction_b, - comm_bytes=attn_out_size_b, + comm_bytes=int(attn_out_size_b or 0), ) transformer_timings["attention_output"] = OperationTiming( "attention_output", @@ -1808,21 +1625,17 @@ def _extract_forward(ret: Sequence[Any]) -> Tuple[float, float, float, float, An attention_mem = self._combine_mem(attn_score_mem_f, attn_out_mem_f) transformer_timings["attention"] = OperationTiming( "attention", - forward=_make_direction( - "attention", - "forward", + forward=DirectionTiming( compute_time=attention_forward_compute, comm_time=attention_forward_comm, - comm_bytes=attention_comm_bytes_f, - memory=self._mem_levels(attention_mem), + comm_bytes=int(attention_comm_bytes_f or 0), + memory_accesses=dict(self._mem_levels(attention_mem)), flops=(attn_score_flops_f + attn_out_flops_f), ), - backward=_make_direction( - "attention", - "backward", + backward=DirectionTiming( compute_time=attention_backward_compute, comm_time=attention_backward_comm, - comm_bytes=attention_comm_bytes_b, + comm_bytes=int(attention_comm_bytes_b or 0), ), ) else: @@ -1845,20 +1658,16 @@ def _extract_forward(ret: Sequence[Any]) -> Tuple[float, float, float, float, An transformer_timings["attention"] = OperationTiming( "attention", - forward=_make_direction( - "attention", - "forward", + forward=DirectionTiming( compute_time=attention_gemm_f, comm_time=attention_reduction_f, - comm_bytes=attention_size_f, - memory=self._mem_levels(attention_mem), + comm_bytes=int(attention_size_f or 0), + memory_accesses=dict(self._mem_levels(attention_mem)), ), - backward=_make_direction( - "attention", - "backward", + backward=DirectionTiming( compute_time=attention_gemm_b, comm_time=attention_reduction_b, - comm_bytes=attention_size_b, + comm_bytes=int(attention_size_b or 0), ), ) @@ -1870,21 +1679,17 @@ def _extract_forward(ret: Sequence[Any]) -> Tuple[float, float, float, float, An ) transformer_timings["output_proj"] = OperationTiming( "output_proj", - forward=_make_direction( - "output_proj", - "forward", + forward=DirectionTiming( compute_time=out_proj_gemm_f, comm_time=out_proj_reduction_f, - comm_bytes=out_proj_size_f, + comm_bytes=int(out_proj_size_f or 0), flops=out_proj_flops_f, - memory=self._mem_levels(out_proj_mem_f), + memory_accesses=dict(self._mem_levels(out_proj_mem_f)), ), - backward=_make_direction( - "output_proj", - "backward", + backward=DirectionTiming( compute_time=out_proj_gemm_b, comm_time=out_proj_reduction_b, - comm_bytes=out_proj_size_b, + comm_bytes=int(out_proj_size_b or 0), ), ) if not self.use_moe: @@ -1915,60 +1720,48 @@ def _extract_forward(ret: Sequence[Any]) -> Tuple[float, float, float, float, An if self.use_moe: transformer_timings["router"] = OperationTiming( "router", - forward=_make_direction( - "router", - "forward", + forward=DirectionTiming( compute_time=router_gemm_time_f, comm_time=router_reduction_f, - comm_bytes=router_size_f, + comm_bytes=int(router_size_f or 0), ), - backward=_make_direction( - "router", - "backward", + backward=DirectionTiming( compute_time=router_gemm_time_b, comm_time=router_reduction_b, - comm_bytes=router_size_b, + comm_bytes=int(router_size_b or 0), ), ) transformer_timings["ffn1"] = OperationTiming( "ffn1", - forward=_make_direction( - "ffn1", - "forward", + forward=DirectionTiming( compute_time=ffn1_gemm_f, comm_time=ffn1_reduction_f, - comm_bytes=ffn1_size_f, + comm_bytes=int(ffn1_size_f or 0), flops=ffn1_flops_f, - memory=self._mem_levels(ffn1_mem_f), + memory_accesses=dict(self._mem_levels(ffn1_mem_f)), ), - backward=_make_direction( - "ffn1", - "backward", + backward=DirectionTiming( compute_time=ffn1_gemm_b, comm_time=ffn1_reduction_b, - comm_bytes=ffn1_size_b, + comm_bytes=int(ffn1_size_b or 0), ), ) transformer_timings["ffn2"] = OperationTiming( "ffn2", - forward=_make_direction( - "ffn2", - "forward", + forward=DirectionTiming( compute_time=ffn2_gemm_f, comm_time=ffn2_reduction_f, - comm_bytes=ffn2_size_f, + comm_bytes=int(ffn2_size_f or 0), flops=ffn2_flops_f, - memory=self._mem_levels(ffn2_mem_f), + memory_accesses=dict(self._mem_levels(ffn2_mem_f)), ), - backward=_make_direction( - "ffn2", - "backward", + backward=DirectionTiming( compute_time=ffn2_gemm_b, comm_time=ffn2_reduction_b, - comm_bytes=ffn2_size_b, + comm_bytes=int(ffn2_size_b or 0), ), ) @@ -1976,20 +1769,16 @@ def _extract_forward(ret: Sequence[Any]) -> Tuple[float, float, float, float, An embedding_b = self.get_embedding_b(vocab_size=vocab_size, seq_len=seq_len, hidden_dim=hidden_dim) transformer_timings["embedding"] = OperationTiming( "embedding", - forward=_make_direction( - "embedding", - "forward", + forward=DirectionTiming( compute_time=embedding_f, comm_time=0.0, - comm_bytes=0.0, - memory=self._mem_levels(embedding_mem), + comm_bytes=0, + memory_accesses=dict(self._mem_levels(embedding_mem)), ), - backward=_make_direction( - "embedding", - "backward", + backward=DirectionTiming( compute_time=embedding_b, comm_time=0.0, - comm_bytes=0.0, + comm_bytes=0, ), ) @@ -2008,19 +1797,15 @@ def _extract_forward(ret: Sequence[Any]) -> Tuple[float, float, float, float, An ) transformer_timings["layernorm1"] = OperationTiming( "layernorm1", - forward=_make_direction( - "layernorm1", - "forward", + forward=DirectionTiming( compute_time=layernorm1_f + residual1_f, comm_time=layernorm1_reduction_f, - comm_bytes=LN1_comm_bytes_f, + comm_bytes=int(LN1_comm_bytes_f or 0), ), - backward=_make_direction( - "layernorm1", - "backward", + backward=DirectionTiming( compute_time=layernorm1_b + residual1_b, comm_time=layernorm1_reduction_b, - comm_bytes=LN1_comm_bytes_b, + comm_bytes=int(LN1_comm_bytes_b or 0), ), ) @@ -2038,19 +1823,15 @@ def _extract_forward(ret: Sequence[Any]) -> Tuple[float, float, float, float, An ) transformer_timings["layernorm2"] = OperationTiming( "layernorm2", - forward=_make_direction( - "layernorm2", - "forward", + forward=DirectionTiming( compute_time=layernorm2_f + residual2_f, comm_time=layernorm2_reduction_f, - comm_bytes=LN2_comm_bytes_f, + comm_bytes=int(LN2_comm_bytes_f or 0), ), - backward=_make_direction( - "layernorm2", - "backward", + backward=DirectionTiming( compute_time=layernorm2_b + residual2_b, comm_time=layernorm2_reduction_b, - comm_bytes=LN2_comm_bytes_b, + comm_bytes=int(LN2_comm_bytes_b or 0), ), ) @@ -2062,28 +1843,24 @@ def _extract_forward(ret: Sequence[Any]) -> Tuple[float, float, float, float, An act_b = self.get_gelu_b(tensor_shape=gemm_ffn1) transformer_timings["gelu"] = OperationTiming( "gelu", - forward=_make_direction("gelu", "forward", compute_time=act_f, comm_time=0.0, comm_bytes=0.0), - backward=_make_direction("gelu", "backward", compute_time=act_b, comm_time=0.0, comm_bytes=0.0), + forward=DirectionTiming(compute_time=act_f, comm_time=0.0, comm_bytes=0), + backward=DirectionTiming(compute_time=act_b, comm_time=0.0, comm_bytes=0), ) linear_softmax_f, linear_softmax_mem = self.get_linear_softmax_f(gemm=gemm_linear) linear_softmax_b = self.get_linear_softmax_b(gemm=gemm_linear) transformer_timings["linear_softmax"] = OperationTiming( "linear_softmax", - forward=_make_direction( - "linear_softmax", - "forward", + forward=DirectionTiming( compute_time=linear_softmax_f, comm_time=0.0, - comm_bytes=0.0, - memory=self._mem_levels(linear_softmax_mem), + comm_bytes=0, + memory_accesses=dict(self._mem_levels(linear_softmax_mem)), ), - backward=_make_direction( - "linear_softmax", - "backward", + backward=DirectionTiming( compute_time=linear_softmax_b, comm_time=0.0, - comm_bytes=0.0, + comm_bytes=0, ), ) @@ -2117,6 +1894,20 @@ def _extract_forward(ret: Sequence[Any]) -> Tuple[float, float, float, float, An + transformer_timings["layernorm2"].total_backward_time() ) + if include_optimizer: + optimizer_time = self.get_data_parallel_reduction_llm(hidden_dim, intermediate_size) + if optimizer_time > 0.0: + transformer_timings["optimizer"] = OperationTiming( + "optimizer", + forward=None, + backward=DirectionTiming( + compute_time=optimizer_time, + comm_time=0.0, + comm_bytes=0, + ), + ) + transformer_time_b += optimizer_time + node_breakdown = { "transformer_time_f": transformer_time_f, "transformer_time_b": transformer_time_b, @@ -2125,11 +1916,10 @@ def _extract_forward(ret: Sequence[Any]) -> Tuple[float, float, float, float, An "linear_softmax_f": transformer_timings["linear_softmax"].total_forward_time(), "linear_softmax_b": transformer_timings["linear_softmax"].total_backward_time(), } - if self._generate_graphs: - results_path = os.path.join(self.output_dir, "transformer_results.txt") + results_path = os.path.join(self.output_dir, "transformer_timings.yaml") with open(results_path, "w", encoding="utf-8") as results_file: - json.dump( + yaml.dump( { "transformer_results": { name: timing.to_dict() for name, timing in transformer_timings.items() @@ -2137,7 +1927,6 @@ def _extract_forward(ret: Sequence[Any]) -> Tuple[float, float, float, float, An "node_breakdown": node_breakdown, }, results_file, - indent=2, sort_keys=True, ) @@ -2170,11 +1959,11 @@ def _build_comm_metadata( grad_collective = 'reduce_scatter' if (self.zero_stage >= 2 and self.dp > 1) else 'all_reduce' metadata = { 'transformer': { - 'size': reduction_sizes['total_size'], + 'size': reduction_sizes, 'type': grad_collective, 'participants': self.dp, 'interconnect_type': 'dp', - 'local_comp_time': local_comp['total_local'] + 'local_comp_time': local_comp }, 'embedding': { 'size': embedding_size, @@ -2290,21 +2079,11 @@ def _prepare_execution_graphs( if not include_pipeline_backward and not include_transformer_backward: # Forward-only inference: skip training all-reduce bookkeeping. - reduction_sizes = { - 'qkv_size': 0, - 'output_size': 0, - 'ffn_size': 0, - 'total_size': 0, - } - local_comp = { - 'qkv_local': 0.0, - 'output_local': 0.0, - 'ffn_local': 0.0, - 'total_local': 0.0, - } + reduction_sizes = 0.0 + local_comp = 0.0 else: reduction_sizes = self.get_data_parallel_reduction_sizes(hidden_dim, intermediate_size) - local_comp = self.get_data_parallel_local_computation(hidden_dim, intermediate_size) + local_comp = self.get_data_parallel_reduction_llm(hidden_dim, intermediate_size) # these are used for dp all-reduce/reduce-scatter. embedding_size = math.ceil(self.precision.grad_communication * vocab_size * hidden_dim) + math.ceil(self.precision.grad_communication * seq_len * hidden_dim * batch_size) @@ -2404,7 +2183,10 @@ def _make_comm_specs(op_name: str, direction: str, total_bytes: float) -> Tuple[ return () rule_spec = rules_by_mode.get(op_name) or rules_by_mode.get(COMM_RULE_DEFAULT_KEY) if not rule_spec: - return () + raise ValueError( + f"Missing communication rule for op '{op_name}' in mode '{parallelism_mode.value}' " + f"with non-zero comm_bytes ({bytes_int})" + ) rule = rule_spec.get(direction) if not rule: return () @@ -2416,8 +2198,7 @@ def _make_comm_specs(op_name: str, direction: str, total_bytes: float) -> Tuple[ if participants <= 1: return () interconnect = rule.get("interconnect", scope) - suffix = rule.get("suffix", kind) - name = f"{op_name}_{direction}_{suffix}" + name = f"{op_name}_{direction}_{kind}" return ( CommSpec( name=name, @@ -2474,6 +2255,8 @@ def _resolve_operation(name: str) -> OperationTiming: else: raise ValueError(f"Unsupported parallelism mode: {parallelism_mode}") op_names = list(op_names) + if "optimizer" in transformer_timings: + op_names.append("optimizer") if self.use_moe and "MLP" in op_names and "router" not in op_names: # if MOE is enabled, ensure router is included in the graph mlp_index = op_names.index("MLP") op_names.insert(mlp_index, "router") @@ -2526,6 +2309,7 @@ def _resolve_operation(name: str) -> OperationTiming: "linear_softmax_b": node_breakdown.get('linear_softmax_b', 0.0) if include_pipeline_backward else 0.0, "transformer_f": node_breakdown.get('transformer_time_f', 0.0), "transformer_b": node_breakdown.get('transformer_time_b', 0.0) if include_pipeline_backward else 0.0, + "optimizer": self.get_data_parallel_reduction_llm(hidden_dim, intermediate_size), "cross_layer_f": 0.0, "cross_layer_b": 0.0, } @@ -2547,7 +2331,10 @@ def _resolve_operation(name: str) -> OperationTiming: misc_metadata=misc_metadata, ) - graph_root = pipeline_graph_obj.construct_fwd_bwd_graph(include_backward=include_pipeline_backward) + graph_root = pipeline_graph_obj.construct_fwd_bwd_graph( + include_backward=include_pipeline_backward, + include_optimizer=True + ) pipeline_graph_obj_no_dp = None graph_root_no_dp = None need_no_dp_variant = getattr(self, "gradient_accumulation_steps", 1) > 1 @@ -2563,7 +2350,10 @@ def _resolve_operation(name: str) -> OperationTiming: misc_metadata=misc_metadata, ) - graph_root_no_dp = pipeline_graph_obj_no_dp.construct_fwd_bwd_graph(include_backward=include_pipeline_backward) + graph_root_no_dp = pipeline_graph_obj_no_dp.construct_fwd_bwd_graph( + include_backward=include_pipeline_backward, + include_optimizer=False + ) interconnect_params = self._build_interconnect_params() @@ -2636,7 +2426,15 @@ def calc_time_llm(self): num_SMs = self.hw_config.tech_config.core.num_bundles transformer_timings, node_breakdown = self.compute_all_gemm_and_node_times( - batch_size, vocab_size, hidden_dim, seq_len, num_heads, kv_heads, intermediate_size, num_SMs + batch_size, + vocab_size, + hidden_dim, + seq_len, + num_heads, + kv_heads, + intermediate_size, + num_SMs, + include_optimizer=False, ) # transformer mem layer considers zero stage internally diff --git a/timing_model.py b/timing_model.py index d9bb35f5..b0a1ea28 100644 --- a/timing_model.py +++ b/timing_model.py @@ -101,21 +101,23 @@ class OperationTiming: Attributes: name: Logical operation name. - forward: Forward-direction timing (required). + forward: Forward-direction timing (optional). backward: Backward-direction timing (optional). metadata: Additional operation-level metadata (e.g. layer index). """ name: str - forward: DirectionTiming + forward: Optional[DirectionTiming] backward: Optional[DirectionTiming] = None metadata: Mapping[str, Any] = field(default_factory=dict) def __post_init__(self) -> None: if not self.name: raise ValueError("OperationTiming.name must be non-empty") - if not isinstance(self.forward, DirectionTiming): - raise TypeError("OperationTiming.forward must be a DirectionTiming instance") + if self.forward is None and self.backward is None: + raise ValueError(f"OperationTiming '{self.name}' must have at least one direction") + if self.forward is not None and not isinstance(self.forward, DirectionTiming): + raise TypeError("OperationTiming.forward must be a DirectionTiming instance or None") if self.backward is not None and not isinstance(self.backward, DirectionTiming): raise TypeError("OperationTiming.backward must be a DirectionTiming instance or None") @@ -123,6 +125,8 @@ def has_backward(self) -> bool: return self.backward is not None def total_forward_time(self) -> float: + if self.forward is None: + raise RuntimeError(f"Operation '{self.name}' does not have forward timing") return self.forward.total_time() def total_backward_time(self) -> float: @@ -131,7 +135,7 @@ def total_backward_time(self) -> float: return self.backward.total_time() def total_time(self) -> float: - total = self.forward.total_time() + total = self.forward.total_time() if self.forward is not None else 0.0 if self.backward is not None: total += self.backward.total_time() return total @@ -140,8 +144,9 @@ def to_dict(self) -> Dict[str, Any]: """Serialize to dict for reporting or artifact dumps.""" data: Dict[str, Any] = { "name": self.name, - "forward": self.forward.to_dict(), } + if self.forward is not None: + data["forward"] = self.forward.to_dict() if self.backward is not None: data["backward"] = self.backward.to_dict() if self.metadata: diff --git a/validation_scripts/validation_configs/hardware-config/a100_80GB_korthikanti.yaml b/validation_scripts/validation_configs/hardware-config/a100_80GB_korthikanti.yaml index 61500997..a23ccfee 100644 --- a/validation_scripts/validation_configs/hardware-config/a100_80GB_korthikanti.yaml +++ b/validation_scripts/validation_configs/hardware-config/a100_80GB_korthikanti.yaml @@ -19,7 +19,7 @@ sw_param: # -> tensor_plus_fp32_master makes weights/gradients use tensor_format bytes and # keeps a separate FP32 master copy for weights. # -> fp32_params makes weights *and* gradients use FP32 bytes without a master copy. - param_storage_mode: tensor_plus_fp32_master + param_storage_mode: as_tensor_format # kv_cache precision may be set as bytes, dtype string, or "as_tensor_format". kv_cache: as_tensor_format h2d_bandwidth: -1 # Bytes/s; set -1 to disable H2D timing @@ -74,13 +74,13 @@ memory_hierarchy: scope: global parallelism: - dp: 2 # data parallelism. Acts as replica count during inference - tp: 2 # tensor parallelism + dp: 1 # data parallelism. Acts as replica count during inference + tp: 8 # tensor parallelism tp_sp: true # tensor-sequence parallelism - cp: 2 # context parallelism. Not supported for inference. - lp: 4 # pipeline parallelism. Number of pipeline stages + cp: 1 # context parallelism. Not supported for inference. + lp: 1 # pipeline parallelism. Number of pipeline stages # mb is not a true parallelism parameter, does not affect num of GPUs. - mb: 4 # micro-batches (for best results set to >= lp). + mb: 1 # micro-batches (for best results set to >= lp). network: # Global faulty links apply to the active AstraSim dimension: [[src, dst, weight]] @@ -126,7 +126,7 @@ execution_backend: model: astra # analytical | astra astra: backend: astra - mode: full_astrasim_hierarchical # hybrid | full_astrasim_hierarchical | full_astrasim_flattened + mode: full_astrasim_flattened # hybrid | full_astrasim_hierarchical | full_astrasim_flattened collectives: # Default collective algorithms. Per-dimension collective_override entries can override these. # Allowed values per op: auto | ring | direct | halvingDoubling | doubleBinaryTree | mesh | mesh2d | hypercube | torus2d diff --git a/validation_scripts/validation_configs/model-config/GPT_22_B.yaml b/validation_scripts/validation_configs/model-config/GPT_22_B.yaml index 67598d95..bc1e43b6 100644 --- a/validation_scripts/validation_configs/model-config/GPT_22_B.yaml +++ b/validation_scripts/validation_configs/model-config/GPT_22_B.yaml @@ -3,9 +3,9 @@ model_param: run_type: "training" tied_embeddings: false model_type: "gpt" - global_batch_size: 8 + global_batch_size: 4 gradient_accumulation_steps: 1 - recomputation: partial + recomputation: selective seq_len: 2048 hidden_dim: 6144 attention: From 0d065cd9fe32f2398727e7ed5175ee708bfe7f3a Mon Sep 17 00:00:00 2001 From: gkarfakis19 Date: Sat, 22 Nov 2025 18:26:46 -0800 Subject: [PATCH 55/77] minor fix --- time_calculation.py | 2 +- .../hardware-config/a100_80GB_korthikanti.yaml | 2 +- 2 files changed, 2 insertions(+), 2 deletions(-) diff --git a/time_calculation.py b/time_calculation.py index 5c22d99b..29f715b5 100644 --- a/time_calculation.py +++ b/time_calculation.py @@ -1499,7 +1499,7 @@ def apply_grad(self, num_params: int) -> float: apply_grad_comp = num_params * OPT_FLOPS_PER_PARAM apply_grad_mem = num_params * (bytes_per_param_read + bytes_per_param_write) - apply_grad_time = self.roofline(apply_grad_comp, apply_grad_mem, name="pointwise-applyGrad") + apply_grad_time = self.roofline(apply_grad_comp, apply_grad_mem, name="apply_grad", mem_level=self.num_levels-1) clip_time = self.grad_clipping(num_params) return apply_grad_time + clip_time diff --git a/validation_scripts/validation_configs/hardware-config/a100_80GB_korthikanti.yaml b/validation_scripts/validation_configs/hardware-config/a100_80GB_korthikanti.yaml index a23ccfee..a9709e7a 100644 --- a/validation_scripts/validation_configs/hardware-config/a100_80GB_korthikanti.yaml +++ b/validation_scripts/validation_configs/hardware-config/a100_80GB_korthikanti.yaml @@ -19,7 +19,7 @@ sw_param: # -> tensor_plus_fp32_master makes weights/gradients use tensor_format bytes and # keeps a separate FP32 master copy for weights. # -> fp32_params makes weights *and* gradients use FP32 bytes without a master copy. - param_storage_mode: as_tensor_format + param_storage_mode: tensor_plus_fp32_master # kv_cache precision may be set as bytes, dtype string, or "as_tensor_format". kv_cache: as_tensor_format h2d_bandwidth: -1 # Bytes/s; set -1 to disable H2D timing From e5e023cfd766db76bc2dd552018caebf94c5df61 Mon Sep 17 00:00:00 2001 From: gkarfakis19 Date: Sat, 22 Nov 2025 19:19:46 -0800 Subject: [PATCH 56/77] removed vestigial code --- time_calculation_LLM.py | 16 ---------------- 1 file changed, 16 deletions(-) diff --git a/time_calculation_LLM.py b/time_calculation_LLM.py index 1f9ec742..52bb9779 100644 --- a/time_calculation_LLM.py +++ b/time_calculation_LLM.py @@ -1485,7 +1485,6 @@ def compute_all_gemm_and_node_times( kv_heads, intermediate_size, num_SMs, - include_optimizer: bool = False, ): """Compute latency for all GEMM operations and node breakdown times.""" @@ -1894,20 +1893,6 @@ def _extract_forward(ret: Sequence[Any]) -> Tuple[float, float, float, float, An + transformer_timings["layernorm2"].total_backward_time() ) - if include_optimizer: - optimizer_time = self.get_data_parallel_reduction_llm(hidden_dim, intermediate_size) - if optimizer_time > 0.0: - transformer_timings["optimizer"] = OperationTiming( - "optimizer", - forward=None, - backward=DirectionTiming( - compute_time=optimizer_time, - comm_time=0.0, - comm_bytes=0, - ), - ) - transformer_time_b += optimizer_time - node_breakdown = { "transformer_time_f": transformer_time_f, "transformer_time_b": transformer_time_b, @@ -2434,7 +2419,6 @@ def calc_time_llm(self): kv_heads, intermediate_size, num_SMs, - include_optimizer=False, ) # transformer mem layer considers zero stage internally From 82351daa4bb5d7aacf8cab1e52243bbccb12ebfd Mon Sep 17 00:00:00 2001 From: gkarfakis19 Date: Sat, 22 Nov 2025 20:33:27 -0800 Subject: [PATCH 57/77] minor weird bug fix with wrong network analyitc gen --- LLM_execution.py | 2 +- astrasim_lib/config_generation.py | 47 ++++++++++++++++++++----------- astrasim_lib/executor.py | 4 ++- 3 files changed, 35 insertions(+), 18 deletions(-) diff --git a/LLM_execution.py b/LLM_execution.py index e68af4ec..5e20038b 100644 --- a/LLM_execution.py +++ b/LLM_execution.py @@ -804,7 +804,7 @@ def _safe_int(value: Any, default: int = 1) -> int: } def _subset_layout(allowed: Sequence[str]) -> Optional[Dict[str, Any]]: - subset = [axis for axis in axis_order if axis in allowed and axis_sizes.get(axis, 1) > 1] + subset = [axis for axis in axis_order if axis in allowed and axis_sizes.get(axis, 1) >= 1] if not subset: return None strides: Dict[str, int] = {} diff --git a/astrasim_lib/config_generation.py b/astrasim_lib/config_generation.py index 806503d1..43a3a5fb 100644 --- a/astrasim_lib/config_generation.py +++ b/astrasim_lib/config_generation.py @@ -109,6 +109,7 @@ def generate_astrasim_configs_from_hw( transform_2d_to_1d: bool = False, faulty_links_override: Optional[Sequence[Tuple[int, int, float]]] = None, ephemeral_outputs: bool = False, + preferred_axes_for_synthetic: Optional[Sequence[str]] = None, ) -> Dict[str, str]: """Write AstraSim network/system configs derived from ``hw_obj``.""" if npus_count is None: @@ -132,34 +133,48 @@ def generate_astrasim_configs_from_hw( astra_mode = str(getattr(astra_cfg, "mode", "") or "").strip().lower() sch_config = getattr(hw_obj, "sch_config", None) - def _safe_int(value: Any, default: int = 1) -> int: - try: - candidate = int(value) - except (TypeError, ValueError): - candidate = default - return max(1, candidate) axes_filter_original = tuple(axes_filter) if axes_filter else None axes_filter_normalized: Optional[Tuple[str, ...]] = None if axes_filter_original: axes_filter_normalized = tuple(str(axis).strip().lower() for axis in axes_filter_original) + preferred_axes_synth: Tuple[str, ...] = tuple( + str(axis).strip().lower() + for axis in preferred_axes_for_synthetic or () + if str(axis).strip() + ) + axis_sizes_full: Dict[str, int] = { - "tp": _safe_int(getattr(sch_config, "tp", 1)) if sch_config else 1, - "cp": _safe_int(getattr(sch_config, "cp", 1)) if sch_config else 1, - "lp": _safe_int(getattr(sch_config, "lp", 1)) if sch_config else 1, - "dp": _safe_int(getattr(sch_config, "dp", 1)) if sch_config else 1, + "tp": int(sch_config.tp), + "cp": int(sch_config.cp), + "lp": int(sch_config.lp), + "dp": int(sch_config.dp), } synthetic_only = axes_filter_normalized is not None and set(axes_filter_normalized) == {"synthetic2"} if synthetic_only: + preferred_set = set(preferred_axes_synth) + base_dim = None if dimensions: - base_dim = dimensions[0] - base_bw = float(getattr(base_dim, "effective_bandwidth", getattr(base_dim, "bandwidth", None))) - base_latency = float(getattr(base_dim, "latency", None)) - base_topology = getattr(base_dim, "topology_type", None) - base_collectives = getattr(base_dim, "collective_override", {}) or {} - else: + for dim in dimensions: + dim_axes = [ + str(axis).strip().lower() + for axis in getattr(dim, "parallelisms", ()) or () + if str(axis).strip() + ] + if preferred_set and any(axis in preferred_set for axis in dim_axes): + base_dim = dim + break + if base_dim is None and dim_axes: + base_dim = dim + if base_dim is None: + base_dim = dimensions[0] if dimensions else None + if base_dim is None: raise ValueError(f"Synthetic dimension requested but no dimensions found in hardware config. Info: Called with filter {axes_filter_original}, npus_count={npus_count}, axes_sizes={axis_sizes_full}.") + base_bw = float(getattr(base_dim, "effective_bandwidth", getattr(base_dim, "bandwidth", None))) + base_latency = float(getattr(base_dim, "latency", None)) + base_topology = getattr(base_dim, "topology_type", None) + base_collectives = getattr(base_dim, "collective_override", {}) or {} synthetic_dim = SimpleNamespace( label="synthetic2", parallelisms=("synthetic2",), diff --git a/astrasim_lib/executor.py b/astrasim_lib/executor.py index 81b803cb..5bed77bc 100644 --- a/astrasim_lib/executor.py +++ b/astrasim_lib/executor.py @@ -1726,7 +1726,8 @@ def run_astra_simulation_only_onepath( exit() rank_layout = getattr(fwdbwd_root, "_astrasim_rank_layout", None) - axis_order, axis_sizes, _ = _extract_axis_layout(rank_layout) + axis_order, axis_sizes, _ = _extract_axis_layout(rank_layout) + preferred_axes_for_synthetic = tuple(axis_order) if axis_order else tuple() axes_filter = derive_axes_filter(axis_order, axis_sizes, dp_count) if not axes_filter: axes_filter = None @@ -1738,6 +1739,7 @@ def run_astra_simulation_only_onepath( rank_count, axes_filter=axes_filter, faulty_links_override=faulty_links_override, + preferred_axes_for_synthetic=preferred_axes_for_synthetic, ) # Run AstraSim simulation on forward graph (cached via manifest) From 65913575df703b034aed385b2a51ded78be02964 Mon Sep 17 00:00:00 2001 From: NanoMembers Date: Sat, 22 Nov 2025 21:49:41 -0800 Subject: [PATCH 58/77] implemented full recomputation for hybrid/hierachical mode --- LLM_util.py | 6 ++++-- config.py | 15 ++++++++++++++- configs/hardware-config/H100_SXM5_80GB.yaml | 1 + configs/hardware-config/a100_80GB.yaml | 3 ++- .../hardware-config/a100_80GB_no_parallelism.yaml | 1 + configs/hardware-config/a100_80GB_problem.yaml | 1 + configs/hardware-config/gmap_mesh2d_demo.yaml | 1 + configs/hardware-config/v100.yaml | 1 + .../hardware-config/waferscale_20v100_80hbm.yaml | 1 + time_calculation.py | 1 + time_calculation_LLM.py | 15 ++++++++++++--- .../hardware-config/H100_SXM5_80GB.yaml | 1 + .../hardware-config/a100_80GB.yaml | 1 + .../hardware-config/a100_80GB_korthikanti.yaml | 1 + .../hardware-config/a100_80GB_no_parallelism.yaml | 1 + .../hardware-config/a100_80GB_selene_sc.yaml | 1 + 16 files changed, 44 insertions(+), 7 deletions(-) diff --git a/LLM_util.py b/LLM_util.py index 110af2bf..18bd3a00 100644 --- a/LLM_util.py +++ b/LLM_util.py @@ -123,12 +123,14 @@ def process_gemm_shapes(self, batch_size, seq_len, d_model, num_heads, kv_heads, return processed -def get_transformer_mem_layer( dp, tp, lp, mb, batch_size, hidden_dim, seq_len, intermediate_size, n_heads, precision, zero_stage, flash_attention, model_type="gpt"):# https://arxiv.org/pdf/2205.05198. https://shjwudp.github.io/blog/2023/gpt-training-memory-estimation-nemo-training-practice/ +def get_transformer_mem_layer( dp, tp, lp, mb, batch_size, hidden_dim, seq_len, intermediate_size, n_heads, precision, zero_stage, flash_attention, full_recomputation, model_type="gpt"):# https://arxiv.org/pdf/2205.05198. https://shjwudp.github.io/blog/2023/gpt-training-memory-estimation-nemo-training-practice/ """ memory estimation of transformer layer for single gpu case in inference mode is supported. other modes or layers are work in progress.""" #Activations refer to output activations that need to be stored - if flash_attention: + if full_recomputation: + act_memory_layer = seq_len * batch_size * hidden_dim * (2 ) * (precision.activations / 2) #full recompute: + elif flash_attention: act_memory_layer = seq_len * batch_size * hidden_dim * (34 / tp ) * (precision.activations / 2) #assuming selective recompute else: act_memory_layer = seq_len * batch_size * hidden_dim * (34 / tp + 5 * n_heads * seq_len/(hidden_dim * tp) ) * (precision.activations / 2) diff --git a/config.py b/config.py index ce785786..9b5749eb 100644 --- a/config.py +++ b/config.py @@ -829,7 +829,14 @@ def grad_accumulation_steps(self) -> int: ], ) SWConfig = _namedtuple( - "sw_param", ["kernel_launch_overhead", "precision", "h2d_bandwidth", "dp_zero_stage"] + "sw_param", + [ + "kernel_launch_overhead", + "precision", + "h2d_bandwidth", + "dp_zero_stage", + "full_recomputation", + ], ) SchedulingConfig = _namedtuple( @@ -1076,11 +1083,17 @@ def parse_config(filename, config_type): kernel_launch_overhead = sw_block["kernel_launch_overhead"] h2d_bandwidth = sw_block["h2d_bandwidth"] dp_zero_stage = sw_block["dp_zero_stage"] + full_recomp_raw = sw_block.get("full_recomputation", False) + if isinstance(full_recomp_raw, str): + full_recomputation = full_recomp_raw.strip().lower() in {"1", "true", "yes", "y", "on", "full"} + else: + full_recomputation = bool(full_recomp_raw) sw_config = SWConfig( kernel_launch_overhead=kernel_launch_overhead, precision=precision_config, h2d_bandwidth=h2d_bandwidth, dp_zero_stage=dp_zero_stage, + full_recomputation=full_recomputation, ) raw_parallelism_block = config_dict.get("parallelism", {}) if not isinstance(raw_parallelism_block, dict): diff --git a/configs/hardware-config/H100_SXM5_80GB.yaml b/configs/hardware-config/H100_SXM5_80GB.yaml index be3fbdfc..e80004b0 100644 --- a/configs/hardware-config/H100_SXM5_80GB.yaml +++ b/configs/hardware-config/H100_SXM5_80GB.yaml @@ -5,6 +5,7 @@ # [Semi] https://semianalysis.com/2025/06/23/nvidia-tensor-core-evolution-from-volta-to-blackwell/#4th-generation-tensor-core-warpgroup-level-asynchronous-mma sw_param: + full_recomputation: false # Enable full activation recompute during backward instead of selective recomputation kernel_launch_overhead: 9e-6 # (s) # dp_zero_stage selects ZeRO-style data-parallel sharding: 0=DDP, 1=ZeRO-1, 2=ZeRO-2, 3=ZeRO-3 # only used if data parallelism (dp) > 1 diff --git a/configs/hardware-config/a100_80GB.yaml b/configs/hardware-config/a100_80GB.yaml index 7e27d5a5..1f106505 100644 --- a/configs/hardware-config/a100_80GB.yaml +++ b/configs/hardware-config/a100_80GB.yaml @@ -1,4 +1,5 @@ sw_param: + full_recomputation: true # Enable full activation recompute during backward instead of selective recomputation kernel_launch_overhead: 9e-6 # (s) # dp_zero_stage selects ZeRO-style data-parallel sharding: 0=DDP, 1=ZeRO-1, 2=ZeRO-2, 3=ZeRO-3 # only used if data parallelism (dp) > 1 @@ -126,7 +127,7 @@ execution_backend: model: astra # analytical | astra astra: backend: astra - mode: full_astrasim_flattened # hybrid | full_astrasim_hierarchical | full_astrasim_flattened + mode: full_astrasim_hierarchica # hybrid | full_astrasim_hierarchical | full_astrasim_flattened collectives: # Default collective algorithms. Per-dimension collective_override entries can override these. # Allowed values per op: auto | ring | direct | halvingDoubling | doubleBinaryTree | mesh | mesh2d | hypercube | torus2d diff --git a/configs/hardware-config/a100_80GB_no_parallelism.yaml b/configs/hardware-config/a100_80GB_no_parallelism.yaml index 798101a2..f57e6052 100644 --- a/configs/hardware-config/a100_80GB_no_parallelism.yaml +++ b/configs/hardware-config/a100_80GB_no_parallelism.yaml @@ -2,6 +2,7 @@ # DO NOT MODIFY. sw_param: + full_recomputation: false # Enable full activation recompute during backward instead of selective recomputation kernel_launch_overhead: 9e-6 # dp_zero_stage selects ZeRO-style data-parallel sharding: 0=DDP, 1=ZeRO-1, 2=ZeRO-2, 3=ZeRO-3 # only used if data parallelism (dp) > 1 diff --git a/configs/hardware-config/a100_80GB_problem.yaml b/configs/hardware-config/a100_80GB_problem.yaml index 333a61af..e284e5df 100644 --- a/configs/hardware-config/a100_80GB_problem.yaml +++ b/configs/hardware-config/a100_80GB_problem.yaml @@ -1,4 +1,5 @@ sw_param: + full_recomputation: false # Enable full activation recompute during backward instead of selective recomputation kernel_launch_overhead: 9e-6 # (s) # dp_zero_stage selects ZeRO-style data-parallel sharding: 0=DDP, 1=ZeRO-1, 2=ZeRO-2, 3=ZeRO-3 # only used if data parallelism (dp) > 1 diff --git a/configs/hardware-config/gmap_mesh2d_demo.yaml b/configs/hardware-config/gmap_mesh2d_demo.yaml index 25708ae0..113ad29b 100644 --- a/configs/hardware-config/gmap_mesh2d_demo.yaml +++ b/configs/hardware-config/gmap_mesh2d_demo.yaml @@ -1,4 +1,5 @@ sw_param: + full_recomputation: false # Enable full activation recompute during backward instead of selective recomputation kernel_launch_overhead: 5e-6 dp_zero_stage: 0 precision: diff --git a/configs/hardware-config/v100.yaml b/configs/hardware-config/v100.yaml index a90f36a8..258290b0 100644 --- a/configs/hardware-config/v100.yaml +++ b/configs/hardware-config/v100.yaml @@ -1,5 +1,6 @@ sw_param: + full_recomputation: false # Enable full activation recompute during backward instead of selective recomputation kernel_launch_overhead: 9e-6 #9us # dp_zero_stage selects ZeRO-style data-parallel sharding: 0=DDP, 1=ZeRO-1, 2=ZeRO-2, 3=ZeRO-3 # only used if data parallelism (dp) > 1 diff --git a/configs/hardware-config/waferscale_20v100_80hbm.yaml b/configs/hardware-config/waferscale_20v100_80hbm.yaml index cfa0ecd0..cbbd980c 100644 --- a/configs/hardware-config/waferscale_20v100_80hbm.yaml +++ b/configs/hardware-config/waferscale_20v100_80hbm.yaml @@ -3,6 +3,7 @@ sw_param: + full_recomputation: false # Enable full activation recompute during backward instead of selective recomputation kernel_launch_overhead: 9e-6 #9us # dp_zero_stage selects ZeRO-style data-parallel sharding: 0=DDP, 1=ZeRO-1, 2=ZeRO-2, 3=ZeRO-3 # only used if data parallelism (dp) > 1 diff --git a/time_calculation.py b/time_calculation.py index 29f715b5..5cbc5ecc 100644 --- a/time_calculation.py +++ b/time_calculation.py @@ -203,6 +203,7 @@ def __init__(self, hw_config, model_config, mode, *, astra_policy_override: Opti # If we want proper backported LSTM support we need to fix this at *some* point. self.h2d_bandwidth = getattr(hw_config.sw_config, "h2d_bandwidth", -1) self.zero_stage = getattr(hw_config.sw_config, "dp_zero_stage", 0) + self.full_recomputation = getattr(hw_config.sw_config, "full_recomputation", False) self.attached = True # Hardware Parameters diff --git a/time_calculation_LLM.py b/time_calculation_LLM.py index 52bb9779..e896cc3e 100644 --- a/time_calculation_LLM.py +++ b/time_calculation_LLM.py @@ -433,6 +433,7 @@ def flash_attention_kernel_backward(self, batch_size, hidden_dim, seq_len, num_h attention_backward_reduction_time = 0 attention_size_b = 0 + recompute_time = 0 # assuming key and value are preloaded into shared memory before attention backward computation load_kv_bytes = Bc * d * 2 * num_SMs * self.precision.activations #load key and value for one tile from HBM to SRAM initial_load_time = self.roofline(0, load_kv_bytes, "flash_attention_initial_load_b", mem_level=self.num_levels - 1) #assume key and value of one attention head is loaded from HBM to SRAM @@ -471,8 +472,11 @@ def flash_attention_kernel_backward(self, batch_size, hidden_dim, seq_len, num_h act_dP_time *= batch_size * num_heads / max(1, self.tp) + self.O softmax_time_backward *= batch_size * num_heads / max(1, self.tp) act_dQ_time *= batch_size * num_heads / max(1, self.tp) + self.O - - attention_backward_gemm_time = initial_load_time + attn_score_time + attn_scale_softmax_time + act_dO_time + act_dP_time + softmax_time_backward + act_dQ_time + if self.full_recomputation: #attention recompute is already included in full recomputation + recompute_time = 0 + else: + recompute_time = attn_score_time + attn_scale_softmax_time #selective recomputation only recompute attention score and softmax + attention_backward_gemm_time = initial_load_time + recompute_time + act_dO_time + act_dP_time + softmax_time_backward + act_dQ_time attention_backward_time = attention_backward_gemm_time + attention_backward_reduction_time @@ -1892,6 +1896,9 @@ def _extract_forward(ret: Sequence[Any]) -> Tuple[float, float, float, float, An + transformer_timings["layernorm1"].total_backward_time() + transformer_timings["layernorm2"].total_backward_time() ) + # Full recomputation during backward requires redoing the forward compute. + if getattr(self, "full_recomputation", False): + transformer_time_b += transformer_time_f node_breakdown = { "transformer_time_f": transformer_time_f, @@ -2285,7 +2292,8 @@ def _resolve_operation(name: str) -> OperationTiming: ) transformer_forward_root = transformer_graph.construct_transformer_graph(direction="forward") if include_transformer_backward: - transformer_backward_root = transformer_graph.construct_transformer_graph(direction="backward") + bwd_direction = "both" if getattr(self, "full_recomputation", False) else "backward" + transformer_backward_root = transformer_graph.construct_transformer_graph(direction=bwd_direction) comp_times = { "embedding_f": node_breakdown.get('embedding_f', 0.0), @@ -2437,6 +2445,7 @@ def calc_time_llm(self): model_type=self.model_type, zero_stage=self.zero_stage, flash_attention=self.flash_attention, + full_recomputation=self.full_recomputation, ) ) diff --git a/validation_scripts/validation_configs/hardware-config/H100_SXM5_80GB.yaml b/validation_scripts/validation_configs/hardware-config/H100_SXM5_80GB.yaml index be3fbdfc..e80004b0 100644 --- a/validation_scripts/validation_configs/hardware-config/H100_SXM5_80GB.yaml +++ b/validation_scripts/validation_configs/hardware-config/H100_SXM5_80GB.yaml @@ -5,6 +5,7 @@ # [Semi] https://semianalysis.com/2025/06/23/nvidia-tensor-core-evolution-from-volta-to-blackwell/#4th-generation-tensor-core-warpgroup-level-asynchronous-mma sw_param: + full_recomputation: false # Enable full activation recompute during backward instead of selective recomputation kernel_launch_overhead: 9e-6 # (s) # dp_zero_stage selects ZeRO-style data-parallel sharding: 0=DDP, 1=ZeRO-1, 2=ZeRO-2, 3=ZeRO-3 # only used if data parallelism (dp) > 1 diff --git a/validation_scripts/validation_configs/hardware-config/a100_80GB.yaml b/validation_scripts/validation_configs/hardware-config/a100_80GB.yaml index 7e27d5a5..29a05637 100644 --- a/validation_scripts/validation_configs/hardware-config/a100_80GB.yaml +++ b/validation_scripts/validation_configs/hardware-config/a100_80GB.yaml @@ -1,4 +1,5 @@ sw_param: + full_recomputation: false # Enable full activation recompute during backward instead of selective recomputation kernel_launch_overhead: 9e-6 # (s) # dp_zero_stage selects ZeRO-style data-parallel sharding: 0=DDP, 1=ZeRO-1, 2=ZeRO-2, 3=ZeRO-3 # only used if data parallelism (dp) > 1 diff --git a/validation_scripts/validation_configs/hardware-config/a100_80GB_korthikanti.yaml b/validation_scripts/validation_configs/hardware-config/a100_80GB_korthikanti.yaml index a9709e7a..c4a2c9ad 100644 --- a/validation_scripts/validation_configs/hardware-config/a100_80GB_korthikanti.yaml +++ b/validation_scripts/validation_configs/hardware-config/a100_80GB_korthikanti.yaml @@ -1,4 +1,5 @@ sw_param: + full_recomputation: false # Enable full activation recompute during backward instead of selective recomputation kernel_launch_overhead: 9e-6 # (s) # dp_zero_stage selects ZeRO-style data-parallel sharding: 0=DDP, 1=ZeRO-1, 2=ZeRO-2, 3=ZeRO-3 # only used if data parallelism (dp) > 1 diff --git a/validation_scripts/validation_configs/hardware-config/a100_80GB_no_parallelism.yaml b/validation_scripts/validation_configs/hardware-config/a100_80GB_no_parallelism.yaml index 798101a2..f57e6052 100644 --- a/validation_scripts/validation_configs/hardware-config/a100_80GB_no_parallelism.yaml +++ b/validation_scripts/validation_configs/hardware-config/a100_80GB_no_parallelism.yaml @@ -2,6 +2,7 @@ # DO NOT MODIFY. sw_param: + full_recomputation: false # Enable full activation recompute during backward instead of selective recomputation kernel_launch_overhead: 9e-6 # dp_zero_stage selects ZeRO-style data-parallel sharding: 0=DDP, 1=ZeRO-1, 2=ZeRO-2, 3=ZeRO-3 # only used if data parallelism (dp) > 1 diff --git a/validation_scripts/validation_configs/hardware-config/a100_80GB_selene_sc.yaml b/validation_scripts/validation_configs/hardware-config/a100_80GB_selene_sc.yaml index 8aad58ad..0fee2d50 100644 --- a/validation_scripts/validation_configs/hardware-config/a100_80GB_selene_sc.yaml +++ b/validation_scripts/validation_configs/hardware-config/a100_80GB_selene_sc.yaml @@ -1,4 +1,5 @@ sw_param: + full_recomputation: false # Enable full activation recompute during backward instead of selective recomputation kernel_launch_overhead: 9e-6 # (s) # dp_zero_stage selects ZeRO-style data-parallel sharding: 0=DDP, 1=ZeRO-1, 2=ZeRO-2, 3=ZeRO-3 # only used if data parallelism (dp) > 1 From 0ce7b9d697b9e0358bab424dbd6cb2befde971db Mon Sep 17 00:00:00 2001 From: binglu Chen Date: Mon, 24 Nov 2025 13:50:00 -0800 Subject: [PATCH 59/77] Revert "implemented full recomputation for hybrid/hierachical mode" This reverts commit 65913575df703b034aed385b2a51ded78be02964. --- LLM_util.py | 6 ++---- config.py | 15 +-------------- configs/hardware-config/H100_SXM5_80GB.yaml | 1 - configs/hardware-config/a100_80GB.yaml | 3 +-- .../hardware-config/a100_80GB_no_parallelism.yaml | 1 - configs/hardware-config/a100_80GB_problem.yaml | 1 - configs/hardware-config/gmap_mesh2d_demo.yaml | 1 - configs/hardware-config/v100.yaml | 1 - .../hardware-config/waferscale_20v100_80hbm.yaml | 1 - time_calculation.py | 1 - time_calculation_LLM.py | 15 +++------------ .../hardware-config/H100_SXM5_80GB.yaml | 1 - .../hardware-config/a100_80GB.yaml | 1 - .../hardware-config/a100_80GB_korthikanti.yaml | 1 - .../hardware-config/a100_80GB_no_parallelism.yaml | 1 - .../hardware-config/a100_80GB_selene_sc.yaml | 1 - 16 files changed, 7 insertions(+), 44 deletions(-) diff --git a/LLM_util.py b/LLM_util.py index 18bd3a00..110af2bf 100644 --- a/LLM_util.py +++ b/LLM_util.py @@ -123,14 +123,12 @@ def process_gemm_shapes(self, batch_size, seq_len, d_model, num_heads, kv_heads, return processed -def get_transformer_mem_layer( dp, tp, lp, mb, batch_size, hidden_dim, seq_len, intermediate_size, n_heads, precision, zero_stage, flash_attention, full_recomputation, model_type="gpt"):# https://arxiv.org/pdf/2205.05198. https://shjwudp.github.io/blog/2023/gpt-training-memory-estimation-nemo-training-practice/ +def get_transformer_mem_layer( dp, tp, lp, mb, batch_size, hidden_dim, seq_len, intermediate_size, n_heads, precision, zero_stage, flash_attention, model_type="gpt"):# https://arxiv.org/pdf/2205.05198. https://shjwudp.github.io/blog/2023/gpt-training-memory-estimation-nemo-training-practice/ """ memory estimation of transformer layer for single gpu case in inference mode is supported. other modes or layers are work in progress.""" #Activations refer to output activations that need to be stored - if full_recomputation: - act_memory_layer = seq_len * batch_size * hidden_dim * (2 ) * (precision.activations / 2) #full recompute: - elif flash_attention: + if flash_attention: act_memory_layer = seq_len * batch_size * hidden_dim * (34 / tp ) * (precision.activations / 2) #assuming selective recompute else: act_memory_layer = seq_len * batch_size * hidden_dim * (34 / tp + 5 * n_heads * seq_len/(hidden_dim * tp) ) * (precision.activations / 2) diff --git a/config.py b/config.py index 9b5749eb..ce785786 100644 --- a/config.py +++ b/config.py @@ -829,14 +829,7 @@ def grad_accumulation_steps(self) -> int: ], ) SWConfig = _namedtuple( - "sw_param", - [ - "kernel_launch_overhead", - "precision", - "h2d_bandwidth", - "dp_zero_stage", - "full_recomputation", - ], + "sw_param", ["kernel_launch_overhead", "precision", "h2d_bandwidth", "dp_zero_stage"] ) SchedulingConfig = _namedtuple( @@ -1083,17 +1076,11 @@ def parse_config(filename, config_type): kernel_launch_overhead = sw_block["kernel_launch_overhead"] h2d_bandwidth = sw_block["h2d_bandwidth"] dp_zero_stage = sw_block["dp_zero_stage"] - full_recomp_raw = sw_block.get("full_recomputation", False) - if isinstance(full_recomp_raw, str): - full_recomputation = full_recomp_raw.strip().lower() in {"1", "true", "yes", "y", "on", "full"} - else: - full_recomputation = bool(full_recomp_raw) sw_config = SWConfig( kernel_launch_overhead=kernel_launch_overhead, precision=precision_config, h2d_bandwidth=h2d_bandwidth, dp_zero_stage=dp_zero_stage, - full_recomputation=full_recomputation, ) raw_parallelism_block = config_dict.get("parallelism", {}) if not isinstance(raw_parallelism_block, dict): diff --git a/configs/hardware-config/H100_SXM5_80GB.yaml b/configs/hardware-config/H100_SXM5_80GB.yaml index e80004b0..be3fbdfc 100644 --- a/configs/hardware-config/H100_SXM5_80GB.yaml +++ b/configs/hardware-config/H100_SXM5_80GB.yaml @@ -5,7 +5,6 @@ # [Semi] https://semianalysis.com/2025/06/23/nvidia-tensor-core-evolution-from-volta-to-blackwell/#4th-generation-tensor-core-warpgroup-level-asynchronous-mma sw_param: - full_recomputation: false # Enable full activation recompute during backward instead of selective recomputation kernel_launch_overhead: 9e-6 # (s) # dp_zero_stage selects ZeRO-style data-parallel sharding: 0=DDP, 1=ZeRO-1, 2=ZeRO-2, 3=ZeRO-3 # only used if data parallelism (dp) > 1 diff --git a/configs/hardware-config/a100_80GB.yaml b/configs/hardware-config/a100_80GB.yaml index 1f106505..7e27d5a5 100644 --- a/configs/hardware-config/a100_80GB.yaml +++ b/configs/hardware-config/a100_80GB.yaml @@ -1,5 +1,4 @@ sw_param: - full_recomputation: true # Enable full activation recompute during backward instead of selective recomputation kernel_launch_overhead: 9e-6 # (s) # dp_zero_stage selects ZeRO-style data-parallel sharding: 0=DDP, 1=ZeRO-1, 2=ZeRO-2, 3=ZeRO-3 # only used if data parallelism (dp) > 1 @@ -127,7 +126,7 @@ execution_backend: model: astra # analytical | astra astra: backend: astra - mode: full_astrasim_hierarchica # hybrid | full_astrasim_hierarchical | full_astrasim_flattened + mode: full_astrasim_flattened # hybrid | full_astrasim_hierarchical | full_astrasim_flattened collectives: # Default collective algorithms. Per-dimension collective_override entries can override these. # Allowed values per op: auto | ring | direct | halvingDoubling | doubleBinaryTree | mesh | mesh2d | hypercube | torus2d diff --git a/configs/hardware-config/a100_80GB_no_parallelism.yaml b/configs/hardware-config/a100_80GB_no_parallelism.yaml index f57e6052..798101a2 100644 --- a/configs/hardware-config/a100_80GB_no_parallelism.yaml +++ b/configs/hardware-config/a100_80GB_no_parallelism.yaml @@ -2,7 +2,6 @@ # DO NOT MODIFY. sw_param: - full_recomputation: false # Enable full activation recompute during backward instead of selective recomputation kernel_launch_overhead: 9e-6 # dp_zero_stage selects ZeRO-style data-parallel sharding: 0=DDP, 1=ZeRO-1, 2=ZeRO-2, 3=ZeRO-3 # only used if data parallelism (dp) > 1 diff --git a/configs/hardware-config/a100_80GB_problem.yaml b/configs/hardware-config/a100_80GB_problem.yaml index e284e5df..333a61af 100644 --- a/configs/hardware-config/a100_80GB_problem.yaml +++ b/configs/hardware-config/a100_80GB_problem.yaml @@ -1,5 +1,4 @@ sw_param: - full_recomputation: false # Enable full activation recompute during backward instead of selective recomputation kernel_launch_overhead: 9e-6 # (s) # dp_zero_stage selects ZeRO-style data-parallel sharding: 0=DDP, 1=ZeRO-1, 2=ZeRO-2, 3=ZeRO-3 # only used if data parallelism (dp) > 1 diff --git a/configs/hardware-config/gmap_mesh2d_demo.yaml b/configs/hardware-config/gmap_mesh2d_demo.yaml index 113ad29b..25708ae0 100644 --- a/configs/hardware-config/gmap_mesh2d_demo.yaml +++ b/configs/hardware-config/gmap_mesh2d_demo.yaml @@ -1,5 +1,4 @@ sw_param: - full_recomputation: false # Enable full activation recompute during backward instead of selective recomputation kernel_launch_overhead: 5e-6 dp_zero_stage: 0 precision: diff --git a/configs/hardware-config/v100.yaml b/configs/hardware-config/v100.yaml index 258290b0..a90f36a8 100644 --- a/configs/hardware-config/v100.yaml +++ b/configs/hardware-config/v100.yaml @@ -1,6 +1,5 @@ sw_param: - full_recomputation: false # Enable full activation recompute during backward instead of selective recomputation kernel_launch_overhead: 9e-6 #9us # dp_zero_stage selects ZeRO-style data-parallel sharding: 0=DDP, 1=ZeRO-1, 2=ZeRO-2, 3=ZeRO-3 # only used if data parallelism (dp) > 1 diff --git a/configs/hardware-config/waferscale_20v100_80hbm.yaml b/configs/hardware-config/waferscale_20v100_80hbm.yaml index cbbd980c..cfa0ecd0 100644 --- a/configs/hardware-config/waferscale_20v100_80hbm.yaml +++ b/configs/hardware-config/waferscale_20v100_80hbm.yaml @@ -3,7 +3,6 @@ sw_param: - full_recomputation: false # Enable full activation recompute during backward instead of selective recomputation kernel_launch_overhead: 9e-6 #9us # dp_zero_stage selects ZeRO-style data-parallel sharding: 0=DDP, 1=ZeRO-1, 2=ZeRO-2, 3=ZeRO-3 # only used if data parallelism (dp) > 1 diff --git a/time_calculation.py b/time_calculation.py index 5cbc5ecc..29f715b5 100644 --- a/time_calculation.py +++ b/time_calculation.py @@ -203,7 +203,6 @@ def __init__(self, hw_config, model_config, mode, *, astra_policy_override: Opti # If we want proper backported LSTM support we need to fix this at *some* point. self.h2d_bandwidth = getattr(hw_config.sw_config, "h2d_bandwidth", -1) self.zero_stage = getattr(hw_config.sw_config, "dp_zero_stage", 0) - self.full_recomputation = getattr(hw_config.sw_config, "full_recomputation", False) self.attached = True # Hardware Parameters diff --git a/time_calculation_LLM.py b/time_calculation_LLM.py index e896cc3e..52bb9779 100644 --- a/time_calculation_LLM.py +++ b/time_calculation_LLM.py @@ -433,7 +433,6 @@ def flash_attention_kernel_backward(self, batch_size, hidden_dim, seq_len, num_h attention_backward_reduction_time = 0 attention_size_b = 0 - recompute_time = 0 # assuming key and value are preloaded into shared memory before attention backward computation load_kv_bytes = Bc * d * 2 * num_SMs * self.precision.activations #load key and value for one tile from HBM to SRAM initial_load_time = self.roofline(0, load_kv_bytes, "flash_attention_initial_load_b", mem_level=self.num_levels - 1) #assume key and value of one attention head is loaded from HBM to SRAM @@ -472,11 +471,8 @@ def flash_attention_kernel_backward(self, batch_size, hidden_dim, seq_len, num_h act_dP_time *= batch_size * num_heads / max(1, self.tp) + self.O softmax_time_backward *= batch_size * num_heads / max(1, self.tp) act_dQ_time *= batch_size * num_heads / max(1, self.tp) + self.O - if self.full_recomputation: #attention recompute is already included in full recomputation - recompute_time = 0 - else: - recompute_time = attn_score_time + attn_scale_softmax_time #selective recomputation only recompute attention score and softmax - attention_backward_gemm_time = initial_load_time + recompute_time + act_dO_time + act_dP_time + softmax_time_backward + act_dQ_time + + attention_backward_gemm_time = initial_load_time + attn_score_time + attn_scale_softmax_time + act_dO_time + act_dP_time + softmax_time_backward + act_dQ_time attention_backward_time = attention_backward_gemm_time + attention_backward_reduction_time @@ -1896,9 +1892,6 @@ def _extract_forward(ret: Sequence[Any]) -> Tuple[float, float, float, float, An + transformer_timings["layernorm1"].total_backward_time() + transformer_timings["layernorm2"].total_backward_time() ) - # Full recomputation during backward requires redoing the forward compute. - if getattr(self, "full_recomputation", False): - transformer_time_b += transformer_time_f node_breakdown = { "transformer_time_f": transformer_time_f, @@ -2292,8 +2285,7 @@ def _resolve_operation(name: str) -> OperationTiming: ) transformer_forward_root = transformer_graph.construct_transformer_graph(direction="forward") if include_transformer_backward: - bwd_direction = "both" if getattr(self, "full_recomputation", False) else "backward" - transformer_backward_root = transformer_graph.construct_transformer_graph(direction=bwd_direction) + transformer_backward_root = transformer_graph.construct_transformer_graph(direction="backward") comp_times = { "embedding_f": node_breakdown.get('embedding_f', 0.0), @@ -2445,7 +2437,6 @@ def calc_time_llm(self): model_type=self.model_type, zero_stage=self.zero_stage, flash_attention=self.flash_attention, - full_recomputation=self.full_recomputation, ) ) diff --git a/validation_scripts/validation_configs/hardware-config/H100_SXM5_80GB.yaml b/validation_scripts/validation_configs/hardware-config/H100_SXM5_80GB.yaml index e80004b0..be3fbdfc 100644 --- a/validation_scripts/validation_configs/hardware-config/H100_SXM5_80GB.yaml +++ b/validation_scripts/validation_configs/hardware-config/H100_SXM5_80GB.yaml @@ -5,7 +5,6 @@ # [Semi] https://semianalysis.com/2025/06/23/nvidia-tensor-core-evolution-from-volta-to-blackwell/#4th-generation-tensor-core-warpgroup-level-asynchronous-mma sw_param: - full_recomputation: false # Enable full activation recompute during backward instead of selective recomputation kernel_launch_overhead: 9e-6 # (s) # dp_zero_stage selects ZeRO-style data-parallel sharding: 0=DDP, 1=ZeRO-1, 2=ZeRO-2, 3=ZeRO-3 # only used if data parallelism (dp) > 1 diff --git a/validation_scripts/validation_configs/hardware-config/a100_80GB.yaml b/validation_scripts/validation_configs/hardware-config/a100_80GB.yaml index 29a05637..7e27d5a5 100644 --- a/validation_scripts/validation_configs/hardware-config/a100_80GB.yaml +++ b/validation_scripts/validation_configs/hardware-config/a100_80GB.yaml @@ -1,5 +1,4 @@ sw_param: - full_recomputation: false # Enable full activation recompute during backward instead of selective recomputation kernel_launch_overhead: 9e-6 # (s) # dp_zero_stage selects ZeRO-style data-parallel sharding: 0=DDP, 1=ZeRO-1, 2=ZeRO-2, 3=ZeRO-3 # only used if data parallelism (dp) > 1 diff --git a/validation_scripts/validation_configs/hardware-config/a100_80GB_korthikanti.yaml b/validation_scripts/validation_configs/hardware-config/a100_80GB_korthikanti.yaml index c4a2c9ad..a9709e7a 100644 --- a/validation_scripts/validation_configs/hardware-config/a100_80GB_korthikanti.yaml +++ b/validation_scripts/validation_configs/hardware-config/a100_80GB_korthikanti.yaml @@ -1,5 +1,4 @@ sw_param: - full_recomputation: false # Enable full activation recompute during backward instead of selective recomputation kernel_launch_overhead: 9e-6 # (s) # dp_zero_stage selects ZeRO-style data-parallel sharding: 0=DDP, 1=ZeRO-1, 2=ZeRO-2, 3=ZeRO-3 # only used if data parallelism (dp) > 1 diff --git a/validation_scripts/validation_configs/hardware-config/a100_80GB_no_parallelism.yaml b/validation_scripts/validation_configs/hardware-config/a100_80GB_no_parallelism.yaml index f57e6052..798101a2 100644 --- a/validation_scripts/validation_configs/hardware-config/a100_80GB_no_parallelism.yaml +++ b/validation_scripts/validation_configs/hardware-config/a100_80GB_no_parallelism.yaml @@ -2,7 +2,6 @@ # DO NOT MODIFY. sw_param: - full_recomputation: false # Enable full activation recompute during backward instead of selective recomputation kernel_launch_overhead: 9e-6 # dp_zero_stage selects ZeRO-style data-parallel sharding: 0=DDP, 1=ZeRO-1, 2=ZeRO-2, 3=ZeRO-3 # only used if data parallelism (dp) > 1 diff --git a/validation_scripts/validation_configs/hardware-config/a100_80GB_selene_sc.yaml b/validation_scripts/validation_configs/hardware-config/a100_80GB_selene_sc.yaml index 0fee2d50..8aad58ad 100644 --- a/validation_scripts/validation_configs/hardware-config/a100_80GB_selene_sc.yaml +++ b/validation_scripts/validation_configs/hardware-config/a100_80GB_selene_sc.yaml @@ -1,5 +1,4 @@ sw_param: - full_recomputation: false # Enable full activation recompute during backward instead of selective recomputation kernel_launch_overhead: 9e-6 # (s) # dp_zero_stage selects ZeRO-style data-parallel sharding: 0=DDP, 1=ZeRO-1, 2=ZeRO-2, 3=ZeRO-3 # only used if data parallelism (dp) > 1 From 9d1105ced7ffb117b1fc30897af15e32b956d3e3 Mon Sep 17 00:00:00 2001 From: gkarfakis19 Date: Mon, 24 Nov 2025 16:56:52 -0800 Subject: [PATCH 60/77] new style astrasim config gen by Faraz --- astrasim_lib/config_generation.py | 59 +++++++++++++++++++++++++++---- 1 file changed, 53 insertions(+), 6 deletions(-) diff --git a/astrasim_lib/config_generation.py b/astrasim_lib/config_generation.py index 43a3a5fb..ab17d7fb 100644 --- a/astrasim_lib/config_generation.py +++ b/astrasim_lib/config_generation.py @@ -100,6 +100,38 @@ def _normalize_topology_name(topo: str) -> str: return "FullyConnected" +def _expand_network_entries(entries: Sequence[Dict[str, Any]]) -> List[Dict[str, Any]]: + """Expand multi-dimensional mesh/torus entries into per-dimension entries.""" + + expanded: List[Dict[str, Any]] = [] + for entry in entries: + topo = entry.get("topology") + npus_value = entry.get("npus") + bw_value = entry.get("bandwidth") + lat_value = entry.get("latency") + + if isinstance(topo, str) and (topo.startswith("Torus") or topo.startswith("Mesh")): + dim_count = int(topo[-2]) if len(topo) >= 2 and topo[-2].isdigit() else 2 + base_name = "Ring" if topo.startswith("Torus") else "Mesh" + for dim_idx in range(dim_count): + expanded.append( + { + **entry, + "topology": base_name, + "npus": ( + npus_value[dim_idx] + if isinstance(npus_value, (list, tuple)) and dim_idx < len(npus_value) + else npus_value + ), + "bandwidth": bw_value, + "latency": lat_value, + } + ) + else: + expanded.append(dict(entry)) + return expanded + + def generate_astrasim_configs_from_hw( hw_obj, out_dir: str = "./astra_cache", @@ -294,9 +326,10 @@ def generate_astrasim_configs_from_hw( ) topo_list: List[str] = [] - npus_list: List[int] = [] + npus_list: List[Any] = [] bw_list: List[float] = [] lat_list: List[float] = [] + network_entries: List[Dict[str, Any]] = [] for dim, axes_selected, product_size, _ in selected_dims: topo = _normalize_topology_name(dim.topology_type) @@ -312,10 +345,22 @@ def generate_astrasim_configs_from_hw( topo = "FullyConnected" effective_bw = float(getattr(dim, "effective_bandwidth", dim.bandwidth)) latency_s = float(dim.latency) - topo_list.append(topo) - npus_list.append(size) - bw_list.append(round(_gbps_from_bps(effective_bw), 6)) - lat_list.append(round(_ns_from_s(latency_s), 3)) + entry = { + "dim": dim, + "axes": axes_selected, + "npus": size, + "topology": topo, + "bandwidth": round(_gbps_from_bps(effective_bw), 6), + "latency": round(_ns_from_s(latency_s), 3), + } + network_entries.append(entry) + + network_entries = _expand_network_entries(network_entries) + + topo_list = [entry["topology"] for entry in network_entries] + npus_list = [entry["npus"] for entry in network_entries] + bw_list = [entry["bandwidth"] for entry in network_entries] + lat_list = [entry["latency"] for entry in network_entries] signature_parts = [str(size) for _dim, _axes, size, _ in selected_dims] dim_signature = "_".join(signature_parts) if signature_parts else f"{target}" @@ -380,7 +425,9 @@ def _collective_for_dimension(dim, topo_name: str, op: str, default_alg: str) -> rs_impl: List[str] = [] a2a_impl: List[str] = [] - for (dim, axes, _, _), topo_name in zip(selected_dims, topo_list): + for entry in network_entries: + dim = entry["dim"] + topo_name = entry["topology"] ag_impl.append(_collective_for_dimension(dim, topo_name, "all-gather", default_ag)) ar_impl.append(_collective_for_dimension(dim, topo_name, "all-reduce", default_ar)) rs_impl.append(_collective_for_dimension(dim, topo_name, "reduce-scatter", default_rs)) From 4360457309a98badf3f3d8a6685319b38a570c9a Mon Sep 17 00:00:00 2001 From: gkarfakis19 Date: Mon, 24 Nov 2025 22:37:37 -0800 Subject: [PATCH 61/77] disaggregated MHA always. needed for next hidden comm refinements --- time_calculation_LLM.py | 64 ++++++++++++++++++++++++++++++----------- time_calculation_inf.py | 12 +------- 2 files changed, 48 insertions(+), 28 deletions(-) diff --git a/time_calculation_LLM.py b/time_calculation_LLM.py index 52bb9779..5aa88d1b 100644 --- a/time_calculation_LLM.py +++ b/time_calculation_LLM.py @@ -54,6 +54,30 @@ class ParallelismMode(Enum): 'forward': {'kind': 'all_reduce', 'participants': 'tp', 'interconnect': 'tp'}, 'backward': {'kind': 'all_reduce', 'participants': 'tp', 'interconnect': 'tp'}, }, + 'layernorm1': { + 'forward': {'kind': 'all_reduce', 'participants': 'tp', 'interconnect': 'tp'}, + 'backward': {'kind': 'all_reduce', 'participants': 'tp', 'interconnect': 'tp'}, + }, + 'layernorm2': { + 'forward': {'kind': 'all_reduce', 'participants': 'tp', 'interconnect': 'tp'}, + 'backward': {'kind': 'all_reduce', 'participants': 'tp', 'interconnect': 'tp'}, + }, + 'qkv_proj': { + 'forward': {'kind': 'all_reduce', 'participants': 'tp', 'interconnect': 'tp'}, + 'backward': {'kind': 'all_reduce', 'participants': 'tp', 'interconnect': 'tp'}, + }, + 'attention': { + 'forward': {'kind': 'all_reduce', 'participants': 'tp', 'interconnect': 'tp'}, + 'backward': {'kind': 'all_reduce', 'participants': 'tp', 'interconnect': 'tp'}, + }, + 'output_proj': { + 'forward': {'kind': 'all_reduce', 'participants': 'tp', 'interconnect': 'tp'}, + 'backward': {'kind': 'all_reduce', 'participants': 'tp', 'interconnect': 'tp'}, + }, + 'MLP': { + 'forward': {'kind': 'all_reduce', 'participants': 'tp', 'interconnect': 'tp'}, + 'backward': {'kind': 'all_reduce', 'participants': 'tp', 'interconnect': 'tp'}, + }, }, ParallelismMode.TENSOR_SEQUENCE: { COMM_RULE_DEFAULT_KEY: {'forward': None, 'backward': None}, # <- dangerous @@ -65,7 +89,15 @@ class ParallelismMode(Enum): 'forward': {'kind': 'all_gather', 'participants': 'tp', 'interconnect': 'tp'}, 'backward': {'kind': 'all_gather', 'participants': 'tp', 'interconnect': 'tp'}, }, - 'MHA': { + 'qkv_proj': { + 'forward': {'kind': 'reduce_scatter', 'participants': 'tp', 'interconnect': 'tp'}, + 'backward': {'kind': 'reduce_scatter', 'participants': 'tp', 'interconnect': 'tp'}, + }, + 'attention': { + 'forward': {'kind': 'reduce_scatter', 'participants': 'tp', 'interconnect': 'tp'}, + 'backward': {'kind': 'reduce_scatter', 'participants': 'tp', 'interconnect': 'tp'}, + }, + 'output_proj': { 'forward': {'kind': 'reduce_scatter', 'participants': 'tp', 'interconnect': 'tp'}, 'backward': {'kind': 'reduce_scatter', 'participants': 'tp', 'interconnect': 'tp'}, }, @@ -1863,14 +1895,6 @@ def _extract_forward(ret: Sequence[Any]) -> Tuple[float, float, float, float, An ), ) - mha_group = OperationGroup( - "MHA", - operations=( - transformer_timings["qkv_proj"], - transformer_timings["attention"], - transformer_timings["output_proj"], - ), - ) mlp_group = OperationGroup( "MLP", operations=( @@ -1881,13 +1905,17 @@ def _extract_forward(ret: Sequence[Any]) -> Tuple[float, float, float, float, An ) transformer_time_f = ( - mha_group.forward_total_time() + transformer_timings["qkv_proj"].total_forward_time() + + transformer_timings["attention"].total_forward_time() + + transformer_timings["output_proj"].total_forward_time() + mlp_group.forward_total_time() + transformer_timings["layernorm1"].total_forward_time() + transformer_timings["layernorm2"].total_forward_time() ) transformer_time_b = ( - mha_group.backward_total_time() + transformer_timings["qkv_proj"].total_backward_time() + + transformer_timings["attention"].total_backward_time() + + transformer_timings["output_proj"].total_backward_time() + mlp_group.backward_total_time() + transformer_timings["layernorm1"].total_backward_time() + transformer_timings["layernorm2"].total_backward_time() @@ -2123,7 +2151,6 @@ def _resolve_comm_rules(mode: ParallelismMode) -> Dict[str, Dict[str, Optional[D } participants_lookup['moe'] = max(1, participants_lookup['tp'] * participants_lookup['cp']) group_members = { - "MHA": ("qkv_proj", "attention", "output_proj"), "MLP": ("ffn1", "gelu", "ffn2"), } @@ -2233,12 +2260,15 @@ def _resolve_operation(name: str) -> OperationTiming: return _build_group_operation(name, group_members[name]) return transformer_timings[name] - if parallelism_mode in (ParallelismMode.CONTEXT, ParallelismMode.TENSOR_CONTEXT_HYBRID): - op_names: Sequence[str] = ("layernorm1", "qkv_proj", "attention", "output_proj", "layernorm2", "MLP") - elif parallelism_mode in (ParallelismMode.TENSOR, ParallelismMode.TENSOR_SEQUENCE, ParallelismMode.SINGLE): - op_names = ("layernorm1", "MHA", "layernorm2", "MLP") - else: + if parallelism_mode not in ( + ParallelismMode.CONTEXT, + ParallelismMode.TENSOR_CONTEXT_HYBRID, + ParallelismMode.TENSOR, + ParallelismMode.TENSOR_SEQUENCE, + ParallelismMode.SINGLE, + ): raise ValueError(f"Unsupported parallelism mode: {parallelism_mode}") + op_names: Sequence[str] = ("layernorm1", "qkv_proj", "attention", "output_proj", "layernorm2", "MLP") op_names = list(op_names) if "optimizer" in transformer_timings: op_names.append("optimizer") diff --git a/time_calculation_inf.py b/time_calculation_inf.py index c71a56f8..f6bf7aba 100644 --- a/time_calculation_inf.py +++ b/time_calculation_inf.py @@ -287,15 +287,6 @@ def _make_forward( backward=None, ) - # Aggregate operations - mha_group = OperationGroup( - "MHA", - operations=( - transformer_timings["qkv_proj"], - transformer_timings["attention"], - transformer_timings["output_proj"], - ), - ) mlp_group = OperationGroup( "MLP", operations=( @@ -393,10 +384,9 @@ def calc_energy(self, transformer_timings: Dict[str, OperationTiming], cross_lay inter_comm_bytes = 0.0 # data parallelism? aggregate_groups = { - "MHA": ("qkv_proj", "attention", "output_proj"), "MLP": ("ffn1", "gelu", "ffn2"), } - solo_ops = ("layernorm1", "layernorm2", "embedding", "linear_softmax") + solo_ops = ("layernorm1", "layernorm2", "embedding", "linear_softmax", "qkv_proj", "attention", "output_proj") for members in aggregate_groups.values(): for name in members: From 9b96b87c3420f0fad58c7a2eb3364b9a87c3f14d Mon Sep 17 00:00:00 2001 From: gkarfakis19 Date: Mon, 24 Nov 2025 22:47:15 -0800 Subject: [PATCH 62/77] added zero 3 inference check --- time_calculation_inf.py | 2 ++ 1 file changed, 2 insertions(+) diff --git a/time_calculation_inf.py b/time_calculation_inf.py index f6bf7aba..f2dbd0e9 100644 --- a/time_calculation_inf.py +++ b/time_calculation_inf.py @@ -33,6 +33,8 @@ class TimeCalculationLLMInference(TimeCalculationLLM): def __init__(self, hw_config, model_config, mode, output_dir: Optional[str] = None): super().__init__(hw_config, model_config, mode, output_dir) + if self.zero_stage >= 3 and self.dp > 1: + raise ValueError("ZeRO-3 data parallelism is not supported for inference runs (dp_zero_stage must be <3 or dp=1).") self._raw_model_config = model_config def _build_decode_transformer_results( From d1a55b519e145c365501a9c9d678aaf587d3c2f6 Mon Sep 17 00:00:00 2001 From: gkarfakis19 Date: Tue, 25 Nov 2025 01:34:20 -0800 Subject: [PATCH 63/77] first attempt at split comm/comp. passes all tests, reduces worst case error around ~2% --- AGENTS.md | 12 + LLM_execution.py | 275 ++++++++++++++++++ astra-sim | 2 +- config.py | 56 +++- configs/hardware-config/H100_SXM5_80GB.yaml | 7 + configs/hardware-config/a100_80GB.yaml | 4 + .../a100_80GB_no_parallelism.yaml | 7 + configs/hardware-config/gmap_mesh2d_demo.yaml | 7 + configs/hardware-config/v100.yaml | 7 + .../waferscale_20v100_80hbm.yaml | 7 + examples/korthi.sh | 2 +- time_calculation_LLM.py | 25 +- tools/mod_field.py | 236 +++++++++++++++ tools/overlap_finder.py | 243 ++++++++++++++++ .../hardware-config/H100_SXM5_80GB.yaml | 7 + .../hardware-config/a100_80GB.yaml | 7 + .../a100_80GB_korthikanti.yaml | 10 +- .../a100_80GB_no_parallelism.yaml | 7 + .../hardware-config/a100_80GB_selene_sc.yaml | 7 + 19 files changed, 910 insertions(+), 18 deletions(-) create mode 100644 AGENTS.md create mode 100644 tools/mod_field.py create mode 100644 tools/overlap_finder.py diff --git a/AGENTS.md b/AGENTS.md new file mode 100644 index 00000000..ce5a1188 --- /dev/null +++ b/AGENTS.md @@ -0,0 +1,12 @@ +DeepFlow quick brief for future agents + +- Purpose: Python simulator for LSTM/GEMM/LLM performance, memory, and (rudimentary) energy. Entry point `run_perf.py` loads YAML configs (hardware + model) and builds compute/comm graphs, then simulates analytically or via AstraSim. +- Execution backends (set in hardware config `execution_backend`): analytical (ring-only, no congestion); hybrid (AstraSim for transformer block graph, analytical pipeline); full_astrasim_hierarchical (AstraSim separately for transformer and pipeline graphs, assumes no congestion between them); full_astrasim_flattened (one big combined graph in AstraSim, slowest/most detailed). +- Parallelism model: supports DP, TP (+optional sequence parallel), CP, GPipe-style pipeline (`lp` stages), and hybrid TP+CP. Pipeline graph assumes uniform per-layer timing and GPipe scheduling; gradient accumulation builds a no-DP variant so only the last accumulation step does DP all-reduce. +- Communication assumptions: blocking collectives; analytical path uses ring formulas; AstraSim collectives flatten 2D topologies to 1D for collectives. Hierarchical mode requires the first active network dimension map exactly to {tp,cp}; non-ring topologies are rejected in analytical mode. Faulty links only valid in full AstraSim modes. +- ZeRO assumptions: stage 0/1 ≈ DDP; stage 2 adds per-layer param AG + grad RS sizes; stage 3 adds per-layer param materialization (counted as comm + ephemeral memory, no overlap). MoE assumes perfect balance and TP all-to-all; multi-GPU MoE is WIP. +- Attention: dense by default. FlashAttention supported for training forward/prefill when `attention.use_flashattention=true` and tile size is set manually; decode path ignores FlashAttention. Sliding-window/MLA are WIP. +- Memory model: static formulas for activations/weights/optimizer; KV-cache memory estimation WIP. Peak memory compared to DRAM per device; ZeRO-3 ephemeral gather tracked separately. No fragmentation/eviction modeling. +- Inference: `TimeCalculationLLMInference` handles prefill + sampled decode steps (linear interpolation between samples). DP acts as replica multiplier only; CP decode is noted WIP. FlashAttention disabled for decode. +- Artifacts/logs: outputs under `output/` (mode subdir). AstraSim artifacts live in `astra_cache/` and optionally `output/LLM/astra_hier` or `astra_flat` when persistence flags are set. Current logs/marker: `log/this_folder_contains_astrasim_logs`, `log/log.log` hold prior runs. Debug hardware snapshots in `dbg_hw_conf/20251105_173209/` (e.g., `baseline_001_tp64_cp1_dp1_lp1.yaml`, `baseline_053_tp1_cp1_dp64_lp1.yaml`). +- Key environment flags: `DEEPFLOW_VISUALIZE_GRAPHS`, `DEEPFLOW_PERSIST_ASTRASIM_ARTIFACTS`, `DEEPFLOW_PERSIST_ARTIFACT_VIZ`, `DEEPFLOW_ASTRA_CACHE_MODE` (default CACHE_READWRITE). diff --git a/LLM_execution.py b/LLM_execution.py index 5e20038b..43e95150 100644 --- a/LLM_execution.py +++ b/LLM_execution.py @@ -14,6 +14,274 @@ from astrasim_lib.layout_utils import axis_layout_from_descriptor from simulate_LLM import Graph from util import log_message +from typing import Iterable + + +def _mode_label(mode: Any) -> str: + for attr in ("value", "name"): + if hasattr(mode, attr): + return str(getattr(mode, attr)).lower() + return str(mode).lower() + + +def _copy_node_metadata(source: simulate_LLM.Node, target: simulate_LLM.Node) -> None: + for attr in ( + "micro_batch_index", + "layer_index", + "direction", + "stage_id", + "tp_rank", + "cp_rank", + ): + if hasattr(source, attr): + setattr(target, attr, getattr(source, attr)) + + +def _copy_edge_metadata(source: simulate_LLM.Edge, target: simulate_LLM.Edge) -> None: + for attr in ( + "local_hw_id", + "stage_id", + "micro_batch_index", + "layer_index", + "direction", + "tp_rank", + "cp_rank", + ): + if hasattr(source, attr): + setattr(target, attr, getattr(source, attr)) + + +def _detach_edge(parent: Any, child: Any) -> None: + children = getattr(parent, "children", None) + if isinstance(children, list): + try: + children.remove(child) + except ValueError: + pass + parents = getattr(child, "parents", None) + if isinstance(parents, list): + try: + parents.remove(parent) + except ValueError: + pass + + +def _connect_edge(parent: Any, child: Any) -> None: + if child in getattr(parent, "children", []): + return + parent.add_child(child) + + +def _split_tp_node( + node: simulate_LLM.Node, + tp_children: List[simulate_LLM.Edge], + overlap: float, +) -> None: + if overlap <= 0.0 or not tp_children: + return + + duration = float(getattr(node, "duration", 0.0) or 0.0) + if duration <= 0.0: + return + + if overlap >= 1.0: + parents = list(getattr(node, "parents", [])) + tp_succs: List[Any] = [] + for tp_edge in tp_children: + _detach_edge(node, tp_edge) + for parent in parents: + _connect_edge(parent, tp_edge) + for succ in list(getattr(tp_edge, "children", []) or []): + tp_succs.append(succ) + for succ in tp_succs: + _connect_edge(node, succ) + return + + head_duration = duration * (1.0 - overlap) + tail_duration = duration * overlap + if head_duration <= 0.0 or tail_duration <= 0.0: + return + + tail = node + tail.duration = tail_duration + head = simulate_LLM.Node( + name=f"{tail.name}_head", + op_id=getattr(tail, "op_id", 0), + hw_id=tail.hw_id, + duration=head_duration, + fwd=tail.fwd, + ) + _copy_node_metadata(tail, head) + + parents = list(getattr(tail, "parents", [])) + for parent in parents: + _detach_edge(parent, tail) + _connect_edge(parent, head) + + _connect_edge(head, tail) + + for tp_edge in tp_children: + _detach_edge(tail, tp_edge) + _connect_edge(head, tp_edge) + for succ in list(getattr(tp_edge, "children", []) or []): + _connect_edge(tail, succ) + + +def _apply_tp_overlap_transforms(root: Any, parallelism_mode: Any, tp_overlap: float, tp_sp_overlap: float) -> Any: + mode = _mode_label(parallelism_mode) + if mode == "tensor_sequence": + overlap = tp_sp_overlap + elif mode in {"tensor", "tensor_context_hybrid"}: + overlap = tp_overlap + else: + return root + if overlap <= 0.0: + return root + + nodes_to_process: List[simulate_LLM.Node] = [] + visited: Set[int] = set() + stack: List[Any] = list(root) if isinstance(root, (list, tuple)) else [root] + while stack: + obj = stack.pop() + obj_id = id(obj) + if obj_id in visited: + continue + visited.add(obj_id) + if isinstance(obj, simulate_LLM.Node): + nodes_to_process.append(obj) + for child in getattr(obj, "children", []): + stack.append(child) + + for node in nodes_to_process: + tp_children = [ + child for child in getattr(node, "children", []) + if isinstance(child, simulate_LLM.Edge) and getattr(child, "comm_interconnect_type", None) == "tp" + ] + if not tp_children: + continue + _split_tp_node(node, tp_children, overlap) + return root + + +def _apply_cp_overlap_transforms(root: Any, parallelism_mode: Any, cp_overlap: float) -> Any: + mode = _mode_label(parallelism_mode) + if mode not in {"context", "tensor_context_hybrid"}: + return root + overlap = cp_overlap + if overlap <= 0.0: + return root + + cp_edges: List[simulate_LLM.Edge] = [] + visited: Set[int] = set() + stack: List[Any] = list(root) if isinstance(root, (list, tuple)) else [root] + while stack: + obj = stack.pop() + obj_id = id(obj) + if obj_id in visited: + continue + visited.add(obj_id) + if isinstance(obj, simulate_LLM.Edge) and getattr(obj, "comm_interconnect_type", None) == "cp": + cp_edges.append(obj) + for child in getattr(obj, "children", []): + stack.append(child) + + for edge in cp_edges: + _split_cp_edge(edge, overlap) + return root + + +def _split_cp_edge(edge: simulate_LLM.Edge, overlap: float) -> None: + attention_children = [ + child for child in getattr(edge, "children", []) + if isinstance(child, simulate_LLM.Node) and "attention" in str(getattr(child, "name", "")).lower() + ] + if not attention_children: + return + + total_bytes = int(getattr(edge, "comm_size_bytes", 0) or 0) + preds = list(getattr(edge, "parents", [])) + succs = list(getattr(edge, "children", [])) + + if overlap >= 1.0 or total_bytes <= 0: + for attention in attention_children: + _detach_edge(edge, attention) + for pred in preds: + _connect_edge(pred, attention) + for succ in succs: + if succ in attention_children: + continue + for attention in attention_children: + _connect_edge(attention, succ) + _connect_edge(edge, succ) + return + + block_bytes = int(math.ceil(total_bytes * (1.0 - overlap))) + ovlp_bytes = max(0, total_bytes - block_bytes) + if block_bytes <= 0: + _split_cp_edge(edge, 1.0) + return + + block_edge = simulate_LLM.Edge( + name=f"{edge.name}_block", + op_id=getattr(edge, "op_id", 0), + duration=0, + is_dp=edge.is_dp, + comm_size_bytes=block_bytes, + comm_type=edge.comm_type, + participants=edge.participants, + comm_interconnect_type=edge.comm_interconnect_type, + ) + ovlp_edge = None + if ovlp_bytes > 0: + ovlp_edge = simulate_LLM.Edge( + name=f"{edge.name}_ovlp", + op_id=getattr(edge, "op_id", 0), + duration=0, + is_dp=edge.is_dp, + comm_size_bytes=ovlp_bytes, + comm_type=edge.comm_type, + participants=edge.participants, + comm_interconnect_type=edge.comm_interconnect_type, + ) + _copy_edge_metadata(edge, ovlp_edge) + _copy_edge_metadata(edge, block_edge) + + for pred in preds: + _detach_edge(pred, edge) + _connect_edge(pred, block_edge) + + for attention in attention_children: + _detach_edge(edge, attention) + _connect_edge(block_edge, attention) + + if ovlp_edge is not None: + _connect_edge(block_edge, ovlp_edge) + + for succ in succs: + if succ in attention_children: + continue + _detach_edge(edge, succ) + for attention in attention_children: + _connect_edge(attention, succ) + if ovlp_edge is not None: + _connect_edge(ovlp_edge, succ) + else: + _connect_edge(block_edge, succ) + + +def apply_overlap_transforms( + root: Any, + parallelism_mode: Any, + tp_overlap: float, + tp_sp_overlap: float, + cp_overlap: float, +) -> Any: + """Apply TP/TP_SP/CP overlap rewrites to a constructed graph.""" + if root is None: + return None + root = _apply_tp_overlap_transforms(root, parallelism_mode, tp_overlap, tp_sp_overlap) + root = _apply_cp_overlap_transforms(root, parallelism_mode, cp_overlap) + return root def _env_flag(name: str) -> bool: @@ -1197,6 +1465,13 @@ def _run_full_astrasim_flattened(self) -> ExecutionResult: if flattened_root is None: raise RuntimeError("Pipeline flattening produced an empty graph") + flattened_root = apply_overlap_transforms( + flattened_root, + self.time_calc.get_parallelism_mode(), + getattr(self.time_calc, "tp_overlap", 0.0), + getattr(self.time_calc, "tp_sp_overlap", 0.0), + getattr(self.time_calc, "cp_overlap", 0.0), + ) flattener._propagate_local_hw_ids(flattened_root) setattr(flattened_root, "_astrasim_rank_layout", self._rank_layout) self._attach_optimize_hint(flattened_root) diff --git a/astra-sim b/astra-sim index 3f56a2f0..ee9dabd2 160000 --- a/astra-sim +++ b/astra-sim @@ -1 +1 @@ -Subproject commit 3f56a2f016cc9f8718789e9d2a401fb84dc3a847 +Subproject commit ee9dabd2496ada3825ba7b388d7f4016fb213e7f diff --git a/config.py b/config.py index ce785786..684e3862 100644 --- a/config.py +++ b/config.py @@ -580,17 +580,18 @@ def _compute_dimension_parallelism_product( def _parse_network_layout( network_spec, parallelism_params: Dict[str, object], -) -> Tuple[Tuple[NetworkDimensionLayout, ...], Tuple[Tuple[int, int, float], ...]]: +) -> Tuple[Tuple[NetworkDimensionLayout, ...], Tuple[Tuple[int, int, float], ...], "NetworkOverlapConfig"]: if network_spec is None: - return tuple(), tuple() + raise ValueError("network section must be specified and include overlap settings") - faulty_links: Tuple[Tuple[int, int, float], ...] = tuple() - dimensions_spec = network_spec - if isinstance(network_spec, dict): - faulty_links = _parse_faulty_links("network", network_spec.get("faulty_links", [])) - dimensions_spec = network_spec.get("dimensions") - if dimensions_spec is None: - raise ValueError("network.dimensions must be specified when network is a mapping") + if not isinstance(network_spec, dict): + raise ValueError("network must be provided as a mapping to supply overlap settings") + + faulty_links: Tuple[Tuple[int, int, float], ...] = _parse_faulty_links("network", network_spec.get("faulty_links", [])) + overlap_config = _parse_network_overlap(network_spec.get("overlap")) + dimensions_spec = network_spec.get("dimensions") + if dimensions_spec is None: + raise ValueError("network.dimensions must be specified when network is a mapping") if not isinstance(dimensions_spec, Sequence) or isinstance(dimensions_spec, (str, bytes)): raise ValueError("network.dimensions must be a sequence of dimension mappings") @@ -604,7 +605,14 @@ def _parse_network_layout( index=index, ) ) - return tuple(dimensions), faulty_links + return tuple(dimensions), faulty_links, overlap_config + + +@dataclass(frozen=True) +class NetworkOverlapConfig: + tp_overlap: float + tp_sp_overlap: float + cp_overlap: float @dataclass(frozen=True) @@ -612,6 +620,7 @@ class NetworkLayoutConfig: dimensions: Tuple[NetworkDimensionLayout, ...] faulty_links: Tuple[Tuple[int, int, float], ...] = field(default_factory=tuple) parallelism_map: Dict[str, NetworkDimensionLayout] = field(default_factory=dict) + overlap_config: "NetworkOverlapConfig" = None def primary_dimension(self) -> Optional[NetworkDimensionLayout]: return self.dimensions[0] if self.dimensions else None @@ -672,10 +681,32 @@ def _parse_faulty_links(owner_label: str, faulty_links_raw) -> Tuple[Tuple[int, entries.append((src, dst, weight)) return tuple(entries) +def _parse_network_overlap(overlap_raw) -> "NetworkOverlapConfig": + if not isinstance(overlap_raw, dict): + raise ValueError("network.overlap must be a mapping with tp_overlap, tp_sp_overlap, and cp_overlap") + required_fields = ("tp_overlap", "tp_sp_overlap", "cp_overlap") + values = {} + for field in required_fields: + if field not in overlap_raw: + raise ValueError(f"network.overlap missing required field '{field}'") + try: + val = float(overlap_raw[field]) + except (TypeError, ValueError) as exc: + raise ValueError(f"network.overlap.{field} must be numeric") from exc + if val < 0.0 or val > 1.0: + raise ValueError(f"network.overlap.{field} must be between 0.0 and 1.0") + values[field] = val + return NetworkOverlapConfig( + tp_overlap=values["tp_overlap"], + tp_sp_overlap=values["tp_sp_overlap"], + cp_overlap=values["cp_overlap"], + ) + def _build_network_layout_config( dimensions: Sequence[NetworkDimensionLayout], faulty_links: Sequence[Tuple[int, int, float]] = (), + overlap_config: Optional["NetworkOverlapConfig"] = None, ) -> NetworkLayoutConfig: parallelism_map: Dict[str, NetworkDimensionLayout] = {} for dim in dimensions: @@ -689,6 +720,7 @@ def _build_network_layout_config( dimensions=tuple(dimensions), faulty_links=tuple(faulty_links), parallelism_map=parallelism_map, + overlap_config=overlap_config, ) @@ -1093,10 +1125,10 @@ def parse_config(filename, config_type): scheduling_for_network = {str(k).lower(): v for k, v in parallelism_params.items()} network_spec = config_dict.get("network") - network_dimensions, network_faults = _parse_network_layout(network_spec, scheduling_for_network) + network_dimensions, network_faults, network_overlap = _parse_network_layout(network_spec, scheduling_for_network) if not network_dimensions: raise ValueError("network section must define at least one dimension") - network_layout_config = _build_network_layout_config(network_dimensions, network_faults) + network_layout_config = _build_network_layout_config(network_dimensions, network_faults, network_overlap) tech_config = TechConfig.from_dict(config_dict["tech_param"]) diff --git a/configs/hardware-config/H100_SXM5_80GB.yaml b/configs/hardware-config/H100_SXM5_80GB.yaml index be3fbdfc..f7db00c0 100644 --- a/configs/hardware-config/H100_SXM5_80GB.yaml +++ b/configs/hardware-config/H100_SXM5_80GB.yaml @@ -127,6 +127,13 @@ network: parallelisms: [] collective_override: {} # overrides execution_backend.collectives when provided + overlap: + # fraction of TP+SP compute that can overlap TP/SP comm + tp_sp_overlap: 0.63 + # fraction of TP compute that can overlap TP comm + tp_overlap: 0.6 + # fraction of CP comm that can overlap CP compute + cp_overlap: 0.0 execution_backend: model: analytical # analytical | astra astra: diff --git a/configs/hardware-config/a100_80GB.yaml b/configs/hardware-config/a100_80GB.yaml index 7e27d5a5..dd80a02f 100644 --- a/configs/hardware-config/a100_80GB.yaml +++ b/configs/hardware-config/a100_80GB.yaml @@ -83,6 +83,10 @@ parallelism: mb: 4 # micro-batches (for best results set to >= lp). network: + overlap: + tp_overlap: 0.6 + tp_sp_overlap: 0.63 + cp_overlap: 1.0 # Global faulty links apply to the active AstraSim dimension: [[src, dst, weight]] # faulty_links: [[0,1,0.6],[16,17,0.1]] dimensions: diff --git a/configs/hardware-config/a100_80GB_no_parallelism.yaml b/configs/hardware-config/a100_80GB_no_parallelism.yaml index 798101a2..6810ae40 100644 --- a/configs/hardware-config/a100_80GB_no_parallelism.yaml +++ b/configs/hardware-config/a100_80GB_no_parallelism.yaml @@ -110,6 +110,13 @@ network: + overlap: + # fraction of TP+SP compute that can overlap TP/SP comm + tp_sp_overlap: 0.63 + # fraction of TP compute that can overlap TP comm + tp_overlap: 0.6 + # fraction of CP comm that can overlap CP compute + cp_overlap: 0.0 execution_backend: model: analytical # analytical | astra astra: diff --git a/configs/hardware-config/gmap_mesh2d_demo.yaml b/configs/hardware-config/gmap_mesh2d_demo.yaml index 25708ae0..ce38e619 100644 --- a/configs/hardware-config/gmap_mesh2d_demo.yaml +++ b/configs/hardware-config/gmap_mesh2d_demo.yaml @@ -120,6 +120,13 @@ network: util: 1.0 parallelisms: [dp] + overlap: + # fraction of TP+SP compute that can overlap TP/SP comm + tp_sp_overlap: 0.63 + # fraction of TP compute that can overlap TP comm + tp_overlap: 0.6 + # fraction of CP comm that can overlap CP compute + cp_overlap: 0.0 power_breakdown: TDP: 400 core: 200 diff --git a/configs/hardware-config/v100.yaml b/configs/hardware-config/v100.yaml index a90f36a8..6e390c71 100644 --- a/configs/hardware-config/v100.yaml +++ b/configs/hardware-config/v100.yaml @@ -109,6 +109,13 @@ network: parallelisms: [tp, cp, lp, dp] + overlap: + # fraction of TP+SP compute that can overlap TP/SP comm + tp_sp_overlap: 0.63 + # fraction of TP compute that can overlap TP comm + tp_overlap: 0.6 + # fraction of CP comm that can overlap CP compute + cp_overlap: 0.0 execution_backend: model: astra # analytical | astra astra: diff --git a/configs/hardware-config/waferscale_20v100_80hbm.yaml b/configs/hardware-config/waferscale_20v100_80hbm.yaml index cfa0ecd0..1e9649dc 100644 --- a/configs/hardware-config/waferscale_20v100_80hbm.yaml +++ b/configs/hardware-config/waferscale_20v100_80hbm.yaml @@ -125,6 +125,13 @@ network: collective_override: {} # overrides execution_backend.collectives when provided parallelisms: [tp, cp, lp, dp] + overlap: + # fraction of TP+SP compute that can overlap TP/SP comm + tp_sp_overlap: 0.63 + # fraction of TP compute that can overlap TP comm + tp_overlap: 0.6 + # fraction of CP comm that can overlap CP compute + cp_overlap: 0.0 execution_backend: model: astra # analytical | astra astra: diff --git a/examples/korthi.sh b/examples/korthi.sh index d2232137..2cb3c7f7 100755 --- a/examples/korthi.sh +++ b/examples/korthi.sh @@ -6,6 +6,6 @@ DEEPFLOW_PERSIST_ASTRASIM_ARTIFACTS=1 \ DEEPFLOW_VISUALIZE_GRAPHS=1 \ DEEPFLOW_PERSIST_ARTIFACT_VIZ=1 \ -python3 run_perf.py \ +uv run run_perf.py \ --hardware_config validation_scripts/validation_configs/hardware-config/a100_80GB_korthikanti.yaml \ --model_config validation_scripts/validation_configs/model-config/GPT_22_B.yaml diff --git a/time_calculation_LLM.py b/time_calculation_LLM.py index 5aa88d1b..a9afa074 100644 --- a/time_calculation_LLM.py +++ b/time_calculation_LLM.py @@ -6,7 +6,7 @@ from typing import Any, Dict, Tuple, Optional, List, Mapping, Sequence, Set from collections.abc import Mapping as MappingABC, Sequence as SequenceABC import simulate_LLM -from LLM_execution import ExecutionMode, LLMExecutionDispatcher +from LLM_execution import ExecutionMode, LLMExecutionDispatcher, apply_overlap_transforms from simulate_LLM import Graph import LLM_util from time_calculation import TimeCalculation @@ -212,6 +212,12 @@ def __init__(self, hw_config, model_config, mode, output_dir: Optional[str] = No self.output_dir = os.path.abspath(output_dir) if output_dir else os.getcwd() os.makedirs(self.output_dir, exist_ok=True) + overlap_cfg = getattr(getattr(hw_config, "network_layout", None), "overlap_config", None) + if overlap_cfg is None: + raise ValueError("network.overlap must be provided in the hardware config") + self.tp_overlap = float(overlap_cfg.tp_overlap) + self.tp_sp_overlap = float(overlap_cfg.tp_sp_overlap) + self.cp_overlap = float(overlap_cfg.cp_overlap) self._generate_graphs = _env_flag("DEEPFLOW_VISUALIZE_GRAPHS") self.persist_astrasim_artifacts = _env_flag("DEEPFLOW_PERSIST_ASTRASIM_ARTIFACTS") self._debug_memory = _env_flag("DEEPFLOW_DEBUG_MEMORY") @@ -2059,6 +2065,7 @@ def _build_interconnect_params(self) -> Dict[str, Tuple[float, float]]: 'cp': (self.IBTP, self.LLTP), } + def _prepare_execution_graphs( self, *, @@ -2317,6 +2324,22 @@ def _resolve_operation(name: str) -> OperationTiming: if include_transformer_backward: transformer_backward_root = transformer_graph.construct_transformer_graph(direction="backward") + transformer_forward_root = apply_overlap_transforms( + transformer_forward_root, + parallelism_mode, + self.tp_overlap, + self.tp_sp_overlap, + self.cp_overlap, + ) + if include_transformer_backward: + transformer_backward_root = apply_overlap_transforms( + transformer_backward_root, + parallelism_mode, + self.tp_overlap, + self.tp_sp_overlap, + self.cp_overlap, + ) + comp_times = { "embedding_f": node_breakdown.get('embedding_f', 0.0), "embedding_b": node_breakdown.get('embedding_b', 0.0) if include_pipeline_backward else 0.0, diff --git a/tools/mod_field.py b/tools/mod_field.py new file mode 100644 index 00000000..6f8105d9 --- /dev/null +++ b/tools/mod_field.py @@ -0,0 +1,236 @@ +#!/usr/bin/env python3 +""" +Modify or insert a scalar field in YAML configs (PyYAML only, minimal text edit). + +- If the dotted path exists, overwrite its scalar value (preserving inline comment). +- If it does not exist, create it (with optional comment) under existing ancestor mappings. +- Skips files where an ancestor is not a mapping. +""" + +from __future__ import annotations + +import argparse +from pathlib import Path +from typing import Any, Dict, List, Mapping, Optional, Tuple + +import yaml + + +def parse_args() -> argparse.Namespace: + parser = argparse.ArgumentParser(description="Modify/insert a YAML field.") + parser.add_argument("--scope", choices=["hardware", "software"], required=True, help="which config family to update") + parser.add_argument("--path", required=True, help="dotted path to set (e.g., network.overlap.tp_overlap)") + parser.add_argument("--value", required=True, help='YAML literal for the value (e.g., 0.0, true, "text")') + parser.add_argument("--comment", default="", help="comment to place above the inserted field (add-only)") + parser.add_argument("--dry-run", action="store_true", help="print planned changes without writing files") + return parser.parse_args() + + +def parse_value(raw: str) -> Any: + try: + return yaml.safe_load(raw) + except Exception: + return raw + + +def find_roots(scope: str) -> List[Path]: + repo_root = Path(__file__).resolve().parents[1] + if scope == "hardware": + roots = [ + repo_root / "configs" / "hardware-config", + repo_root / "validation_scripts" / "validation_configs" / "hardware-config", + ] + else: + roots = [ + repo_root / "configs" / "model-config", + repo_root / "validation_scripts" / "validation_configs" / "model-config", + ] + return [p for p in roots if p.exists()] + + +def key_from_line(line: str) -> Optional[Tuple[int, str]]: + if ":" not in line: + return None + stripped = line.split("#", 1)[0].rstrip() + if ":" not in stripped: + return None + prefix, _ = stripped.split(":", 1) + if not prefix.strip(): + return None + indent = len(prefix) - len(prefix.lstrip(" ")) + key = prefix.strip() + return indent, key + + +def locate_block(lines: List[str], path_parts: List[str]) -> Optional[Tuple[int, int]]: + idx = 0 + indent = 0 + last_idx = None + last_indent = None + for part in path_parts: + found = False + for i in range(idx, len(lines)): + res = key_from_line(lines[i]) + if res is None: + continue + ind, key = res + if ind < indent: + return None + if ind == indent and key == part: + last_idx, last_indent = i, ind + idx = i + 1 + indent = ind + 2 + found = True + break + if not found: + return None + if last_idx is None or last_indent is None: + return None + return last_idx, last_indent + + +def find_insertion_index(lines: List[str], parent_line_idx: int, parent_indent: int) -> int: + for i in range(parent_line_idx + 1, len(lines)): + res = key_from_line(lines[i]) + if res is None: + continue + ind, _ = res + if ind <= parent_indent: + return i + return len(lines) + + +def render_scalar(value: Any) -> str: + dumped = yaml.safe_dump(value, default_flow_style=True) + # Take only the first line to avoid document end marker ('...') + first_line = dumped.split('\n', 1)[0] + return first_line + + +def ancestors_are_mappings(data: Dict[str, Any], path_parts: List[str]) -> bool: + cursor = data + for part in path_parts[:-1]: + if not isinstance(cursor, dict): + return False + if part not in cursor: + return True + cursor = cursor[part] + return isinstance(cursor, dict) + + +def insert_block(lines: List[str], path_parts: List[str], value: Any, comment: str) -> List[str]: + # Find deepest existing ancestor + deepest_idx = -1 + deepest_indent = -2 + missing_at = 0 + for i in range(len(path_parts)): + loc = locate_block(lines, path_parts[: i + 1]) + if loc is None: + missing_at = i + break + deepest_idx, deepest_indent = loc + missing_at = i + 1 + + # All parts exist + if missing_at >= len(path_parts): + return lines + + missing_parts = path_parts[missing_at:] + insertion_idx = find_insertion_index(lines, deepest_idx, deepest_indent) if deepest_idx >= 0 else len(lines) + indent = deepest_indent + 2 + + snippet_lines: List[str] = [] + for depth, key in enumerate(missing_parts): + pad = " " * (indent + depth * 2) + is_leaf = depth == len(missing_parts) - 1 + if is_leaf: + if comment: + snippet_lines.append(f"{pad}# {comment}") + snippet_lines.append(f"{pad}{key}: {render_scalar(value)}") + else: + snippet_lines.append(f"{pad}{key}:") + + return lines[:insertion_idx] + snippet_lines + lines[insertion_idx:] + + +def update_leaf_value(lines: List[str], path_parts: List[str], value: Any) -> List[str]: + loc = locate_block(lines, path_parts) + if loc is None: + return lines + leaf_idx, leaf_indent = loc + key = path_parts[-1] + line = lines[leaf_idx] + comment = "" + if "#" in line: + _, comment = line.split("#", 1) + new_line = f"{' ' * leaf_indent}{key}: {render_scalar(value)}" + if comment: + new_line += f" #{comment.lstrip()}" + lines[leaf_idx] = new_line + return lines + + +def set_field(path: Path, dotted_path: str, value: Any, comment: str = "", dry_run: bool = False) -> str: + path_parts = dotted_path.split(".") + try: + text = path.read_text() + except Exception as exc: + return f"[SKIP] {path}: read error: {exc}" + + lines = [ln for ln in text.splitlines() if ln.strip() != "..."] + + try: + data = yaml.safe_load("\n".join(lines)) or {} + except Exception as exc: + return f"[SKIP] {path}: yaml parse error: {exc}" + + if not isinstance(data, dict): + return f"[SKIP] {path}: top-level is not a mapping" + + if not ancestors_are_mappings(data, path_parts): + return f"[SKIP] {path}: ancestor path is missing or not a mapping; not inserting" + + exists = True + cursor: Any = data + for part in path_parts: + if isinstance(cursor, dict) and part in cursor: + cursor = cursor[part] + else: + exists = False + break + + if exists: + new_lines = update_leaf_value(lines, path_parts, value) + else: + new_lines = insert_block(lines, path_parts, value, comment) + + new_text = "\n".join(new_lines) + ("\n" if text.endswith("\n") else "") + if dry_run: + return f"[DRY-RUN] {path}" + try: + path.write_text(new_text) + except Exception as exc: + return f"[SKIP] {path}: write error: {exc}" + return f"[UPDATED] {path}" + + +def main() -> None: + args = parse_args() + value = parse_value(args.value) + roots = find_roots(args.scope) + path_parts = args.path + modified = 0 + skipped = 0 + for root in roots: + for yaml_path in sorted(root.rglob("*.yaml")): + result = set_field(yaml_path, path_parts, value, args.comment, args.dry_run) + if result.startswith("[UPDATED]") or result.startswith("[DRY-RUN]"): + modified += 1 + else: + skipped += 1 + print(result) + print(f"Done. Modified: {modified}, skipped: {skipped}") + + +if __name__ == "__main__": + main() diff --git a/tools/overlap_finder.py b/tools/overlap_finder.py new file mode 100644 index 00000000..19d0659b --- /dev/null +++ b/tools/overlap_finder.py @@ -0,0 +1,243 @@ +#!/usr/bin/env python3 +""" +Lightweight overlap tuner using golden-section search (no SciPy). + +Global knobs (edit below): + - RUN_INFERENCE / RUN_TRAINING: enable which suites to optimize + - INF_DEVICES / TRAIN_DEVICES: which devices to run from validation CSVs + - GOLD_TOL / GOLD_MAX_ITER: search tolerance and iteration cap + +Current behavior: + - Inference: optimize network.overlap.tp_overlap using validation_scripts/nvidia_inf.py + - Training: optimize network.overlap.tp_sp_overlap using validation_scripts/nvidia_train.py + +Extensible: add more evaluators that return an average pct error for a target overlap key. +""" + +from __future__ import annotations + +import io +import math +import sys +from contextlib import redirect_stdout +from dataclasses import dataclass +from pathlib import Path +from typing import Any, Callable, Dict, Iterable, List, Mapping, Optional, Sequence, Tuple + +# Ensure project root on path +PROJECT_ROOT = Path(__file__).resolve().parents[1] +if str(PROJECT_ROOT) not in sys.path: + sys.path.insert(0, str(PROJECT_ROOT)) + +from validation_scripts import nvidia_inf, nvidia_train # type: ignore +from validation_scripts.validation_helpers import ValidationSpec, run_validation_suite # type: ignore +from tools import mod_field + +# ======= Global config ======= +RUN_INFERENCE = True +RUN_TRAINING = True +# Set to True to write best overlaps back to hardware config files (rounded to 2 decimals). +APPLY_BEST = True + +INF_DEVICES: Sequence[str] = ("A100",) +TRAIN_DEVICES: Sequence[str] = ("A100_korthi",) + +GOLD_TOL = 0.01 # stop when interval width < GOLD_TOL +GOLD_MAX_ITER = 20 + +# ============================= + + +def _merge_dict(target: Dict[str, Any], updates: Mapping[str, Any]) -> Dict[str, Any]: + for key, val in updates.items(): + if isinstance(val, Mapping): + current = target.get(key, {}) + if not isinstance(current, dict): + current = {} + target[key] = _merge_dict(dict(current), val) + else: + target[key] = val + return target + + +def _with_overlap(spec: ValidationSpec, overlap_key: str, overlap_value: float) -> ValidationSpec: + base_overrides: Dict[str, Any] = {} + if spec.hardware_overrides and isinstance(spec.hardware_overrides, Mapping): + base_overrides = _merge_dict({}, spec.hardware_overrides) + base_overrides = _merge_dict(base_overrides, {"network": {"overlap": {overlap_key: float(overlap_value)}}}) + return ValidationSpec( + label=spec.label, + model_overrides=spec.model_overrides, + hardware_overrides=base_overrides, + metadata=spec.metadata, + model_config_path=spec.model_config_path, + hardware_config_path=spec.hardware_config_path, + order=spec.order, + ) + + +def _avg_pct_error(errors: Iterable[float]) -> float: + vals = [e for e in errors if not math.isnan(e)] + if not vals: + return float("inf") + return sum(vals) / len(vals) + + +def _all_hardware_configs() -> List[Path]: + roots = [ + PROJECT_ROOT / "configs" / "hardware-config", + PROJECT_ROOT / "validation_scripts" / "validation_configs" / "hardware-config", + ] + paths: List[Path] = [] + for root in roots: + if root.exists(): + paths.extend(sorted(root.glob("*.yaml"))) + return paths + + +def eval_inference(tp_overlap: float) -> float: + """Average pct error across INF_DEVICES for given tp_overlap.""" + all_errors: List[float] = [] + for device in INF_DEVICES: + specs, actual_lookup, base_model_path, base_hw_path = nvidia_inf.build_specs_for_device( + device, network_ignored=True + ) + specs = [_with_overlap(spec, "tp_overlap", tp_overlap) for spec in specs] + with io.StringIO() as buf, redirect_stdout(buf): + results = run_validation_suite( + specs, + base_model_config_path=base_model_path, + base_hardware_config_path=base_hw_path, + result_parser=nvidia_inf.parse_inference_time, + run_perf_path=nvidia_inf.RUN_PERF, + ) + rows = nvidia_inf.compute_pct_errors(results, actual_lookup) + all_errors.extend(float(row["pct_error"]) for row in rows) + return _avg_pct_error(all_errors) + + +def eval_training(tp_sp_overlap: float) -> float: + """Average pct error across TRAIN_DEVICES for given tp_sp_overlap.""" + all_errors: List[float] = [] + for device in TRAIN_DEVICES: + specs, actual_lookup, base_model_path, base_hw_path = nvidia_train.build_specs_for_device(device) + specs = [_with_overlap(spec, "tp_sp_overlap", tp_sp_overlap) for spec in specs] + with io.StringIO() as buf, redirect_stdout(buf): + results = run_validation_suite( + specs, + base_model_config_path=base_model_path, + base_hardware_config_path=base_hw_path, + result_parser=nvidia_train.parse_training_time, + run_perf_path=nvidia_train.RUN_PERF, + ) + rows = nvidia_train.compute_pct_errors(results, actual_lookup) + all_errors.extend(float(row["pct_error"]) for row in rows) + return _avg_pct_error(all_errors) + + +@dataclass +class SearchResult: + best_x: float + best_y: float + evaluations: List[Tuple[float, float]] + + +def golden_section_search( + fn: Callable[[float], float], + lo: float = 0.0, + hi: float = 1.0, + tol: float = GOLD_TOL, + max_iter: int = GOLD_MAX_ITER, + log_fn: Optional[Callable[[int, int, float, float], None]] = None, +) -> SearchResult: + phi = (1 + 5 ** 0.5) / 2 + inv_phi = 1 / phi + inv_phi_sq = inv_phi ** 2 + + a, b = lo, hi + h = b - a + if h <= tol: + mid = (a + b) / 2 + val = fn(mid) + return SearchResult(best_x=mid, best_y=val, evaluations=[(mid, val)]) + + n = int(math.ceil(math.log(tol / h) / math.log(inv_phi))) + c = a + inv_phi_sq * h + d = a + inv_phi * h + yc = fn(c) + yd = fn(d) + if log_fn: + log_fn(1, max_iter, c, yc) + log_fn(2, max_iter, d, yd) + evals: List[Tuple[float, float]] = [(c, yc), (d, yd)] + + iter_idx = 2 + for _ in range(min(n, max_iter)): + if yc < yd: + b = d + d, yd = c, yc + h = inv_phi * h + c = a + inv_phi_sq * h + yc = fn(c) + evals.append((c, yc)) + else: + a = c + c, yc = d, yd + h = inv_phi * h + d = a + inv_phi * h + yd = fn(d) + evals.append((d, yd)) + iter_idx += 1 + if log_fn: + log_fn(iter_idx, max_iter, evals[-1][0], evals[-1][1]) + if h < tol: + break + + best_x, best_y = min(evals, key=lambda t: t[1]) + return SearchResult(best_x=best_x, best_y=best_y, evaluations=evals) + + +def main() -> None: + if RUN_INFERENCE: + print("=== Optimizing tp_overlap for inference ===") + res = golden_section_search( + eval_inference, + log_fn=lambda i, m, x, y: print(f"[{i}/{m}] Tried {x:.5f} --> {y:.2f}"), + ) + print(f"Best tp_overlap={res.best_x:.4f}, avg pct error={res.best_y:.4f}") + for x, y in res.evaluations: + print(f" tried {x:.4f} -> {y:.4f}") + err0 = eval_inference(0.0) + err1 = eval_inference(1.0) + print(f" error at 0.0: {err0:.4f}") + print(f" error at 1.0: {err1:.4f}") + if APPLY_BEST: + best_val = round(res.best_x, 2) + print(f"Applying tp_overlap={best_val:.2f} to all hardware configs") + for cfg_path in _all_hardware_configs(): + msg = mod_field.set_field(cfg_path, "network.overlap.tp_overlap", best_val, dry_run=False) + print(f" {msg}") + + if RUN_TRAINING: + print("\n=== Optimizing tp_sp_overlap for training ===") + res = golden_section_search( + eval_training, + log_fn=lambda i, m, x, y: print(f"[{i}/{m}] Tried {x:.5f} --> {y:.2f}"), + ) + print(f"Best tp_sp_overlap={res.best_x:.4f}, avg pct error={res.best_y:.4f}") + for x, y in res.evaluations: + print(f" tried {x:.4f} -> {y:.4f}") + err0 = eval_training(0.0) + err1 = eval_training(1.0) + print(f" error at 0.0: {err0:.4f}") + print(f" error at 1.0: {err1:.4f}") + if APPLY_BEST: + best_val = round(res.best_x, 2) + print(f"Applying tp_sp_overlap={best_val:.2f} to all hardware configs") + for cfg_path in _all_hardware_configs(): + msg = mod_field.set_field(cfg_path, "network.overlap.tp_sp_overlap", best_val, dry_run=False) + print(f" {msg}") + + +if __name__ == "__main__": + main() diff --git a/validation_scripts/validation_configs/hardware-config/H100_SXM5_80GB.yaml b/validation_scripts/validation_configs/hardware-config/H100_SXM5_80GB.yaml index be3fbdfc..f7db00c0 100644 --- a/validation_scripts/validation_configs/hardware-config/H100_SXM5_80GB.yaml +++ b/validation_scripts/validation_configs/hardware-config/H100_SXM5_80GB.yaml @@ -127,6 +127,13 @@ network: parallelisms: [] collective_override: {} # overrides execution_backend.collectives when provided + overlap: + # fraction of TP+SP compute that can overlap TP/SP comm + tp_sp_overlap: 0.63 + # fraction of TP compute that can overlap TP comm + tp_overlap: 0.6 + # fraction of CP comm that can overlap CP compute + cp_overlap: 0.0 execution_backend: model: analytical # analytical | astra astra: diff --git a/validation_scripts/validation_configs/hardware-config/a100_80GB.yaml b/validation_scripts/validation_configs/hardware-config/a100_80GB.yaml index 7e27d5a5..8bff7cbc 100644 --- a/validation_scripts/validation_configs/hardware-config/a100_80GB.yaml +++ b/validation_scripts/validation_configs/hardware-config/a100_80GB.yaml @@ -122,6 +122,13 @@ network: parallelisms: [] collective_override: {} # overrides execution_backend.collectives when provided + overlap: + # fraction of TP+SP compute that can overlap TP/SP comm + tp_sp_overlap: 0.63 + # fraction of TP compute that can overlap TP comm + tp_overlap: 0.6 + # fraction of CP comm that can overlap CP compute + cp_overlap: 0.0 execution_backend: model: astra # analytical | astra astra: diff --git a/validation_scripts/validation_configs/hardware-config/a100_80GB_korthikanti.yaml b/validation_scripts/validation_configs/hardware-config/a100_80GB_korthikanti.yaml index a9709e7a..e4ed3d89 100644 --- a/validation_scripts/validation_configs/hardware-config/a100_80GB_korthikanti.yaml +++ b/validation_scripts/validation_configs/hardware-config/a100_80GB_korthikanti.yaml @@ -39,7 +39,7 @@ tech_param: bandwidth: 1986 GB # (/s) dynamic_energy_per_bit: 4.5e-12 # (J) latency: 100e-9 # (s) - util: 1.0 + util: 0.7 SRAM-L2: size: 40 MB bandwidth: 7050 GB @@ -83,6 +83,10 @@ parallelism: mb: 1 # micro-batches (for best results set to >= lp). network: + overlap: + tp_overlap: 0.6 + tp_sp_overlap: 0.63 + cp_overlap: 0.8 # Global faulty links apply to the active AstraSim dimension: [[src, dst, weight]] # faulty_links: [[0,1,0.6],[16,17,0.1]] dimensions: @@ -92,7 +96,7 @@ network: size: auto topology: type: FullyConnected - bandwidth: 300 GB # based on num nvlinks per GPU * 25 GB/s per link + bandwidth: 125 GB # based on num nvlinks per GPU * 25 GB/s per link + 50GB/s derate + uni energy_per_bit: 8e-12 # (J/bit) latency: 1e-6 util: 1.0 @@ -126,7 +130,7 @@ execution_backend: model: astra # analytical | astra astra: backend: astra - mode: full_astrasim_flattened # hybrid | full_astrasim_hierarchical | full_astrasim_flattened + mode: full_astrasim_hierarchical # hybrid | full_astrasim_hierarchical | full_astrasim_flattened collectives: # Default collective algorithms. Per-dimension collective_override entries can override these. # Allowed values per op: auto | ring | direct | halvingDoubling | doubleBinaryTree | mesh | mesh2d | hypercube | torus2d diff --git a/validation_scripts/validation_configs/hardware-config/a100_80GB_no_parallelism.yaml b/validation_scripts/validation_configs/hardware-config/a100_80GB_no_parallelism.yaml index 798101a2..6810ae40 100644 --- a/validation_scripts/validation_configs/hardware-config/a100_80GB_no_parallelism.yaml +++ b/validation_scripts/validation_configs/hardware-config/a100_80GB_no_parallelism.yaml @@ -110,6 +110,13 @@ network: + overlap: + # fraction of TP+SP compute that can overlap TP/SP comm + tp_sp_overlap: 0.63 + # fraction of TP compute that can overlap TP comm + tp_overlap: 0.6 + # fraction of CP comm that can overlap CP compute + cp_overlap: 0.0 execution_backend: model: analytical # analytical | astra astra: diff --git a/validation_scripts/validation_configs/hardware-config/a100_80GB_selene_sc.yaml b/validation_scripts/validation_configs/hardware-config/a100_80GB_selene_sc.yaml index 8aad58ad..4110819d 100644 --- a/validation_scripts/validation_configs/hardware-config/a100_80GB_selene_sc.yaml +++ b/validation_scripts/validation_configs/hardware-config/a100_80GB_selene_sc.yaml @@ -122,6 +122,13 @@ network: parallelisms: [] collective_override: {} # overrides execution_backend.collectives when provided + overlap: + # fraction of TP+SP compute that can overlap TP/SP comm + tp_sp_overlap: 0.63 + # fraction of TP compute that can overlap TP comm + tp_overlap: 0.6 + # fraction of CP comm that can overlap CP compute + cp_overlap: 0.0 execution_backend: model: astra # analytical | astra astra: From ccb03ebc30f6ff93f0b9b073eed2ff8ea49a4fb8 Mon Sep 17 00:00:00 2001 From: NanoMembers Date: Sat, 22 Nov 2025 21:49:41 -0800 Subject: [PATCH 64/77] implemented full recomputation for hybrid/hierachical mode --- LLM_util.py | 6 ++++-- config.py | 15 ++++++++++++++- configs/hardware-config/H100_SXM5_80GB.yaml | 1 + configs/hardware-config/a100_80GB.yaml | 3 ++- .../hardware-config/a100_80GB_no_parallelism.yaml | 1 + configs/hardware-config/a100_80GB_problem.yaml | 1 + configs/hardware-config/gmap_mesh2d_demo.yaml | 1 + configs/hardware-config/v100.yaml | 1 + .../hardware-config/waferscale_20v100_80hbm.yaml | 1 + time_calculation.py | 1 + time_calculation_LLM.py | 15 ++++++++++++--- .../hardware-config/H100_SXM5_80GB.yaml | 1 + .../hardware-config/a100_80GB.yaml | 1 + .../hardware-config/a100_80GB_korthikanti.yaml | 1 + .../hardware-config/a100_80GB_no_parallelism.yaml | 1 + .../hardware-config/a100_80GB_selene_sc.yaml | 1 + 16 files changed, 44 insertions(+), 7 deletions(-) diff --git a/LLM_util.py b/LLM_util.py index 110af2bf..18bd3a00 100644 --- a/LLM_util.py +++ b/LLM_util.py @@ -123,12 +123,14 @@ def process_gemm_shapes(self, batch_size, seq_len, d_model, num_heads, kv_heads, return processed -def get_transformer_mem_layer( dp, tp, lp, mb, batch_size, hidden_dim, seq_len, intermediate_size, n_heads, precision, zero_stage, flash_attention, model_type="gpt"):# https://arxiv.org/pdf/2205.05198. https://shjwudp.github.io/blog/2023/gpt-training-memory-estimation-nemo-training-practice/ +def get_transformer_mem_layer( dp, tp, lp, mb, batch_size, hidden_dim, seq_len, intermediate_size, n_heads, precision, zero_stage, flash_attention, full_recomputation, model_type="gpt"):# https://arxiv.org/pdf/2205.05198. https://shjwudp.github.io/blog/2023/gpt-training-memory-estimation-nemo-training-practice/ """ memory estimation of transformer layer for single gpu case in inference mode is supported. other modes or layers are work in progress.""" #Activations refer to output activations that need to be stored - if flash_attention: + if full_recomputation: + act_memory_layer = seq_len * batch_size * hidden_dim * (2 ) * (precision.activations / 2) #full recompute: + elif flash_attention: act_memory_layer = seq_len * batch_size * hidden_dim * (34 / tp ) * (precision.activations / 2) #assuming selective recompute else: act_memory_layer = seq_len * batch_size * hidden_dim * (34 / tp + 5 * n_heads * seq_len/(hidden_dim * tp) ) * (precision.activations / 2) diff --git a/config.py b/config.py index 684e3862..f8443ad9 100644 --- a/config.py +++ b/config.py @@ -861,7 +861,14 @@ def grad_accumulation_steps(self) -> int: ], ) SWConfig = _namedtuple( - "sw_param", ["kernel_launch_overhead", "precision", "h2d_bandwidth", "dp_zero_stage"] + "sw_param", + [ + "kernel_launch_overhead", + "precision", + "h2d_bandwidth", + "dp_zero_stage", + "full_recomputation", + ], ) SchedulingConfig = _namedtuple( @@ -1108,11 +1115,17 @@ def parse_config(filename, config_type): kernel_launch_overhead = sw_block["kernel_launch_overhead"] h2d_bandwidth = sw_block["h2d_bandwidth"] dp_zero_stage = sw_block["dp_zero_stage"] + full_recomp_raw = sw_block.get("full_recomputation", False) + if isinstance(full_recomp_raw, str): + full_recomputation = full_recomp_raw.strip().lower() in {"1", "true", "yes", "y", "on", "full"} + else: + full_recomputation = bool(full_recomp_raw) sw_config = SWConfig( kernel_launch_overhead=kernel_launch_overhead, precision=precision_config, h2d_bandwidth=h2d_bandwidth, dp_zero_stage=dp_zero_stage, + full_recomputation=full_recomputation, ) raw_parallelism_block = config_dict.get("parallelism", {}) if not isinstance(raw_parallelism_block, dict): diff --git a/configs/hardware-config/H100_SXM5_80GB.yaml b/configs/hardware-config/H100_SXM5_80GB.yaml index f7db00c0..de3dfc12 100644 --- a/configs/hardware-config/H100_SXM5_80GB.yaml +++ b/configs/hardware-config/H100_SXM5_80GB.yaml @@ -5,6 +5,7 @@ # [Semi] https://semianalysis.com/2025/06/23/nvidia-tensor-core-evolution-from-volta-to-blackwell/#4th-generation-tensor-core-warpgroup-level-asynchronous-mma sw_param: + full_recomputation: false # Enable full activation recompute during backward instead of selective recomputation kernel_launch_overhead: 9e-6 # (s) # dp_zero_stage selects ZeRO-style data-parallel sharding: 0=DDP, 1=ZeRO-1, 2=ZeRO-2, 3=ZeRO-3 # only used if data parallelism (dp) > 1 diff --git a/configs/hardware-config/a100_80GB.yaml b/configs/hardware-config/a100_80GB.yaml index dd80a02f..7551dfea 100644 --- a/configs/hardware-config/a100_80GB.yaml +++ b/configs/hardware-config/a100_80GB.yaml @@ -1,4 +1,5 @@ sw_param: + full_recomputation: true # Enable full activation recompute during backward instead of selective recomputation kernel_launch_overhead: 9e-6 # (s) # dp_zero_stage selects ZeRO-style data-parallel sharding: 0=DDP, 1=ZeRO-1, 2=ZeRO-2, 3=ZeRO-3 # only used if data parallelism (dp) > 1 @@ -130,7 +131,7 @@ execution_backend: model: astra # analytical | astra astra: backend: astra - mode: full_astrasim_flattened # hybrid | full_astrasim_hierarchical | full_astrasim_flattened + mode: full_astrasim_hierarchica # hybrid | full_astrasim_hierarchical | full_astrasim_flattened collectives: # Default collective algorithms. Per-dimension collective_override entries can override these. # Allowed values per op: auto | ring | direct | halvingDoubling | doubleBinaryTree | mesh | mesh2d | hypercube | torus2d diff --git a/configs/hardware-config/a100_80GB_no_parallelism.yaml b/configs/hardware-config/a100_80GB_no_parallelism.yaml index 6810ae40..be423bb2 100644 --- a/configs/hardware-config/a100_80GB_no_parallelism.yaml +++ b/configs/hardware-config/a100_80GB_no_parallelism.yaml @@ -2,6 +2,7 @@ # DO NOT MODIFY. sw_param: + full_recomputation: false # Enable full activation recompute during backward instead of selective recomputation kernel_launch_overhead: 9e-6 # dp_zero_stage selects ZeRO-style data-parallel sharding: 0=DDP, 1=ZeRO-1, 2=ZeRO-2, 3=ZeRO-3 # only used if data parallelism (dp) > 1 diff --git a/configs/hardware-config/a100_80GB_problem.yaml b/configs/hardware-config/a100_80GB_problem.yaml index 333a61af..e284e5df 100644 --- a/configs/hardware-config/a100_80GB_problem.yaml +++ b/configs/hardware-config/a100_80GB_problem.yaml @@ -1,4 +1,5 @@ sw_param: + full_recomputation: false # Enable full activation recompute during backward instead of selective recomputation kernel_launch_overhead: 9e-6 # (s) # dp_zero_stage selects ZeRO-style data-parallel sharding: 0=DDP, 1=ZeRO-1, 2=ZeRO-2, 3=ZeRO-3 # only used if data parallelism (dp) > 1 diff --git a/configs/hardware-config/gmap_mesh2d_demo.yaml b/configs/hardware-config/gmap_mesh2d_demo.yaml index ce38e619..e3fb8e9a 100644 --- a/configs/hardware-config/gmap_mesh2d_demo.yaml +++ b/configs/hardware-config/gmap_mesh2d_demo.yaml @@ -1,4 +1,5 @@ sw_param: + full_recomputation: false # Enable full activation recompute during backward instead of selective recomputation kernel_launch_overhead: 5e-6 dp_zero_stage: 0 precision: diff --git a/configs/hardware-config/v100.yaml b/configs/hardware-config/v100.yaml index 6e390c71..d185ab8e 100644 --- a/configs/hardware-config/v100.yaml +++ b/configs/hardware-config/v100.yaml @@ -1,5 +1,6 @@ sw_param: + full_recomputation: false # Enable full activation recompute during backward instead of selective recomputation kernel_launch_overhead: 9e-6 #9us # dp_zero_stage selects ZeRO-style data-parallel sharding: 0=DDP, 1=ZeRO-1, 2=ZeRO-2, 3=ZeRO-3 # only used if data parallelism (dp) > 1 diff --git a/configs/hardware-config/waferscale_20v100_80hbm.yaml b/configs/hardware-config/waferscale_20v100_80hbm.yaml index 1e9649dc..2812d62b 100644 --- a/configs/hardware-config/waferscale_20v100_80hbm.yaml +++ b/configs/hardware-config/waferscale_20v100_80hbm.yaml @@ -3,6 +3,7 @@ sw_param: + full_recomputation: false # Enable full activation recompute during backward instead of selective recomputation kernel_launch_overhead: 9e-6 #9us # dp_zero_stage selects ZeRO-style data-parallel sharding: 0=DDP, 1=ZeRO-1, 2=ZeRO-2, 3=ZeRO-3 # only used if data parallelism (dp) > 1 diff --git a/time_calculation.py b/time_calculation.py index 29f715b5..5cbc5ecc 100644 --- a/time_calculation.py +++ b/time_calculation.py @@ -203,6 +203,7 @@ def __init__(self, hw_config, model_config, mode, *, astra_policy_override: Opti # If we want proper backported LSTM support we need to fix this at *some* point. self.h2d_bandwidth = getattr(hw_config.sw_config, "h2d_bandwidth", -1) self.zero_stage = getattr(hw_config.sw_config, "dp_zero_stage", 0) + self.full_recomputation = getattr(hw_config.sw_config, "full_recomputation", False) self.attached = True # Hardware Parameters diff --git a/time_calculation_LLM.py b/time_calculation_LLM.py index a9afa074..7a2a1a9c 100644 --- a/time_calculation_LLM.py +++ b/time_calculation_LLM.py @@ -471,6 +471,7 @@ def flash_attention_kernel_backward(self, batch_size, hidden_dim, seq_len, num_h attention_backward_reduction_time = 0 attention_size_b = 0 + recompute_time = 0 # assuming key and value are preloaded into shared memory before attention backward computation load_kv_bytes = Bc * d * 2 * num_SMs * self.precision.activations #load key and value for one tile from HBM to SRAM initial_load_time = self.roofline(0, load_kv_bytes, "flash_attention_initial_load_b", mem_level=self.num_levels - 1) #assume key and value of one attention head is loaded from HBM to SRAM @@ -509,8 +510,11 @@ def flash_attention_kernel_backward(self, batch_size, hidden_dim, seq_len, num_h act_dP_time *= batch_size * num_heads / max(1, self.tp) + self.O softmax_time_backward *= batch_size * num_heads / max(1, self.tp) act_dQ_time *= batch_size * num_heads / max(1, self.tp) + self.O - - attention_backward_gemm_time = initial_load_time + attn_score_time + attn_scale_softmax_time + act_dO_time + act_dP_time + softmax_time_backward + act_dQ_time + if self.full_recomputation: #attention recompute is already included in full recomputation + recompute_time = 0 + else: + recompute_time = attn_score_time + attn_scale_softmax_time #selective recomputation only recompute attention score and softmax + attention_backward_gemm_time = initial_load_time + recompute_time + act_dO_time + act_dP_time + softmax_time_backward + act_dQ_time attention_backward_time = attention_backward_gemm_time + attention_backward_reduction_time @@ -1926,6 +1930,9 @@ def _extract_forward(ret: Sequence[Any]) -> Tuple[float, float, float, float, An + transformer_timings["layernorm1"].total_backward_time() + transformer_timings["layernorm2"].total_backward_time() ) + # Full recomputation during backward requires redoing the forward compute. + if getattr(self, "full_recomputation", False): + transformer_time_b += transformer_time_f node_breakdown = { "transformer_time_f": transformer_time_f, @@ -2322,7 +2329,8 @@ def _resolve_operation(name: str) -> OperationTiming: ) transformer_forward_root = transformer_graph.construct_transformer_graph(direction="forward") if include_transformer_backward: - transformer_backward_root = transformer_graph.construct_transformer_graph(direction="backward") + bwd_direction = "both" if getattr(self, "full_recomputation", False) else "backward" + transformer_backward_root = transformer_graph.construct_transformer_graph(direction=bwd_direction) transformer_forward_root = apply_overlap_transforms( transformer_forward_root, @@ -2490,6 +2498,7 @@ def calc_time_llm(self): model_type=self.model_type, zero_stage=self.zero_stage, flash_attention=self.flash_attention, + full_recomputation=self.full_recomputation, ) ) diff --git a/validation_scripts/validation_configs/hardware-config/H100_SXM5_80GB.yaml b/validation_scripts/validation_configs/hardware-config/H100_SXM5_80GB.yaml index f7db00c0..de3dfc12 100644 --- a/validation_scripts/validation_configs/hardware-config/H100_SXM5_80GB.yaml +++ b/validation_scripts/validation_configs/hardware-config/H100_SXM5_80GB.yaml @@ -5,6 +5,7 @@ # [Semi] https://semianalysis.com/2025/06/23/nvidia-tensor-core-evolution-from-volta-to-blackwell/#4th-generation-tensor-core-warpgroup-level-asynchronous-mma sw_param: + full_recomputation: false # Enable full activation recompute during backward instead of selective recomputation kernel_launch_overhead: 9e-6 # (s) # dp_zero_stage selects ZeRO-style data-parallel sharding: 0=DDP, 1=ZeRO-1, 2=ZeRO-2, 3=ZeRO-3 # only used if data parallelism (dp) > 1 diff --git a/validation_scripts/validation_configs/hardware-config/a100_80GB.yaml b/validation_scripts/validation_configs/hardware-config/a100_80GB.yaml index 8bff7cbc..96d8361d 100644 --- a/validation_scripts/validation_configs/hardware-config/a100_80GB.yaml +++ b/validation_scripts/validation_configs/hardware-config/a100_80GB.yaml @@ -1,4 +1,5 @@ sw_param: + full_recomputation: false # Enable full activation recompute during backward instead of selective recomputation kernel_launch_overhead: 9e-6 # (s) # dp_zero_stage selects ZeRO-style data-parallel sharding: 0=DDP, 1=ZeRO-1, 2=ZeRO-2, 3=ZeRO-3 # only used if data parallelism (dp) > 1 diff --git a/validation_scripts/validation_configs/hardware-config/a100_80GB_korthikanti.yaml b/validation_scripts/validation_configs/hardware-config/a100_80GB_korthikanti.yaml index e4ed3d89..c5bac124 100644 --- a/validation_scripts/validation_configs/hardware-config/a100_80GB_korthikanti.yaml +++ b/validation_scripts/validation_configs/hardware-config/a100_80GB_korthikanti.yaml @@ -1,4 +1,5 @@ sw_param: + full_recomputation: false # Enable full activation recompute during backward instead of selective recomputation kernel_launch_overhead: 9e-6 # (s) # dp_zero_stage selects ZeRO-style data-parallel sharding: 0=DDP, 1=ZeRO-1, 2=ZeRO-2, 3=ZeRO-3 # only used if data parallelism (dp) > 1 diff --git a/validation_scripts/validation_configs/hardware-config/a100_80GB_no_parallelism.yaml b/validation_scripts/validation_configs/hardware-config/a100_80GB_no_parallelism.yaml index 6810ae40..be423bb2 100644 --- a/validation_scripts/validation_configs/hardware-config/a100_80GB_no_parallelism.yaml +++ b/validation_scripts/validation_configs/hardware-config/a100_80GB_no_parallelism.yaml @@ -2,6 +2,7 @@ # DO NOT MODIFY. sw_param: + full_recomputation: false # Enable full activation recompute during backward instead of selective recomputation kernel_launch_overhead: 9e-6 # dp_zero_stage selects ZeRO-style data-parallel sharding: 0=DDP, 1=ZeRO-1, 2=ZeRO-2, 3=ZeRO-3 # only used if data parallelism (dp) > 1 diff --git a/validation_scripts/validation_configs/hardware-config/a100_80GB_selene_sc.yaml b/validation_scripts/validation_configs/hardware-config/a100_80GB_selene_sc.yaml index 4110819d..78ab94cf 100644 --- a/validation_scripts/validation_configs/hardware-config/a100_80GB_selene_sc.yaml +++ b/validation_scripts/validation_configs/hardware-config/a100_80GB_selene_sc.yaml @@ -1,4 +1,5 @@ sw_param: + full_recomputation: false # Enable full activation recompute during backward instead of selective recomputation kernel_launch_overhead: 9e-6 # (s) # dp_zero_stage selects ZeRO-style data-parallel sharding: 0=DDP, 1=ZeRO-1, 2=ZeRO-2, 3=ZeRO-3 # only used if data parallelism (dp) > 1 From 7a113784c480e29a3eb8d3f657366b54d607cb45 Mon Sep 17 00:00:00 2001 From: NanoMembers Date: Tue, 25 Nov 2025 19:16:58 -0800 Subject: [PATCH 65/77] implemented full-recomputation for all modes with pipeline style recompute block --- simulate_LLM.py | 86 ++++++++++++++++++++++++++++++----------- time_calculation_LLM.py | 20 +++++++--- 2 files changed, 79 insertions(+), 27 deletions(-) diff --git a/simulate_LLM.py b/simulate_LLM.py index 08c345cc..f54ec259 100644 --- a/simulate_LLM.py +++ b/simulate_LLM.py @@ -129,6 +129,9 @@ def __init__( self.comp_times = comp_times or {} self.comm_metadata = comm_metadata or {} self.misc_metadata = misc_metadata or {} + self.full_recomputation = bool(self.misc_metadata.get("full_recomputation", False)) + self.flattened_mode = bool(self.misc_metadata.get("flattened_mode", False)) + self.pipeline_style_recompute = bool(self.misc_metadata.get("pipeline_style_recompute", False)) self.num_batch = self.misc_metadata.get("num_batch", 0) self.num_layer = self.misc_metadata.get("num_layer", 0) @@ -586,9 +589,6 @@ def construct_fwd_bwd_graph(self, include_backward: bool = True, include_optimiz layer_entry_nodes = [[] for _ in range(self.num_batch)] # lists of entry nodes per layer layer_exit_nodes = [[] for _ in range(self.num_batch)] # transformer nodes per layer - embedding_b_time = self._time("embedding_b") - linear_softmax_b_time = self._time("linear_softmax_b") - transformer_b_time = self._time("transformer_b") linear_softmax_f_time = self._time("linear_softmax_f") linear_softmax_b_time = self._time("linear_softmax_b") transformer_f_time = self._time("transformer_f") @@ -596,6 +596,11 @@ def construct_fwd_bwd_graph(self, include_backward: bool = True, include_optimiz embedding_f_time = self._time("embedding_f") embedding_b_time = self._time("embedding_b") cross_layer_time = self._time("cross_layer_f") + recompute_enabled = include_backward and self.full_recomputation and ( + self.flattened_mode or self.pipeline_style_recompute + ) + transformer_b_base_time = transformer_b_time + def attach_parallel_edge(target, gather_edge, skip_non_comm_children=None, skip_comm_children=None): parents = list(getattr(target, "parents", [])) @@ -622,6 +627,9 @@ def attach_parallel_edge(target, gather_edge, skip_non_comm_children=None, skip_ transformer_nodes_b = [[] for _ in range(self.num_batch)] # for b in range(self.num_batch): transformer_nodes_b[b] = [[] for _ in range(self.num_layer)] + recompute_nodes_b: Optional[List[List[Optional[Node]]]] = None + if recompute_enabled: + recompute_nodes_b = [[None for _ in range(self.num_layer)] for _ in range(self.num_batch)] op_id = 0 # operation ID, used to distinguish nodes and edges batch_id = 0 # batch ID, used to distinguish data batches @@ -726,6 +734,16 @@ def attach_parallel_edge(target, gather_edge, skip_non_comm_children=None, skip_ if not include_backward: return embedding_node[0] + def _bwd_entry_node(mb_idx: int, layer_idx: int) -> Any: + if recompute_enabled and recompute_nodes_b is not None: + candidate = recompute_nodes_b[mb_idx][layer_idx] + if candidate is not None: + return candidate + return transformer_nodes_b[mb_idx][layer_idx] + + def _bwd_exit_node(mb_idx: int, layer_idx: int) -> Any: + return transformer_nodes_b[mb_idx][layer_idx] + for b in reversed(range(self.num_batch)): #connect each data batch node with corresponding nodes emb_b = Node("embedding_b", op_id, 0, embedding_b_time, fwd=False) # hw_id = 0 op_id += 1 @@ -738,17 +756,41 @@ def attach_parallel_edge(target, gather_edge, skip_non_comm_children=None, skip_ for l in reversed(range(self.num_layer)): hw_id = self._stage_for_layer(l) - transformer_node_b = Node(f"transformer_layer{l}_b", op_id, hw_id, transformer_b_time, fwd=False) + recompute_node = None + if recompute_enabled and recompute_nodes_b is not None: + recompute_node = Node( + f"transformer_layer{l}_recompute", + op_id, + hw_id, + transformer_f_time, + fwd=True, + ) + recompute_node.micro_batch_index = b + recompute_node.layer_index = l + recompute_node.direction = "forward" + recompute_node.stage_id = hw_id + op_id += 1 + recompute_nodes_b[b][l] = recompute_node + + transformer_node_b = Node( + f"transformer_layer{l}_b", + op_id, + hw_id, + transformer_b_base_time, + fwd=False, + ) transformer_node_b.micro_batch_index = b transformer_node_b.layer_index = l transformer_node_b.direction = "backward" transformer_node_b.stage_id = hw_id op_id += 1 transformer_nodes_b[b][l] = transformer_node_b + if recompute_node is not None: + recompute_node.add_child(transformer_node_b) for l in reversed(range(1, self.num_layer)): - curr_node = transformer_nodes_b[b][l] # current layer's qkv_proj. - next_ffn2 = transformer_nodes_b[b][l-1] # next layer's layernorm. + curr_node = _bwd_exit_node(b, l) # current layer's qkv_proj. + next_ffn2 = _bwd_entry_node(b, l-1) # next layer's layernorm. if curr_node.hw_id == next_ffn2.hw_id: edge = Edge("cross_layer", op_id, 0, comm_type="pipeline") else: @@ -757,7 +799,7 @@ def attach_parallel_edge(target, gather_edge, skip_non_comm_children=None, skip_ curr_node.add_child(edge); edge.add_child(next_ffn2) - qkv_0_b = transformer_nodes_b[b][0] # first layer's qkv_proj + qkv_0_b = _bwd_exit_node(b, 0) # first layer's qkv_proj if qkv_0_b.hw_id == emb_b.hw_id: edge = Edge("Emb_node0", op_id, 0, comm_type="pipeline") else: @@ -767,7 +809,7 @@ def attach_parallel_edge(target, gather_edge, skip_non_comm_children=None, skip_ edge.add_child(emb_b) - prev_layer_norm2 = transformer_nodes_b[b][self.num_layer-1] # last layer's layernorm2 + prev_layer_norm2 = _bwd_entry_node(b, self.num_layer-1) # last layer's layernorm2 if prev_layer_norm2.hw_id == softmax_node_b[b].hw_id: layernorm_Softmax = Edge("layernorm2_Softmax", op_id, 0, comm_type="pipeline") # same GPU else: @@ -831,7 +873,7 @@ def attach_parallel_edge(target, gather_edge, skip_non_comm_children=None, skip_ op_id, "transformer", is_dp=True, - local_hw_id=transformer_nodes_b[b][layer_idx].hw_id, + local_hw_id=_bwd_exit_node(b, layer_idx).hw_id, ) R_edge[b].append(reducer) op_id += 1 @@ -842,7 +884,7 @@ def attach_parallel_edge(target, gather_edge, skip_non_comm_children=None, skip_ op_id, zero2_transformer_key, is_dp=True, - local_hw_id=transformer_nodes_b[b][layer_idx].hw_id, + local_hw_id=_bwd_exit_node(b, layer_idx).hw_id, ) op_id += 1 reducer.add_child(gather_edge) @@ -926,7 +968,7 @@ def attach_parallel_edge(target, gather_edge, skip_non_comm_children=None, skip_ softmax_node_b[b].add_child(R_edge[b][-1]) embedding_node_b[b].add_child(R_edge[b][0]) for layer_idx in range(self.num_layer): - transformer_nodes_b[b][layer_idx].add_child(R_edge[b][layer_idx + 1]) + _bwd_exit_node(b, layer_idx).add_child(R_edge[b][layer_idx + 1]) last_transformer_layer = [-1] * self.lp # Initialize with -1 for all GPUs @@ -935,8 +977,8 @@ def attach_parallel_edge(target, gather_edge, skip_non_comm_children=None, skip_ # first_transformer_layer.append(0) gpu_index = self.lp - 1 for l in range(self.num_layer - 1, 0, -1): - if transformer_nodes_b[0][l].hw_id != transformer_nodes_b[0][l-1].hw_id: # Check if on different GPU - # print("Layer ", l, " is on GPU ", transformer_nodes_b[0][l].hw_id) + if _bwd_exit_node(0, l).hw_id != _bwd_exit_node(0, l-1).hw_id: # Check if on different GPU + # print("Layer ", l, " is on GPU ", _bwd_exit_node(0, l).hw_id) first_transformer_layer[gpu_index-1] = l-1 # Record first layer on each GPU last_transformer_layer[gpu_index] = l # Record last layer on each GPU gpu_index -= 1 @@ -948,18 +990,18 @@ def attach_parallel_edge(target, gather_edge, skip_non_comm_children=None, skip_ for b in range(self.num_batch-1, 0, -1): gpu_index = self.lp - 1 for l in range(self.num_layer - 1, 0, -1): - if transformer_nodes_b[b][l].hw_id != transformer_nodes_b[b][l-1].hw_id: # Check if on different GPUs + if _bwd_exit_node(b, l).hw_id != _bwd_exit_node(b, l-1).hw_id: # Check if on different GPUs # last_transformer_layer.append(l) # Record last layer on each GPU # first_transformer_layer.append(l-1) # Record first layer on each GPU - if transformer_nodes_b[b][l].hw_id == self.lp - 1: - transformer_nodes_b[b][l].add_child(softmax_node_b[b-1]) # Add dependency edge + if _bwd_exit_node(b, l).hw_id == self.lp - 1: + _bwd_exit_node(b, l).add_child(softmax_node_b[b-1]) # Add dependency edge else: - transformer_nodes_b[b][l].add_child(transformer_nodes_b[b-1][first_transformer_layer[gpu_index]]) # Add dependency edge + _bwd_exit_node(b, l).add_child(_bwd_entry_node(b-1, first_transformer_layer[gpu_index])) # Add dependency edge gpu_index -= 1 # Ensure embedding_node_b[b] is connected to the correct transformer node # if first_transformer_layer: - embedding_node_b[b].add_child(transformer_nodes_b[b-1][first_transformer_layer[0]]) # Add dependency edge + embedding_node_b[b].add_child(_bwd_entry_node(b-1, first_transformer_layer[0])) # Add dependency edge zero3_embedding_key = "zero3_embedding_gather" zero3_transformer_key = "zero3_transformer_gather" @@ -993,8 +1035,8 @@ def attach_parallel_edge(target, gather_edge, skip_non_comm_children=None, skip_ if zero3_transformer_key in self.comm_metadata and self.num_layer > 0: for layer_idx in reversed(range(self.num_layer)): - host = softmax_node_b[b] if layer_idx == self.num_layer - 1 else transformer_nodes_b[b][layer_idx + 1] - target = transformer_nodes_b[b][layer_idx] + host = softmax_node_b[b] if layer_idx == self.num_layer - 1 else _bwd_exit_node(b, layer_idx + 1) + target = _bwd_entry_node(b, layer_idx) gather_edge = self.create_comm_edge( f"{zero3_transformer_key}_b{b}_layer{layer_idx}_bwd", op_id, @@ -1014,7 +1056,7 @@ def attach_parallel_edge(target, gather_edge, skip_non_comm_children=None, skip_ if zero3_embedding_key in self.comm_metadata and self.num_layer > 0: - host = transformer_nodes_b[b][0] + host = _bwd_exit_node(b, 0) gather_edge = self.create_comm_edge( f"{zero3_embedding_key}_b{b}_bwd", op_id, @@ -1049,7 +1091,7 @@ def attach_parallel_edge(target, gather_edge, skip_non_comm_children=None, skip_ else: min_l = stage_min_layer.get(stage) if min_l is not None: - last_node = transformer_nodes_b[0][min_l] + last_node = _bwd_exit_node(0, min_l) if last_node: opt_node = Node(f"optimizer_stage{stage}", op_id, stage, optimizer_time, fwd=False) diff --git a/time_calculation_LLM.py b/time_calculation_LLM.py index 7a2a1a9c..36e8a768 100644 --- a/time_calculation_LLM.py +++ b/time_calculation_LLM.py @@ -223,6 +223,8 @@ def __init__(self, hw_config, model_config, mode, output_dir: Optional[str] = No self._debug_memory = _env_flag("DEEPFLOW_DEBUG_MEMORY") self._memory_breakdown_debug = None self.execution_mode = execution_mode + # Use pipeline-style recomputation (explicit recompute nodes in the pipeline graph). + self.pipeline_style_recompute = bool(self.full_recomputation) inference_cfg = getattr(hw_config, "inference_config", None) self.all_reduce = "every layer" @@ -240,6 +242,7 @@ def __init__(self, hw_config, model_config, mode, output_dir: Optional[str] = No self.transformer_backward_root: Optional[Any] = None self.transformer_analytical_time_forward: Optional[float] = None self.transformer_analytical_time_backward: Optional[float] = None + self.transformer_analytical_time_backward_combined: Optional[float] = None self.transformer_astrasim_time_forward: Optional[float] = None self.transformer_astrasim_time_backward: Optional[float] = None self.transformer_astrasim_per_rank_forward: Optional[List[float]] = None @@ -1930,13 +1933,13 @@ def _extract_forward(ret: Sequence[Any]) -> Tuple[float, float, float, float, An + transformer_timings["layernorm1"].total_backward_time() + transformer_timings["layernorm2"].total_backward_time() ) - # Full recomputation during backward requires redoing the forward compute. - if getattr(self, "full_recomputation", False): - transformer_time_b += transformer_time_f + # Pipeline-style recompute uses explicit recompute nodes, so keep backward time as-is. + transformer_time_b_combined = transformer_time_b node_breakdown = { "transformer_time_f": transformer_time_f, "transformer_time_b": transformer_time_b, + "transformer_time_b_combined": transformer_time_b_combined, "embedding_f": transformer_timings["embedding"].total_forward_time(), "embedding_b": transformer_timings["embedding"].total_backward_time(), "linear_softmax_f": transformer_timings["linear_softmax"].total_forward_time(), @@ -2329,7 +2332,7 @@ def _resolve_operation(name: str) -> OperationTiming: ) transformer_forward_root = transformer_graph.construct_transformer_graph(direction="forward") if include_transformer_backward: - bwd_direction = "both" if getattr(self, "full_recomputation", False) else "backward" + bwd_direction = "backward" transformer_backward_root = transformer_graph.construct_transformer_graph(direction=bwd_direction) transformer_forward_root = apply_overlap_transforms( @@ -2359,11 +2362,16 @@ def _resolve_operation(name: str) -> OperationTiming: "cross_layer_f": 0.0, "cross_layer_b": 0.0, } + flattened_mode = self.execution_mode == ExecutionMode.FULL_ASTRASIM_FLATTENED + pipeline_style_recompute_flag = bool(getattr(self, "full_recomputation", False)) misc_metadata = { "num_batch": self.mb, "num_layer": self.num_layers, "all_reduce": "every_layer", "dp_zero_stage": self.zero_stage, + "full_recomputation": self.full_recomputation, + "flattened_mode": flattened_mode, + "pipeline_style_recompute": pipeline_style_recompute_flag, } pipeline_graph_obj = simulate_LLM.Graph( @@ -2564,7 +2572,9 @@ def calc_time_llm(self): self.transformer_forward_root = transformer_forward_root self.transformer_backward_root = transformer_backward_root self.transformer_analytical_time_forward = node_breakdown['transformer_time_f'] - self.transformer_analytical_time_backward = node_breakdown['transformer_time_b'] + # Report backward including recompute overhead when enabled. + self.transformer_analytical_time_backward_combined = node_breakdown['transformer_time_b_combined'] + self.transformer_analytical_time_backward = self.transformer_analytical_time_backward_combined self.pipeline_graph = pipeline_graph_obj self.pipeline_root = graph_root From bc2fee5acf347c946f0eef1643ee52557d9772bd Mon Sep 17 00:00:00 2001 From: gkarfakis19 Date: Tue, 25 Nov 2025 21:59:20 -0800 Subject: [PATCH 66/77] small update to faraz config parser --- astrasim_lib/config_generation.py | 21 ++++++++++++++------- 1 file changed, 14 insertions(+), 7 deletions(-) diff --git a/astrasim_lib/config_generation.py b/astrasim_lib/config_generation.py index ab17d7fb..ce4e85b6 100644 --- a/astrasim_lib/config_generation.py +++ b/astrasim_lib/config_generation.py @@ -6,7 +6,7 @@ import uuid from typing import Any, Dict, List, Optional, Sequence, Tuple from types import SimpleNamespace - +import math from hw_component import Network from .bootstrap import ensure_chakra_available @@ -114,15 +114,22 @@ def _expand_network_entries(entries: Sequence[Dict[str, Any]]) -> List[Dict[str, dim_count = int(topo[-2]) if len(topo) >= 2 and topo[-2].isdigit() else 2 base_name = "Ring" if topo.startswith("Torus") else "Mesh" for dim_idx in range(dim_count): + curr_npu = ( + npus_value[dim_idx] + if isinstance(npus_value, (list, tuple)) + else npus_value + ) + if dim_count == 2: + root = int(math.isqrt(int(curr_npu))) + if root * root != int(curr_npu): + raise ValueError(f"npus ({curr_npu}) must be a perfect square for 2D topology {topo}.") + curr_npu = root + topology_name = "Mesh" if dim_count == 2 and curr_npu <= 2 else base_name expanded.append( { **entry, - "topology": base_name, - "npus": ( - npus_value[dim_idx] - if isinstance(npus_value, (list, tuple)) and dim_idx < len(npus_value) - else npus_value - ), + "topology": topology_name, + "npus": curr_npu, "bandwidth": bw_value, "latency": lat_value, } From 1ec43107998086c9bab62442c598b1219cb3dcab Mon Sep 17 00:00:00 2001 From: gkarfakis19 Date: Tue, 25 Nov 2025 22:03:21 -0800 Subject: [PATCH 67/77] updated astrasim?? --- astra-sim | 2 +- 1 file changed, 1 insertion(+), 1 deletion(-) diff --git a/astra-sim b/astra-sim index ee9dabd2..32bec13d 160000 --- a/astra-sim +++ b/astra-sim @@ -1 +1 @@ -Subproject commit ee9dabd2496ada3825ba7b388d7f4016fb213e7f +Subproject commit 32bec13d1fee017c0184edba43454c1542fb1190 From 536b764ff0b48c7436edee0a3bf782c32015bf2e Mon Sep 17 00:00:00 2001 From: gkarfakis19 Date: Tue, 25 Nov 2025 22:13:30 -0800 Subject: [PATCH 68/77] annoying error. hopefully this fixes it. --- .gitmodules | 5 ++++- 1 file changed, 4 insertions(+), 1 deletion(-) diff --git a/.gitmodules b/.gitmodules index 34f1cb50..33f7ddab 100644 --- a/.gitmodules +++ b/.gitmodules @@ -1,3 +1,6 @@ [submodule "astra-sim"] path = astra-sim - url = https://github.com/gkarfakis19/astra-sim \ No newline at end of file + url = https://github.com/gkarfakis19/astra-sim +[submodule "extern/graph_frontend/param"] + path = extern/graph_frontend/param + url = https://github.com/facebookresearch/param.git From 89881272d5b55a75569573e712dc0dfa364ce0fa Mon Sep 17 00:00:00 2001 From: gkarfakis19 Date: Thu, 4 Dec 2025 16:26:53 -0800 Subject: [PATCH 69/77] added dp microbatch sync option to fix selene tests. nvidia train tests in much better shape now. --- .gitignore | 1 + config.py | 9 ++++++++ configs/hardware-config/H100_SXM5_80GB.yaml | 2 ++ configs/hardware-config/a100_80GB.yaml | 2 ++ .../a100_80GB_no_parallelism.yaml | 2 ++ .../hardware-config/a100_80GB_problem.yaml | 2 ++ configs/hardware-config/gmap_mesh2d_demo.yaml | 2 ++ configs/hardware-config/v100.yaml | 2 ++ .../waferscale_20v100_80hbm.yaml | 2 ++ simulate_LLM.py | 5 +++++ time_calculation.py | 1 + time_calculation_LLM.py | 1 + validation_scripts/imec_data/A100_train.csv | 7 +++++++ validation_scripts/nvidia_train.py | 21 +++++++++++++------ .../hardware-config/H100_SXM5_80GB.yaml | 2 ++ .../hardware-config/a100_80GB.yaml | 2 ++ .../a100_80GB_korthikanti.yaml | 2 ++ .../a100_80GB_no_parallelism.yaml | 2 ++ .../hardware-config/a100_80GB_selene_sc.yaml | 2 ++ 19 files changed, 63 insertions(+), 6 deletions(-) diff --git a/.gitignore b/.gitignore index 62f6ee97..19c70b0c 100644 --- a/.gitignore +++ b/.gitignore @@ -1,5 +1,6 @@ *.swp *.swo +.cursor/ __pycache__/ output/ venv/ diff --git a/config.py b/config.py index f8443ad9..4f43910a 100644 --- a/config.py +++ b/config.py @@ -868,6 +868,7 @@ def grad_accumulation_steps(self) -> int: "h2d_bandwidth", "dp_zero_stage", "full_recomputation", + "dp_microbatch", ], ) @@ -1120,12 +1121,20 @@ def parse_config(filename, config_type): full_recomputation = full_recomp_raw.strip().lower() in {"1", "true", "yes", "y", "on", "full"} else: full_recomputation = bool(full_recomp_raw) + dp_microbatch_raw = sw_block.get("dp_microbatch", "every_mb") + if isinstance(dp_microbatch_raw, str): + dp_microbatch = dp_microbatch_raw.strip().lower() + else: + dp_microbatch = str(dp_microbatch_raw).strip().lower() + if dp_microbatch not in {"every_mb", "last_mb"}: + raise ValueError("sw_param.dp_microbatch must be 'every_mb' or 'last_mb'") sw_config = SWConfig( kernel_launch_overhead=kernel_launch_overhead, precision=precision_config, h2d_bandwidth=h2d_bandwidth, dp_zero_stage=dp_zero_stage, full_recomputation=full_recomputation, + dp_microbatch=dp_microbatch, ) raw_parallelism_block = config_dict.get("parallelism", {}) if not isinstance(raw_parallelism_block, dict): diff --git a/configs/hardware-config/H100_SXM5_80GB.yaml b/configs/hardware-config/H100_SXM5_80GB.yaml index de3dfc12..f5e72fbc 100644 --- a/configs/hardware-config/H100_SXM5_80GB.yaml +++ b/configs/hardware-config/H100_SXM5_80GB.yaml @@ -30,6 +30,8 @@ sw_param: # kv_cache precision may be set as bytes, dtype string, or "as_tensor_format". kv_cache: as_tensor_format h2d_bandwidth: -1 # Bytes/s; set -1 to disable H2D timing + # DP gradient sync frequency: every_mb (default) or last_mb (sync only on final microbatch; ZeRO-3 stays per-mb) + dp_microbatch: last_mb tech_param: core: operating_frequency: 1.98e9 # (Hz) GPU boost clock frequency [TePU] diff --git a/configs/hardware-config/a100_80GB.yaml b/configs/hardware-config/a100_80GB.yaml index 7551dfea..2f4503d7 100644 --- a/configs/hardware-config/a100_80GB.yaml +++ b/configs/hardware-config/a100_80GB.yaml @@ -24,6 +24,8 @@ sw_param: # kv_cache precision may be set as bytes, dtype string, or "as_tensor_format". kv_cache: as_tensor_format h2d_bandwidth: -1 # Bytes/s; set -1 to disable H2D timing + # DP gradient sync frequency: every_mb (default) or last_mb (sync only on final microbatch; ZeRO-3 stays per-mb) + dp_microbatch: last_mb tech_param: core: operating_frequency: 1.41e9 # (Hz) diff --git a/configs/hardware-config/a100_80GB_no_parallelism.yaml b/configs/hardware-config/a100_80GB_no_parallelism.yaml index be423bb2..f0476141 100644 --- a/configs/hardware-config/a100_80GB_no_parallelism.yaml +++ b/configs/hardware-config/a100_80GB_no_parallelism.yaml @@ -27,6 +27,8 @@ sw_param: # kv_cache precision may be set as bytes, dtype string, or "as_tensor_format". kv_cache: as_tensor_format h2d_bandwidth: -1 # Bytes/s; set -1 to disable H2D timing + # DP gradient sync frequency: every_mb (default) or last_mb (sync only on final microbatch; ZeRO-3 stays per-mb) + dp_microbatch: last_mb tech_param: core: operating_frequency: 1.41e9 diff --git a/configs/hardware-config/a100_80GB_problem.yaml b/configs/hardware-config/a100_80GB_problem.yaml index e284e5df..3b1bc65c 100644 --- a/configs/hardware-config/a100_80GB_problem.yaml +++ b/configs/hardware-config/a100_80GB_problem.yaml @@ -24,6 +24,8 @@ sw_param: # kv_cache precision may be set as bytes, dtype string, or "as_tensor_format". kv_cache: as_tensor_format h2d_bandwidth: -1 # Bytes/s; set -1 to disable H2D timing + # DP gradient sync frequency: every_mb (default) or last_mb (sync only on final microbatch; ZeRO-3 stays per-mb) + dp_microbatch: last_mb tech_param: core: operating_frequency: 1.41e9 # (Hz) diff --git a/configs/hardware-config/gmap_mesh2d_demo.yaml b/configs/hardware-config/gmap_mesh2d_demo.yaml index e3fb8e9a..cb62cc48 100644 --- a/configs/hardware-config/gmap_mesh2d_demo.yaml +++ b/configs/hardware-config/gmap_mesh2d_demo.yaml @@ -9,6 +9,8 @@ sw_param: kv_cache: as_tensor_format h2d_bandwidth: -1 + # DP gradient sync frequency: every_mb (default) or last_mb (sync only on final microbatch; ZeRO-3 stays per-mb) + dp_microbatch: last_mb tech_param: core: operating_frequency: 1.41e9 diff --git a/configs/hardware-config/v100.yaml b/configs/hardware-config/v100.yaml index d185ab8e..782ddac6 100644 --- a/configs/hardware-config/v100.yaml +++ b/configs/hardware-config/v100.yaml @@ -26,6 +26,8 @@ sw_param: kv_cache: as_tensor_format h2d_bandwidth: 13297674240 # Bytes/s; set -1 to disable H2D timing + # DP gradient sync frequency: every_mb (default) or last_mb (sync only on final microbatch; ZeRO-3 stays per-mb) + dp_microbatch: last_mb tech_param: core: operating_frequency: 1.33e9 #Hz diff --git a/configs/hardware-config/waferscale_20v100_80hbm.yaml b/configs/hardware-config/waferscale_20v100_80hbm.yaml index 2812d62b..c4d4bbfe 100644 --- a/configs/hardware-config/waferscale_20v100_80hbm.yaml +++ b/configs/hardware-config/waferscale_20v100_80hbm.yaml @@ -29,6 +29,8 @@ sw_param: kv_cache: as_tensor_format h2d_bandwidth: 13297674240 # Bytes/s; set -1 to disable H2D timing + # DP gradient sync frequency: every_mb (default) or last_mb (sync only on final microbatch; ZeRO-3 stays per-mb) + dp_microbatch: last_mb tech_param: core: operating_frequency: 1.33e9 #Hz diff --git a/simulate_LLM.py b/simulate_LLM.py index f54ec259..85c6310e 100644 --- a/simulate_LLM.py +++ b/simulate_LLM.py @@ -137,6 +137,8 @@ def __init__( self.num_layer = self.misc_metadata.get("num_layer", 0) self.layer_per_device = self.num_layer / self.lp if self.lp else self.num_layer self.all_reduce = self.misc_metadata.get("all_reduce", "every layer") + self.dp_microbatch_mode = str(self.misc_metadata.get("dp_microbatch_mode", "every_mb")).lower() + self.dp_zero_stage = int(self.misc_metadata.get("dp_zero_stage", 0) or 0) self.transformer_cfg = self.comp_times.get("transformer", {}) self.T_grad_transformer = self.comp_times.get("grad_transformer", 0.0) @@ -836,6 +838,9 @@ def _bwd_exit_node(mb_idx: int, layer_idx: int) -> Any: for b in range(self.num_batch): # Data parallel collectives (all reduce for DDP/ZeRO-1, reduce-scatter + all-gather for ZeRO-2/3) if self.dp > 1: + apply_dp_all_mbs = self.dp_microbatch_mode != "last_mb" or self.dp_zero_stage >= 3 + if not apply_dp_all_mbs and b != 0: # mb is backwards in backpass + continue zero2_embedding_key = "zero2_embedding_gather" zero2_transformer_key = "zero2_transformer_gather" zero2_softmax_key = "zero2_softmax_gather" diff --git a/time_calculation.py b/time_calculation.py index 5cbc5ecc..5ff5b6e9 100644 --- a/time_calculation.py +++ b/time_calculation.py @@ -204,6 +204,7 @@ def __init__(self, hw_config, model_config, mode, *, astra_policy_override: Opti self.h2d_bandwidth = getattr(hw_config.sw_config, "h2d_bandwidth", -1) self.zero_stage = getattr(hw_config.sw_config, "dp_zero_stage", 0) self.full_recomputation = getattr(hw_config.sw_config, "full_recomputation", False) + self.dp_microbatch = getattr(hw_config.sw_config, "dp_microbatch", "every_mb") self.attached = True # Hardware Parameters diff --git a/time_calculation_LLM.py b/time_calculation_LLM.py index 36e8a768..4a8cf15b 100644 --- a/time_calculation_LLM.py +++ b/time_calculation_LLM.py @@ -2372,6 +2372,7 @@ def _resolve_operation(name: str) -> OperationTiming: "full_recomputation": self.full_recomputation, "flattened_mode": flattened_mode, "pipeline_style_recompute": pipeline_style_recompute_flag, + "dp_microbatch_mode": getattr(self, "dp_microbatch", "every_mb"), } pipeline_graph_obj = simulate_LLM.Graph( diff --git a/validation_scripts/imec_data/A100_train.csv b/validation_scripts/imec_data/A100_train.csv index d580d834..f8428204 100644 --- a/validation_scripts/imec_data/A100_train.csv +++ b/validation_scripts/imec_data/A100_train.csv @@ -1,5 +1,12 @@ device,model,batch,mb,dp,tp,pp,cp,tp_sp,recomputation,actual +A100_korthi,GPT 22B,4,1,1,8,1,1,False,full,1.4 +A100_korthi,GPT 175B,64,64,1,8,8,1,False,full,16.9 +A100_korthi,GPT 530B,280,280,1,8,35,1,False,full,46.8 +A100_korthi,GPT 1T,512,512,1,8,64,1,False,full,87.9 A100_korthi,GPT 22B,4,1,1,8,1,1,True,selective,1.1 A100_korthi,GPT 175B,64,64,1,8,8,1,True,selective,12.9 A100_korthi,GPT 530B,280,280,1,8,35,1,True,selective,35.5 A100_korthi,GPT 1T,512,512,1,8,64,1,True,selective,69.1 +A100_selene,GPT 310B,2160,144,15,8,16,1,False,full,34.1 +A100_selene,GPT 530B,2520,280,9,8,35,1,False,full,51.2 +A100_selene,GPT 1T,3072,512,6,8,64,1,False,full,100.7 diff --git a/validation_scripts/nvidia_train.py b/validation_scripts/nvidia_train.py index c87ee18a..3c2e27fd 100644 --- a/validation_scripts/nvidia_train.py +++ b/validation_scripts/nvidia_train.py @@ -113,10 +113,12 @@ def _build_spec( ) -> Tuple[ValidationSpec, str, str]: label = f"{device} {model} bs={batch} mb={mb} dp={dp} tp={tp} pp={pp} cp={cp} tp_sp={tp_sp}" + recomputation_mode = str(recomputation).strip().lower() + full_recompute = recomputation_mode in {"full", "true", "yes", "on", "1"} + model_overrides = { "model_param": { "global_batch_size": int(batch), - "recomputation": str(recomputation), } } @@ -128,7 +130,11 @@ def _build_spec( "cp": int(cp), "lp": int(pp), "mb": max(1, int(mb)), - } + }, + "sw_param": { + # sw_param.full_recomputation toggles full activation recompute during backward. + "full_recomputation": full_recompute, + }, } spec = ValidationSpec( @@ -341,7 +347,10 @@ def run( if __name__ == "__main__": - print("=== Running A100_korthi training validation ===") - run(device="A100_korthi", emit_logs=True, show_progress=True) - # print("\n=== Running A100_selene training validation ===") - # run(device="A100_selene", emit_logs=True, show_progress=True) + devices = [ + "A100_korthi", # Uncomment to include korthi runs. + "A100_selene", + ] + for device in devices: + print(f"=== Running {device} training validation ===") + run(device=device, emit_logs=True, show_progress=True) diff --git a/validation_scripts/validation_configs/hardware-config/H100_SXM5_80GB.yaml b/validation_scripts/validation_configs/hardware-config/H100_SXM5_80GB.yaml index de3dfc12..f5e72fbc 100644 --- a/validation_scripts/validation_configs/hardware-config/H100_SXM5_80GB.yaml +++ b/validation_scripts/validation_configs/hardware-config/H100_SXM5_80GB.yaml @@ -30,6 +30,8 @@ sw_param: # kv_cache precision may be set as bytes, dtype string, or "as_tensor_format". kv_cache: as_tensor_format h2d_bandwidth: -1 # Bytes/s; set -1 to disable H2D timing + # DP gradient sync frequency: every_mb (default) or last_mb (sync only on final microbatch; ZeRO-3 stays per-mb) + dp_microbatch: last_mb tech_param: core: operating_frequency: 1.98e9 # (Hz) GPU boost clock frequency [TePU] diff --git a/validation_scripts/validation_configs/hardware-config/a100_80GB.yaml b/validation_scripts/validation_configs/hardware-config/a100_80GB.yaml index 96d8361d..0a53f4d8 100644 --- a/validation_scripts/validation_configs/hardware-config/a100_80GB.yaml +++ b/validation_scripts/validation_configs/hardware-config/a100_80GB.yaml @@ -24,6 +24,8 @@ sw_param: # kv_cache precision may be set as bytes, dtype string, or "as_tensor_format". kv_cache: as_tensor_format h2d_bandwidth: -1 # Bytes/s; set -1 to disable H2D timing + # DP gradient sync frequency: every_mb (default) or last_mb (sync only on final microbatch; ZeRO-3 stays per-mb) + dp_microbatch: last_mb tech_param: core: operating_frequency: 1.41e9 # (Hz) diff --git a/validation_scripts/validation_configs/hardware-config/a100_80GB_korthikanti.yaml b/validation_scripts/validation_configs/hardware-config/a100_80GB_korthikanti.yaml index c5bac124..a34e6fb8 100644 --- a/validation_scripts/validation_configs/hardware-config/a100_80GB_korthikanti.yaml +++ b/validation_scripts/validation_configs/hardware-config/a100_80GB_korthikanti.yaml @@ -24,6 +24,8 @@ sw_param: # kv_cache precision may be set as bytes, dtype string, or "as_tensor_format". kv_cache: as_tensor_format h2d_bandwidth: -1 # Bytes/s; set -1 to disable H2D timing + # DP gradient sync frequency: every_mb (default) or last_mb (sync only on final microbatch; ZeRO-3 stays per-mb) + dp_microbatch: last_mb tech_param: core: operating_frequency: 1.41e9 # (Hz) diff --git a/validation_scripts/validation_configs/hardware-config/a100_80GB_no_parallelism.yaml b/validation_scripts/validation_configs/hardware-config/a100_80GB_no_parallelism.yaml index be423bb2..f0476141 100644 --- a/validation_scripts/validation_configs/hardware-config/a100_80GB_no_parallelism.yaml +++ b/validation_scripts/validation_configs/hardware-config/a100_80GB_no_parallelism.yaml @@ -27,6 +27,8 @@ sw_param: # kv_cache precision may be set as bytes, dtype string, or "as_tensor_format". kv_cache: as_tensor_format h2d_bandwidth: -1 # Bytes/s; set -1 to disable H2D timing + # DP gradient sync frequency: every_mb (default) or last_mb (sync only on final microbatch; ZeRO-3 stays per-mb) + dp_microbatch: last_mb tech_param: core: operating_frequency: 1.41e9 diff --git a/validation_scripts/validation_configs/hardware-config/a100_80GB_selene_sc.yaml b/validation_scripts/validation_configs/hardware-config/a100_80GB_selene_sc.yaml index 78ab94cf..3ab41ce0 100644 --- a/validation_scripts/validation_configs/hardware-config/a100_80GB_selene_sc.yaml +++ b/validation_scripts/validation_configs/hardware-config/a100_80GB_selene_sc.yaml @@ -24,6 +24,8 @@ sw_param: # kv_cache precision may be set as bytes, dtype string, or "as_tensor_format". kv_cache: as_tensor_format h2d_bandwidth: -1 # Bytes/s; set -1 to disable H2D timing + # DP gradient sync frequency: every_mb (default) or last_mb (sync only on final microbatch; ZeRO-3 stays per-mb) + dp_microbatch: last_mb tech_param: core: operating_frequency: 1.41e9 # (Hz) From 0a2a12f63b99c129abdc236524a6f6f882bb0134 Mon Sep 17 00:00:00 2001 From: gkarfakis19 Date: Thu, 4 Dec 2025 16:29:05 -0800 Subject: [PATCH 70/77] added scotchpy to reqs txt --- requirements.txt | 1 + 1 file changed, 1 insertion(+) diff --git a/requirements.txt b/requirements.txt index 746d674e..03d8cc60 100644 --- a/requirements.txt +++ b/requirements.txt @@ -11,3 +11,4 @@ pydot==3.0.0 graphviz==0.21 joblib>=1.5.2 seaborn +scotchpy64>=1.0.1 From e59304570b2291dd8c7bc08a8ad3e098d624748a Mon Sep 17 00:00:00 2001 From: NanoMembers Date: Thu, 4 Dec 2025 20:29:04 -0800 Subject: [PATCH 71/77] Plot GEMM operation comparison --- gemm_calc.py | 79 ++++ ispass_plot.py | 263 ++++++++++++ tile.py | 22 + validate_gemm.py | 460 +++++++++++++++++++++ validation_scripts/llm_shapes.txt | 6 + validation_scripts/ncu_metrics_results.csv | 49 +++ 6 files changed, 879 insertions(+) create mode 100644 gemm_calc.py create mode 100644 ispass_plot.py create mode 100644 validate_gemm.py create mode 100644 validation_scripts/llm_shapes.txt create mode 100644 validation_scripts/ncu_metrics_results.csv diff --git a/gemm_calc.py b/gemm_calc.py new file mode 100644 index 00000000..8b7478bb --- /dev/null +++ b/gemm_calc.py @@ -0,0 +1,79 @@ +import sys +from tile import TiledGEMM + + +class GEMMCalc: + def __init__(self, core, mem_hierarchy, dtype_bytes, flashattn_enable): + self.core = core + self.mem_hierarchy = mem_hierarchy + self.dtype_bytes = dtype_bytes + self.flashattn_enable = flashattn_enable + + def run(self, M : int, K : int, N : int, name = "", throughput: float | None = None): + best_result = self._search_tile_space(M, K, N, throughput=throughput) + return best_result + + def _search_tile_space(self, M : int, K : int, N : int, throughput: float | None = None): + best_gemm = None + best_time = float("inf") + + for gemm in TiledGEMM.enumerate_candidates( + self.core, + self.mem_hierarchy, + M, K, N, + self.dtype_bytes + ): + time = self._compute_time(gemm, throughput=throughput) + if time < best_time: + best_time = time + best_gemm = gemm + + return best_gemm, best_time + + def _compute_time(self, gemm : TiledGEMM, throughput: float | None = None): + times,_,_ = self._roofline( + gemm.GEMM_flop, + gemm.mem_accesses.totals(), + gemm.bundle_util, + flashattn_enable=self.flashattn_enable, + throughput=throughput + ) + return max(times) + + def _roofline(self, total_flops, mem_access : tuple, util : float, + flashattn_enable : bool, throughput: float | None = None): + + if not throughput: + throughput = self.core.getThroughput() + throughput *= util + num_level = len(mem_access) + try: + if not flashattn_enable: + assert mem_access[num_level - 1] > 0, f"mem_access: {mem_access}" + except Exception as e: + print( + "Number of accesses to the last level of memory hierarchy cannot be zero:\n {}".format(e), + flush=True, + ) + sys.exit(0) + levels_to_compute = [(i, mem_access[i]) for i in range(num_level)] + + # Compute roofline time for each level + times = [] + inf_pts = [] + comp_ints = [] + for level_idx, num_mem in levels_to_compute: + mem_bw = self.mem_hierarchy[level_idx].getThroughput() + mem_latency = self.mem_hierarchy[level_idx].getLatency() + inflection_point = float("inf") if mem_bw == 0 else throughput / mem_bw + # print(f"Level {level_idx}: mem_bw={mem_bw}, mem_latency={mem_latency}, inflection_point={inflection_point}", flush=True) + comp_int = float("inf") if num_mem == 0 else total_flops / num_mem + if comp_int < inflection_point: # mem-bound + level_time = (float("inf") if (mem_bw == 0 or num_mem == 0) else (num_mem / mem_bw)) + mem_latency + else: # compute-bound + level_time = float("inf") if (throughput == 0) else (total_flops / throughput) + times.append(level_time) + inf_pts.append(inflection_point) + comp_ints.append(comp_int) + # print("Roofline: exited {}".format(name)) + return times, inf_pts, comp_ints \ No newline at end of file diff --git a/ispass_plot.py b/ispass_plot.py new file mode 100644 index 00000000..3e39bbc1 --- /dev/null +++ b/ispass_plot.py @@ -0,0 +1,263 @@ +from gemm_calc import GEMMCalc +from hw_component import Core, MemoryHierarchy +from validate_gemm import extract_gemm_shapes, parse_llm_shapes_file +import config +from typing import List, Dict + +import os +import csv +from pathlib import Path +import pprint +import matplotlib.pyplot as plt +import numpy as np + +STACK_ORDER = [ + 'qkv_proj', + 'attention_score', + 'attention_output', + 'output_proj', + 'ffn1', + 'ffn2', + 'linear' +] + +def configs(): + """Load hardware and model configurations from YAML files.""" + hw_config_path = "configs/hardware-config/a100_80GB_no_parallelism.yaml" + model_config_path = "configs/model-config/GEMM.yaml" + + # Make paths absolute relative to the repository root + # base_dir = os.path.dirname(os.path.dirname(os.path.abspath(__file__))) + # hw_config_path = os.path.join(base_dir, hw_config_path) + # model_config_path = os.path.join(base_dir, model_config_path) + + # Load configurations using config module + hw_config = config.parse_config(hw_config_path, config_type="hardware") + model_config = config.parse_config(model_config_path, config_type="GEMM") + + return hw_config, model_config + +def csv_validation_data(): + """Load validation data from ncu_metrics_results.csv.""" + csv_path = Path("validation_scripts") / "ncu_metrics_results.csv" + data = [] + with open(csv_path, 'r') as f: + reader = csv.DictReader(f) + for row in reader: + parsed_row = { + 'm': int(row['m']), + 'k': int(row['k']), + 'n': int(row['n']), + 'latency_ms': float(row['latency_ms']), + 'clock': float(row['clock_Ghz']) * 1e9, # Convert GHz to Hz + } + data.append(parsed_row) + return data + +def _create_gemm_calc(configs): + """Helper function to create GEMMCalc instance from configs. + + Args: + configs: Tuple of (hw_config, model_config) + + Returns: + Initialized GEMMCalc instance + """ + hw_config, model_config = configs + core = Core(hw_config) + mem_hierarchy = MemoryHierarchy(hw_config) + dtype_bytes = hw_config.sw_config.precision.activations + flashattn_enable = False + return GEMMCalc(core, mem_hierarchy.memLayer, dtype_bytes, flashattn_enable) + +def find_row( + csv_data: List[Dict], + m: int, + k: int, + n: int +) -> Dict | None: + """Find a row in CSV data matching given m, k, n dimensions.""" + for row in csv_data: + if row['m'] == m and row['k'] == k and row['n'] == n: + return row + return None + +def compare( + gemm_calc: GEMMCalc, + csv_validation_data: List[Dict], + shapes_file: str, +): + models_data = parse_llm_shapes_file(shapes_file) + gemm_configs = extract_gemm_shapes(models_data) + + data = {} + + for cfg in gemm_configs: + if cfg['op_name'] == "router": + continue # Skip router operations + ncu_data = find_row( + csv_validation_data, + cfg['m'], + cfg['k'], + cfg['n'] + ) + + num_heads = cfg.get('batch', 1) + + ncu_time = ncu_data['latency_ms'] + clock = ncu_data['clock'] + derate = clock / gemm_calc.core.operating_freq + + # Run GEMMCalc + best_gemm, best_time = gemm_calc.run(cfg['m'], cfg['k'], cfg['n'], name=f"{cfg['model']}_{cfg['op_name']}", throughput=gemm_calc.core.getThroughput() * derate) + + best_time = (best_time + 9e-6) * 1e3 + + model_dict = data.get(f"{cfg['model']}", {}) + model_dict[f"{cfg['op_name']}"] = { + 'm': cfg['m'], + 'k': cfg['k'], + 'n': cfg['n'], + 'ncu_time_ms': ncu_time * num_heads, + 'gemmcalc_time_ms': best_time * num_heads, + 'err_pct': abs(ncu_time - best_time) / ncu_time * 100, + } + data[f"{cfg['model']}"] = model_dict + return data + +def plot_stacked_bar_chart(results: Dict[str, Dict[str, Dict]]): + """Create a double stacked bar chart comparing NCU and GEMMCalc times. + + Args: + results: Dictionary with structure {'model': {'op_name': {'ncu_time_ms', 'gemmcalc_time_ms', ...}}} + """ + # Split models into llama2 and llama3.1 groups + llama2_models = {k: v for k, v in results.items() if 'llama2' in k.lower()} + llama31_models = {k: v for k, v in results.items() if 'llama3.1' in k.lower() or 'llama3_1' in k.lower()} + + # Get all unique operation names across all models + all_ops = sorted(set(op for ops in results.values() for op in ops.keys())) + colors = plt.cm.tab20(np.linspace(0, 1, len(all_ops))) + + # Create figure with two subplots + fig, (ax1, ax2) = plt.subplots(1, 2, figsize=(16, 8)) + + # Plot llama2 models + _plot_model_group(ax1, llama2_models, all_ops, colors, 'Llama2 Models (Batch Size 16)') + + # Plot llama3.1 models + _plot_model_group(ax2, llama31_models, all_ops, colors, 'Llama3.1 Models (Batch Size 1)') + + # Create shared legend + legend_elements = [plt.Rectangle((0,0),1,1, fc=colors[i], label=op) + for i, op in enumerate(STACK_ORDER)] + legend_elements.insert(0, plt.Rectangle((0,0),1,1, fc='white', ec='black', + label='Actual')) + legend_elements.insert(1, plt.Rectangle((0,0),1,1, fc='white', ec='black', hatch='///', + label='DeepFlowSim')) + fig.legend(handles=legend_elements, loc='upper center', bbox_to_anchor=(0.5, 0.96), + ncol=len(legend_elements)//2 + len(legend_elements)%2, frameon=True) + + fig.suptitle('aggregated error weighed by time per GEMM', fontsize=16, fontweight='bold', y=0.99) + + plt.subplots_adjust(top=0.85) + + # Save the plot + output_path = Path("validation_scripts") / "comparison_plot.png" + plt.savefig(output_path, dpi=300, bbox_inches='tight') + print(f"Plot saved to {output_path}") + +def _plot_model_group(ax, model_dict: Dict[str, Dict[str, Dict]], all_ops: List[str], colors, title: str): + """Helper function to plot a group of models on a given axis. + + Args: + ax: Matplotlib axis to plot on + model_dict: Dictionary of models and their operations + all_ops: List of all operation names + colors: Color map for operations + title: Title for the subplot + """ + if not model_dict: + ax.set_visible(False) + return + + models = list(model_dict.keys()) + num_models = len(models) + + # Prepare data structures + ncu_times = {} + gemmcalc_times = {} + + for model, ops in model_dict.items(): + ncu_times[model] = {op: data['ncu_time_ms'] for op, data in ops.items()} + gemmcalc_times[model] = {op: data['gemmcalc_time_ms'] for op, data in ops.items()} + + x = np.arange(num_models) + width = 0.25 + gap = 0.05 + + # Create stacked bars for each operation + ncu_bottoms = np.zeros(num_models) + gemmcalc_bottoms = np.zeros(num_models) + + for idx, op in enumerate(reversed(STACK_ORDER)): + ncu_values = [ncu_times[model].get(op, 0) for model in models] + gemmcalc_values = [gemmcalc_times[model].get(op, 0) for model in models] + + ax.bar(x - (width+gap)/2, ncu_values, width, bottom=ncu_bottoms, color=colors[len(all_ops) - idx - 1]) + ax.bar(x + (width+gap)/2, gemmcalc_values, width, bottom=gemmcalc_bottoms, + color=colors[len(all_ops) - idx - 1], hatch='///', edgecolor='white', linewidth=0.5) + + ncu_bottoms += ncu_values + gemmcalc_bottoms += gemmcalc_values + + # Add overall error annotations with vertical brackets + for i, model in enumerate(models): + total_ncu = ncu_bottoms[i] + total_gemmcalc = gemmcalc_bottoms[i] + overall_error = abs(total_ncu - total_gemmcalc) / total_ncu * 100 if total_ncu > 0 else 0 + + # Draw vertical bracket showing difference + min_height = min(total_ncu, total_gemmcalc) + max_height = max(total_ncu, total_gemmcalc) + + # Bracket position (to the right of the bars) + bracket_x = x[i] + bracket_width = 0.02 + + # Draw vertical line + ax.plot([bracket_x, bracket_x], [min_height, max_height], + 'k-', linewidth=1.5, solid_capstyle='butt') + # Draw top horizontal line + ax.plot([bracket_x - bracket_width/2, bracket_x + bracket_width/2], [max_height, max_height], + 'k-', linewidth=1.5) + # Draw bottom horizontal line + ax.plot([bracket_x - bracket_width/2, bracket_x + bracket_width/2], [min_height, min_height], + 'k-', linewidth=1.5) + + # Position error text beside the bracket + mid_height = (min_height + max_height) / 2 + ax.text(bracket_x, max_height, f'err: {overall_error:.1f}%', + ha='center', va='bottom', fontsize=10, fontweight='bold') + + # Customize the subplot + # ax.set_xlabel('Model', fontsize=12, fontweight='bold') + ax.set_ylabel('Time (ms)', fontsize=12, fontweight='bold') + ax.set_title(title, fontsize=14, fontweight='bold') + ax.set_xticks(x) + ax.set_xticklabels(models) + ax.grid(axis='y', alpha=0.3) + +if __name__ == "__main__": + hw_config, model_config = configs() + gemm_calc = _create_gemm_calc((hw_config, model_config)) + csv_data = csv_validation_data() + shapes_file = Path("validation_scripts") / "llm_shapes.txt" + + results = compare(gemm_calc, csv_data, shapes_file) + + # pp = pprint.PrettyPrinter(indent=4) + # pp.pprint(results) + + # Create visualization + plot_stacked_bar_chart(results) \ No newline at end of file diff --git a/tile.py b/tile.py index 5fedf9b9..0208dbf7 100644 --- a/tile.py +++ b/tile.py @@ -31,6 +31,7 @@ from math import ceil from functools import lru_cache from enum import IntEnum +import math from typing import Tuple, Union, NamedTuple class AccessBytes(NamedTuple): @@ -383,6 +384,27 @@ def GEMM_flop(self): """Theoretical FLOP count for the full GEMM of size MxKxN.""" return self.M * self.N * (2 * self.K - 1) + @property + def bundle_util(self): + """Wave utilization of bundles""" + reuse_M = -(-self.M // self.l2_M) + reuse_K = -(-self.K // self.l2_K) + reuse_N = -(-self.N // self.l2_N) + + gemm_waves = self.memLayer[1].calc_waves_per_sm( + self.M, self.K, self.N, + self.l2_M, self.l2_K, self.l2_N + ) + + if gemm_waves == -1: + return 1.0 + + full_waves = math.floor(gemm_waves) + partial_wave = gemm_waves - full_waves # fractional wave + + return (full_waves + partial_wave * partial_wave) / gemm_waves + + def sysarray_accesses(self): """Return systolic-array traffic (bytes) using cached pure function. diff --git a/validate_gemm.py b/validate_gemm.py new file mode 100644 index 00000000..08862490 --- /dev/null +++ b/validate_gemm.py @@ -0,0 +1,460 @@ +import os +import csv +import itertools +import argparse +import logging +from collections import OrderedDict +import re +from validation_scripts.helpers import ( + run_cutlass_profiler, + CUTLASS_ENV_VAR, + run_ncu_profiler, + load_metrics_spec, + parse_int_list, + read_cutlass_results, + ensure_cutlass_profiler, + get_metric_labels, +) + +from run_perf import import_config + +logger = logging.getLogger(__name__) +if not logger.handlers: + h = logging.StreamHandler() + h.setFormatter(logging.Formatter("%(asctime)s - %(levelname)s - %(message)s")) + logger.addHandler(h) +logger.setLevel(logging.INFO) + + +def parse_llm_shapes_file(filepath): + """ + Parse llm_shapes.txt into a list of dicts, each containing model name and OrderedDict of operations. + Each operation maps to a tuple (m,k,n) or (b,m,k,n). + + Returns: List[dict] where each dict has 'model' (str) and 'shapes' (OrderedDict[str, tuple]) + """ + models = [] + with open(filepath, 'r') as f: + for line in f: + line = line.strip() + if not line: + continue + # Expected format: "ModelName: OrderedDict([('op1', (dims)), ...])" + match = re.match(r"^([^:]+):\s*OrderedDict\((.+)\)$", line) + if not match: + logger.warning(f"Skipping malformed line: {line}") + continue + model_name = match.group(1).strip() + shapes_str = match.group(2).strip() + + # Parse the OrderedDict content: [('key', (tuple)), ...] + shapes = OrderedDict() + # Remove outer brackets + shapes_str = shapes_str.strip('[]') + # Split by '), (' pattern to separate entries + entries = re.findall(r"\('([^']+)',\s*\(([^)]+)\)\)", shapes_str) + for op_name, dims_str in entries: + dims = tuple(map(int, dims_str.split(','))) + shapes[op_name] = dims + + models.append({'model': model_name, 'shapes': shapes}) + logger.info(f"Parsed model {model_name} with {len(shapes)} operations") + + return models + + +def extract_gemm_shapes(models_data): + """ + Extract GEMM shapes from parsed models data. + Returns: List[dict] with keys: 'model', 'op_name', 'm', 'k', 'n', 'batch' (optional) + """ + gemm_shapes = [] + for model_info in models_data: + model_name = model_info['model'] + for op_name, dims in model_info['shapes'].items(): + if len(dims) == 3: + m, k, n = dims + gemm_shapes.append({ + 'model': model_name, + 'op_name': op_name, + 'm': m, 'k': k, 'n': n + }) + elif len(dims) == 4: + b, m, k, n = dims + # For batched GEMMs, we profile with the actual m dimension + gemm_shapes.append({ + 'model': model_name, + 'op_name': op_name, + 'm': m, 'k': k, 'n': n, + 'batch': b + }) + else: + logger.warning(f"Skipping {model_name}:{op_name} with unsupported dims: {dims}") + + return gemm_shapes + + +def _safe_err(ncu_val, pred_val): + if ncu_val is None or pred_val is None: + return "", "" + try: + abs_err = pred_val - ncu_val + pct_err = abs_err / ncu_val * 100 if ncu_val != 0 else "" + return f"{abs_err:<.3f}", f"{pct_err:<.2f}" + except Exception: + return "", "" + +def main(): + parser = argparse.ArgumentParser( + description="Search CUTLASS kernels across GEMM and CTA sizes." + ) + # parser.add_argument("--cutlass-path", required=True, help="Path to CUTLASS profiler binary") + parser.add_argument("--m-list", default="256,1024,4096") + parser.add_argument("--n-list", default="256,1024,4096") + parser.add_argument("--k-list", default="256,1024,4096") + parser.add_argument("--shapes-file", default="validation_scripts/llm_shapes.txt", + help="Path to file containing LLM GEMM shapes (e.g., llm_shapes.txt)") + parser.add_argument("--cta-m", default="64,128,256") + parser.add_argument("--cta-n", default="64,128,256") + parser.add_argument("--cta-k", default="32,64") + parser.add_argument("--cutlass-output", default="validation_scripts/kernel_search_results.csv") + parser.add_argument("--ncu-output", default="validation_scripts/cta_sweep/fastest.csv") + parser.add_argument("--ncu-metrics", default="validation_scripts/metrics.yaml") + parser.add_argument("--log-level", default="INFO") + parser.add_argument("--cutlass-path", default=None, help="Path to cutlass_profiler binary; overrides env") + parser.add_argument("--deepflow-compare", action="store_true", + help="Enable DeepFlow model comparison CSV output") + parser.add_argument("--deepflow-output", default="validation_scripts/deepflow_comparison.csv", + help="Output CSV for DeepFlow vs NCU errors") + parser.add_argument("--hw-config", default="configs/hardware-config/a100_80GB_no_parallelism.yaml", + help="Path to A100 hardware config") + parser.add_argument("--sw-config", default="configs/model-config/GEMM.yaml", + help="Path to GEMM software config") + args = parser.parse_args() + + logger.setLevel(getattr(logging, args.log_level.upper(), logging.INFO)) + + # Ensure CUTLASS profiler path (no hardcoded default) + ensure_cutlass_profiler(args.cutlass_path) # added + logger.info(f"{CUTLASS_ENV_VAR}={os.environ.get(CUTLASS_ENV_VAR, '')}") + + # Load NCU metric spec (labels, units) from file or fallback + if args.ncu_metrics and os.path.isfile(args.ncu_metrics): + try: + ncu_metric_spec = load_metrics_spec(args.ncu_metrics) + logger.info(f"Loaded NCU metrics spec from {args.ncu_metrics} ({len(ncu_metric_spec)} metrics)") + except Exception as e: + logger.warning(f"Failed to load metrics spec from {args.ncu_metrics}: {e}.") + exit(1) + else: + logger.info("NCU metrics spec file not found.") + exit(1) + + # Determine GEMM shapes source + if args.shapes_file: + if not os.path.isfile(args.shapes_file): + logger.error(f"Shapes file not found: {args.shapes_file}") + exit(1) + logger.info(f"Loading GEMM shapes from {args.shapes_file}") + models_data = parse_llm_shapes_file(args.shapes_file) + gemm_configs = extract_gemm_shapes(models_data) + logger.info(f"Extracted {len(gemm_configs)} GEMM shapes from {len(models_data)} models") + else: + # Use m/n/k lists + m_vals = parse_int_list(args.m_list) + n_vals = parse_int_list(args.n_list) + k_vals = parse_int_list(args.k_list) + gemm_configs = [ + {'m': m, 'k': k, 'n': n, 'model': 'synthetic', 'op_name': f'm{m}_k{k}_n{n}'} + for m in m_vals for k in k_vals for n in n_vals + ] + logger.info(f"Using {len(gemm_configs)} synthetic GEMM shapes from m/k/n lists") + + cta_m_vals = parse_int_list(args.cta_m) + cta_n_vals = parse_int_list(args.cta_n) + cta_k_vals = parse_int_list(args.cta_k) + + cta_tuples = list(itertools.product(cta_m_vals, cta_n_vals, cta_k_vals)) + logger.info(f"CTA tuples ({len(cta_tuples)}): {cta_tuples}") + + out_path = args.cutlass_output + ncu_out_path = args.ncu_output + ncu_exists = os.path.isfile(ncu_out_path) + perform_ncu = not ncu_exists + ncu_write_header = not ncu_exists + + # DeepFlow TimeCalculation initialization + if args.deepflow_compare: + try: + deepflow_tc = import_config(args.hw_config, args.sw_config, "GEMM") + logger.info("DeepFlow TimeCalculation initialized.") + except Exception as e: + logger.warning(f"Failed to initialize DeepFlow TimeCalculation: {e}") + deepflow_tc = None + else: + deepflow_tc = None + + # Accumulate NCU results for DeepFlow comparison + ncu_rows = [] # each: dict(m,k,n, metrics_dict) + + # If NCU output already exists, load it and skip further NCU profiling + if ncu_exists: + logger.info(f"{ncu_out_path} exists. Skipping NCU profiling and loading existing metrics.") + try: + with open(ncu_out_path, "r", newline="") as f_in: + reader = csv.reader(f_in) + header = next(reader, None) + if header is None: + logger.warning(f"{ncu_out_path} is empty. No NCU metrics to load.") + else: + # First 7 columns are fixed: m,k,n,cta_m,cta_k,cta_n,stages + metric_labels = header[7:] + for row in reader: + if len(row) < 7: + continue + try: + m = int(row[0]) + n = int(row[1]) + k = int(row[2]) + metrics = { + "cta_m": int(row[3]), + "cta_n": int(row[4]), + "cta_k": int(row[5]), + "stages": row[6], + } + # Parse remaining metrics as float when possible + for lab, val in zip(metric_labels, row[7:]): + try: + metrics[lab] = float(val) + except ValueError: + metrics[lab] = val + # Provide canonical aliases if present + if "latency_ms" not in metrics and "latency" in metrics: + metrics["latency_ms"] = metrics["latency"] + ncu_rows.append({"m": m, "k": k, "n": n, "metrics": metrics}) + except Exception as e: + logger.debug(f"Skipping malformed NCU row: {row} ({e})") + logger.info(f"Loaded {len(ncu_rows)} NCU metric rows from {ncu_out_path}") + except Exception as e: + logger.warning(f"Failed to read existing NCU metrics file {ncu_out_path}: {e}") + + # If CUTLASS results already exist, reuse them; only run NCU if perform_ncu + if os.path.isfile(out_path): # broken + logger.info(f"{out_path} exists. Skipping CUTLASS profiling and using existing results.") + existing = read_cutlass_results(out_path) + if not existing: + logger.warning(f"No valid entries found in {out_path}. Nothing to do.") + # If DeepFlow requested and we have NCU rows, still do comparison + if deepflow_tc and ncu_rows: + pass # fall-through to comparison block below + else: + return + else: + if perform_ncu: + with open(ncu_out_path, "a", newline="") as f_ncu: + writer_ncu = csv.writer(f_ncu) + if ncu_write_header: + metric_labels = get_metric_labels(ncu_metric_spec) + writer_ncu.writerow(["m", "k", "n", "cta_m", "cta_k", "cta_n", "stages"] + metric_labels) + for m, k, n, best_kernel in existing: + try: + logger.info(f"Running NCU for m={m} k={k} n={n} kernel={best_kernel.name}") + ncu_metrics = run_ncu_profiler( + m=m, k=k, n=n, kernel_name=best_kernel.name, + metric_spec=ncu_metric_spec + ) + ordered_values = [ncu_metrics.get(spec.label, "") for spec in ncu_metric_spec.values()] + writer_ncu.writerow([m, k, n, best_kernel.cta_m, best_kernel.cta_k, + best_kernel.cta_n, best_kernel.stages] + ordered_values) + ncu_rows.append({"m": m, "k": k, "n": n, "metrics": ncu_metrics}) + except Exception as e: + logger.warning(f"NCU profiling failed for m={m} k={k} n={n} kernel={best_kernel.name}: {e}") + logger.info(f"NCU metrics written to {ncu_out_path}") + else: + logger.info("NCU profiling skipped (existing NCU metrics file present).") + # DeepFlow comparison with existing (and/or loaded) NCU rows + if deepflow_tc and ncu_rows: + try: + write_header_df = not os.path.isfile(args.deepflow_output) + # if deepflow_output already exists, clear it and rewrite + if not write_header_df: + os.remove(args.deepflow_output) + write_header_df = True + with open(args.deepflow_output, "a", newline="") as f_df: + w_df = csv.writer(f_df) + if write_header_df: + w_df.writerow([ + "m","k","n","cta_m","cta_k","cta_n","l1_m","l1_k","l1_n", + "ncu_latency_us","df_latency_us","latency_abs_err","latency_pct_err", + "ncu_l2_read","df_l2_read","l2_read_pct_err", + "ncu_l2_write","df_l2_write","l2_write_pct_err", + "ncu_l2_total","df_l2_total","l2_total_pct_err", + "ncu_dram_read","df_dram_read","dram_read_pct_err", + "ncu_dram_write","df_dram_write","dram_write_pct_err", + "ncu_dram_total","df_dram_total","dram_total_pct_err" + ]) + for row in ncu_rows: + m,k,n = row["m"], row["k"], row["n"] + metrics = row["metrics"] + ncu_latency = metrics.get("latency_ms") + ncu_latency *= 1e3 # ms → us + ncu_l2_read = metrics.get("l2_read") + ncu_l2_write = metrics.get("l2_write") + ncu_l2_total = ncu_l2_read + ncu_l2_write + ncu_dram_read = metrics.get("dram_read") + ncu_dram_write = metrics.get("dram_write") + ncu_dram_total = ncu_dram_read + ncu_dram_write + ncu_cta_m = metrics.get("cta_m") + ncu_cta_n = metrics.get("cta_n") + ncu_cta_k = metrics.get("cta_k") + try: + df_latency, _, df_tile_dims, _, best_rw_access = deepflow_tc.getGEMMTime(m, k, n, name=f"m{m}_k{k}_n{n}") + except Exception as e: + logger.warning(f"DeepFlow getGEMMTime failed for {m},{k},{n}: {e}") + continue + # logger.info(f"({m},{k},{n}): {df_tile_dims}") + df_latency *= 1e6 # s → us + l1_m, l1_k, l1_n = df_tile_dims[1] + df_l2_read = best_rw_access.reads[2] + df_l2_write = best_rw_access.writes[2] + df_l2_total = df_l2_read + df_l2_write + df_dram_read = best_rw_access.reads[3] + df_dram_write = best_rw_access.writes[3] + df_dram_total = df_dram_read + df_dram_write + lat_abs, lat_pct = _safe_err(ncu_latency, df_latency) + _, l2_read_pct = _safe_err(ncu_l2_read, df_l2_read) + _, l2_write_pct = _safe_err(ncu_l2_write, df_l2_write) + _, l2_total_pct = _safe_err(ncu_l2_total, df_l2_total) + _, dram_read_pct = _safe_err(ncu_dram_read, df_dram_read) + _, dram_write_pct = _safe_err(ncu_dram_write, df_dram_write) + _, dram_total_pct = _safe_err(ncu_dram_total, df_dram_total) + w_df.writerow([ + m,k,n,ncu_cta_m,ncu_cta_k,ncu_cta_n,l1_m,l1_k,l1_n, + f"{ncu_latency:.3f}", f"{df_latency:.3f}", lat_abs, lat_pct, + ncu_l2_read, df_l2_read, l2_read_pct, + ncu_l2_write, df_l2_write, l2_write_pct, + ncu_l2_total, df_l2_total, l2_total_pct, + ncu_dram_read, df_dram_read, dram_read_pct, + ncu_dram_write, df_dram_write, dram_write_pct, + ncu_dram_total, df_dram_total, dram_total_pct + ]) + logger.info(f"DeepFlow comparison written to {args.deepflow_output}") + except Exception as e: + logger.warning(f"Failed DeepFlow comparison output: {e}") + return + + # Otherwise (CUTLASS results absent), run CUTLASS search; only run NCU if perform_ncu + write_header = not os.path.isfile(out_path) + with open(out_path, "a", newline="") as f_kernel: + writer_kernel = csv.writer(f_kernel) + if write_header: + writer_kernel.writerow([ + "model","op_name","m","k","n","best_kernel_name","cta_m","cta_k","cta_n", + "stages","inst_m","inst_n","inst_k","best_time_ms", + ]) + total_jobs = len(gemm_configs) + job_idx = 0 + for config in gemm_configs: + m, k, n = config['m'], config['k'], config['n'] + model_name = config.get('model', 'unknown') + op_name = config.get('op_name', 'unknown') + job_idx += 1 + logger.info(f"[{job_idx}/{total_jobs}] Profiling {model_name}:{op_name} GEMM m={m} k={k} n={n}") + best_kernel, best_time = run_cutlass_profiler(m=m, k=k, n=n, cta_tuples=cta_tuples) + if best_kernel is None: + logger.warning(f"No kernel found for {model_name}:{op_name} m={m} k={k} n={n}") + continue + writer_kernel.writerow([ + model_name,op_name,m,k,n,best_kernel.name,best_kernel.cta_m,best_kernel.cta_k, + best_kernel.cta_n,best_kernel.stages,best_kernel.inst_m, + best_kernel.inst_n,best_kernel.inst_k,f"{best_time:.6f}", + ]) + logger.info(f"Best kernel: {best_kernel.name} CTA=({best_kernel.cta_m},{best_kernel.cta_k},{best_kernel.cta_n}) time={best_time:.4f} ms") + if perform_ncu: + # Open (append) NCU file lazily only if needed + if ncu_write_header: + with open(ncu_out_path, "a", newline="") as f_ncu: + writer_ncu = csv.writer(f_ncu) + metric_labels = get_metric_labels(ncu_metric_spec) + writer_ncu.writerow(["m","k","n","cta_m","cta_k","cta_n","stages"] + metric_labels) + ncu_write_header = False + try: + ncu_metrics = run_ncu_profiler( + m=m, k=k, n=n, kernel_name=best_kernel.name, + metric_spec=ncu_metric_spec + ) + ordered_values = [ncu_metrics.get(spec.label, "") for spec in ncu_metric_spec.values()] + with open(ncu_out_path, "a", newline="") as f_ncu: + writer_ncu = csv.writer(f_ncu) + writer_ncu.writerow([m, k, n, best_kernel.cta_m, best_kernel.cta_k, + best_kernel.cta_n, best_kernel.stages] + ordered_values) + ncu_rows.append({"m": m, "k": k, "n": n, "metrics": ncu_metrics}) + logger.info(f"NCU metrics collected for kernel {best_kernel.name}: latency={ncu_metrics.get('latency_ms', 'N/A')}") + except Exception as e: + logger.warning(f"NCU profiling failed for kernel {best_kernel.name}: {e}") + + logger.info(f"Kernel results written to {out_path}") + if perform_ncu: + logger.info(f"NCU metrics written to {ncu_out_path}") + else: + logger.info(f"Skipped writing NCU metrics (existing file retained: {ncu_out_path})") + + # DeepFlow comparison after new profiling (only if NCU metrics available or loaded) + if deepflow_tc and ncu_rows: + try: + write_header_df = not os.path.isfile(args.deepflow_output) + with open(args.deepflow_output, "a", newline="") as f_df: + w_df = csv.writer(f_df) + if write_header_df: + w_df.writerow([ + "m","k","n","cta_m","cta_k","cta_n","l1_m","l1_k","l1_n", + "ncu_latency_ms","df_latency_ms","latency_abs_err","latency_pct_err", + "ncu_l2_read","df_l2_read","l2_read_pct_err", + "ncu_l2_write","df_l2_write","l2_write_pct_err", + "ncu_dram_read","df_dram_read","dram_read_pct_err", + "ncu_dram_write","df_dram_write","dram_write_pct_err" + ]) + for row in ncu_rows: + m,k,n = row["m"], row["k"], row["n"] + metrics = row["metrics"] + ncu_latency = metrics.get("latency_ms") + ncu_l2_read = metrics.get("l2_read") + ncu_l2_write = metrics.get("l2_write") + ncu_dram_read = metrics.get("dram_read") + ncu_dram_write = metrics.get("dram_write") + ncu_cta_m = metrics.get("cta_m") + ncu_cta_n = metrics.get("cta_n") + ncu_cta_k = metrics.get("cta_k") + try: + df_latency, _, df_tile_dims, _, best_rw_access = deepflow_tc.getGEMMTime(m, k, n) + except Exception as e: + logger.warning(f"DeepFlow getGEMMTime failed for {m},{k},{n}: {e}") + continue + l1_m, l1_k, l1_n = df_tile_dims[1] + df_l2_read = best_rw_access.reads[2] + df_l2_write = best_rw_access.writes[2] + df_dram_read = best_rw_access.reads[3] + df_dram_write = best_rw_access.writes[3] + lat_abs, lat_pct = _safe_err(ncu_latency, df_latency * 1e3) + _, l2_read_pct = _safe_err(ncu_l2_read, df_l2_read) + _, l2_write_pct = _safe_err(ncu_l2_write, df_l2_write) + _, dram_read_pct = _safe_err(ncu_dram_read, df_dram_read) + _, dram_write_pct = _safe_err(ncu_dram_write, df_dram_write) + w_df.writerow([ + m,k,n,ncu_cta_m,ncu_cta_k,ncu_cta_n,l1_m,l1_k,l1_n, + ncu_latency, df_latency, lat_abs, lat_pct, + ncu_l2_read, df_l2_read, l2_read_pct, + ncu_l2_write, df_l2_write, l2_write_pct, + ncu_dram_read, df_dram_read, dram_read_pct, + ncu_dram_write, df_dram_write, dram_write_pct + ]) + logger.info(f"DeepFlow comparison written to {args.deepflow_output}") + except Exception as e: + logger.warning(f"Failed DeepFlow comparison output: {e}") + + +if __name__ == "__main__": + os.environ[CUTLASS_ENV_VAR] = ( + "/u1/ee/nanoproj/users/yaoe888/cutlass/build/tools/profiler/cutlass_profiler" + ) + main() diff --git a/validation_scripts/llm_shapes.txt b/validation_scripts/llm_shapes.txt new file mode 100644 index 00000000..47113a44 --- /dev/null +++ b/validation_scripts/llm_shapes.txt @@ -0,0 +1,6 @@ +Llama2-7B: OrderedDict([('qkv_proj', (16384, 4096, 12288)), ('attention_score', (512, 1024, 128, 1024)), ('attention_output', (512, 1024, 1024, 128)), ('output_proj', (16384, 4096, 4096)), ('ffn1', (16384, 4096, 22016)), ('ffn2', (16384, 11008, 4096)), ('linear', (16384, 4096, 32000)), ('router', (16384, 4096, 1))]) +Llama2-13B: OrderedDict([('qkv_proj', (16384, 5120, 15360)), ('attention_score', (640, 1024, 128, 1024)), ('attention_output', (640, 1024, 1024, 128)), ('output_proj', (16384, 5120, 5120)), ('ffn1', (16384, 5120, 27648)), ('ffn2', (16384, 13824, 5120)), ('linear', (16384, 5120, 32000)), ('router', (16384, 5120, 1))]) +Llama2-70B: OrderedDict([('qkv_proj', (16384, 8192, 10240)), ('attention_score', (128, 8192, 128, 1024)), ('attention_output', (128, 8192, 1024, 128)), ('output_proj', (16384, 8192, 8192)), ('ffn1', (16384, 8192, 57344)), ('ffn2', (16384, 28672, 8192)), ('linear', (16384, 8192, 32000)), ('router', (16384, 8192, 1))]) +Llama3.1-8B: OrderedDict([('qkv_proj', (1024, 4096, 6144)), ('attention_score', (8, 4096, 128, 1024)), ('attention_output', (8, 4096, 1024, 128)), ('output_proj', (1024, 4096, 4096)), ('ffn1', (1024, 4096, 28672)), ('ffn2', (1024, 14336, 4096)), ('linear', (1024, 4096, 128256)), ('router', (1024, 4096, 1))]) +Llama3.1-70B: OrderedDict([('qkv_proj', (1024, 8192, 10240)), ('attention_score', (8, 8192, 128, 1024)), ('attention_output', (8, 8192, 1024, 128)), ('output_proj', (1024, 8192, 8192)), ('ffn1', (1024, 8192, 57344)), ('ffn2', (1024, 28672, 8192)), ('linear', (1024, 8192, 128256)), ('router', (1024, 8192, 1))]) +Llama3.1-405B: OrderedDict([('qkv_proj', (1024, 16384, 18432)), ('attention_score', (8, 16384, 128, 1024)), ('attention_output', (8, 16384, 1024, 128)), ('output_proj', (1024, 16384, 16384)), ('ffn1', (1024, 16384, 106496)), ('ffn2', (1024, 53248, 16384)), ('linear', (1024, 16384, 128256)), ('router', (1024, 16384, 1))]) \ No newline at end of file diff --git a/validation_scripts/ncu_metrics_results.csv b/validation_scripts/ncu_metrics_results.csv new file mode 100644 index 00000000..9b32e1d3 --- /dev/null +++ b/validation_scripts/ncu_metrics_results.csv @@ -0,0 +1,49 @@ +m,k,n,cta_m,cta_k,cta_n,stages,latency_ms,clock_Ghz,dram_read,dram_write,l2_hit,l2_total,l2_read,l2_write,l2_load,l2_store,global_st,global_ldgsts,global_ldgsts_sectors,global_ldgsts_requests,shared_ldsm,shared_ldgsts,shared_ldgsts_cache_bypass,grid_size,block_size,regs_per_thread,waves_per_sm,dynamic_shared_mem_per_block,shared_mem_config_size,theoretical_warps_per_sm,achieved_warps_per_sm,sm_efficiency,memory_efficiency +16384,4096,12288,128,32,128,3,7.444512,1.040302,2161039638,418370214,87.638415,32071083439,27670116351,422212465,27670116351,422212465,422212465,20860773194,607123552.0,50331648.0,55556404848,8436628955,20707303275,12288,128,164,37.93,50331,171966,12.0,11.906987,96.281251,66.35622 +1024,128,1024,64,32,256,4,0.010848000000000002,1.036226,565575,0,87.893802,10259905,5497558,2199023,5497558,2199023,2199023,5222680,155648.0,10240.0,9895604,947912,4365834,64,128,162,0.3,83886,171966,8.0,3.940938,10.796474,13.058038 +1024,1024,128,64,64,64,5,0.013344,1.038394,2494033,0,60.83675,11745561,8796093,268435,8796093,268435,268435,7971459,237568.0,16384.0,17867063,2005749,2005749,32,128,80,0.15,83886,171966,8.0,3.959352,8.758664,12.668023 +16384,4096,4096,128,32,128,3,2.487136,1.049378,869111059,136524259,85.753672,10816774203,9223372474,140737488,9223372474,140737488,140737488,6953589633,202374768.0,16777216.0,18518801616,2802948270,8075642323,4096,128,164,12.64,50331,171966,12.0,11.749281,95.232402,66.604182 +16384,4096,22016,128,32,128,3,13.175808,1.052499,3763901032,752867233,87.948712,57659793141,49575624234,756463999,49575624234,756463999,756463999,37375514079,1087762656.0,90177536.0,99538558955,15122685075,36510827641,22016,128,164,67.95,50331,171966,12.0,11.945483,96.33759,67.126956 +16384,11008,4096,128,32,256,3,6.545024,1.052425,2060710276,136070872,82.480282,22647972437,18590859356,140737488,18590859356,140737488,140737488,15510385331,451412288.0,33816576.0,49647681974,7001813525,6881838983,2048,256,218,18.96,75497,104857,8.0,7.997923,96.975567,52.639561 +16384,4096,32000,128,32,128,3,19.085312,1.053524,5325413702,1122295355,87.88347,83858419336,72057594037,1125899906,72057594037,1125899906,1125899906,54324947530,1581051296.0,131072000.0,144678138029,21995077278,51393319892,32000,128,164,98.77,50331,171966,12.0,11.955218,96.574459,67.351305 +16384,4096,1,64,32,64,2,0.21443199999999998,1.050273,140777216,2436857,21.688324,272897826,141665872,33554,141666677,16777,33554,0,0.0,0.0,564049465,0,0,256,128,92,0.47,16777,104857,20.0,9.787831,45.110381,70.066126 +16384,5120,15360,128,32,128,3,11.458784,1.052614,3768151976,524792256,85.992715,50864445135,43234555993,527765581,43234555993,527765581,527765581,32561167125,947648800.0,78643200.0,86739328780,13195250856,32929264076,15360,128,164,47.41,50331,171966,12.0,11.93105,96.593947,67.314009 +1024,128,1024,128,64,128,3,0.010944,1.030935,567017,128,87.118525,9091741,4398046,2199023,4398046,2199023,2199023,4398046,131072.0,8192.0,9895604,394264,4179607,64,128,168,0.59,100663,104857,4.0,4.00477,10.756691,13.041437 +1024,1024,128,64,64,64,5,0.01344,1.036274,2494168,0,60.839544,11745628,8796093,268435,8796093,268435,268435,7971459,237568.0,16384.0,17867063,2005749,2005749,32,128,80,0.15,83886,171966,8.0,3.978921,8.713897,12.558307 +16384,5120,5120,128,32,128,3,3.918304,1.039268,1466355521,173123689,84.373827,17126994915,14411519022,175921860,14411519022,175921860,175921860,10853711995,315882928.0,26214400.0,28913109951,4389332022,12245953562,5120,128,164,15.8,50331,171966,12.0,11.807363,95.369729,65.667447 +16384,5120,27648,128,32,128,3,20.562912,1.053423,6631522920,946931303,86.317107,91676702924,77822201432,949978046,77822201432,949978046,949978046,58610010631,1705765152.0,141557760.0,156130791804,23754743633,59409309414,27648,128,164,85.33,50331,171966,12.0,11.955874,96.81497,67.513753 +16384,13824,5120,128,32,128,3,10.335744,1.052809,5727252989,175648324,79.520888,46680338461,38911101252,175921860,38911101252,175921860,175921860,29228419606,850656656.0,70778880.0,77912273338,11891435150,31561409334,5120,128,164,15.8,50331,171966,12.0,11.811011,96.36264,67.170841 +16384,5120,32000,128,32,128,3,23.775616,1.053654,7625098106,1123201593,86.04038,106586295345,90071992547,1125899906,90071992547,1125899906,1125899906,67835691651,1974266528.0,163840000.0,180706935048,27532093282,62675846791,32000,128,164,98.77,50331,171966,12.0,11.958393,96.891624,67.579012 +16384,5120,1,64,32,64,2,0.265024,1.0517,175964004,2437796,22.555699,343609496,177083045,33554,177082038,16777,33554,0,0.0,0.0,704786953,0,0,256,128,92,0.47,16777,104857,20.0,9.752035,45.375592,70.629169 +16384,8192,10240,128,32,128,3,12.22576,1.053123,6179862012,350559659,81.287171,55614438446,46116860227,351843720,46116860227,351843720,351843720,34677784991,1009253536.0,83886080.0,92413864267,14087177050,36503870867,10240,128,164,31.6,50331,171966,12.0,11.906335,96.523044,67.296226 +8192,128,1024,128,32,128,3,0.027232,1.010449,2500342,581828,92.320141,65372489,35184372,17592186,35184372,17592186,17592186,30811021,917952.0,65536.0,79164837,8224275,32008881,512,128,164,1.58,50331,171966,12.0,10.229684,35.284347,43.133191 +8192,1024,128,128,32,128,5,0.027584,1.019056,17902765,107610,44.176805,48397604,35184372,2199023,35184372,2199023,2199023,35184372,1048576.0,65536.0,71468255,8821862,9027819,64,128,220,0.3,83886,171966,8.0,3.993464,34.539874,32.281163 +16384,8192,8192,128,32,256,3,9.919744,1.04257,3498759816,278153624,79.428101,34873624762,27670116351,281474976,27670116351,281474976,281474976,23094437508,672137072.0,50331648.0,73931090984,10412166809,10226037958,4096,256,218,37.93,75497,104857,8.0,7.999085,96.132498,52.467213 +16384,8192,57344,128,32,128,3,68.45232,1.055575,44016446199,2016652501,76.548944,332141729870,258254417487,2017612426,258254417487,2017612426,2017612426,194195640618,5651812560.0,469762048.0,517517640114,79076574548,173258501831,57344,128,164,176.99,50331,171966,12.0,11.985653,96.315677,67.290225 +16384,28672,8192,128,32,256,3,33.594464,1.054792,13954805084,280242455,78.001315,120923102338,96845406155,281474976,96845406155,281474976,281474976,80740515745,2349858816.0,176160768.0,258398531893,36580164519,35932477008,4096,256,218,37.93,75497,104857,8.0,7.999301,98.199608,52.781711 +16384,8192,32000,128,32,128,3,37.99808,1.053792,21499678570,1125221302,79.021507,180300408877,144115188075,1125899906,144115188075,1125899906,1125899906,108368023873,3153912800.0,262144000.0,288793326105,44079942996,104244490197,32000,128,164,98.77,50331,171966,12.0,11.970652,96.988343,67.650502 +16384,8192,1,64,32,64,2,0.418208,1.053273,281527992,2453634,23.717734,555372255,283333221,33554,283333422,16777,33554,0,0.0,0.0,1154046975,0,0,256,128,92,0.47,16777,104857,20.0,9.73105,45.656654,71.296692 +1024,4096,6144,128,32,128,4,0.27532799999999996,1.036044,73249727,10017608,85.279225,994988364,844424930,13194139,844424930,13194139,13194139,638279921,19022048.0,1572864.0,1736137450,229657538,361850071,384,128,160,1.78,67108,138412,8.0,7.311376,81.688114,56.934269 +4096,128,1024,128,32,128,2,0.01856,1.013349,1400293,0,91.381847,32724530,16636555,8796093,16638569,8796093,8796093,0,0.0,0.0,37383395,0,0,256,128,148,0.79,33554,104857,12.0,9.262791,25.811254,42.585261 +4096,1024,128,64,32,128,6,0.019584,1.027625,9097948,128,55.225712,35732685,26388279,1099511,26388279,1099511,1099511,18966575,565248.0,49152.0,53601191,15470472,15470472,64,128,158,0.3,75497,171966,8.0,4.023061,24.121831,25.748215 +1024,4096,4096,128,32,128,4,0.20668799999999998,1.042899,50164010,5554600,84.81866,661345238,562949953,8796093,562949953,8796093,8796093,425522811,12681312.0,1048576.0,1157424966,153003780,261456620,256,128,160,1.19,67108,138412,8.0,7.175642,72.067339,50.544367 +1024,4096,28672,128,32,128,3,1.110944,1.048413,355381297,58233985,85.360376,4830676419,4035224853,61572651,4035224853,61572651,61572651,3042206940,88539024.0,7340032.0,8101975841,1224356663,3539359068,1792,128,164,5.53,50331,171966,12.0,11.455518,93.361916,65.234636 +1024,14336,4096,128,32,128,4,0.683584,1.048398,224542769,9869566,81.75315,2464348081,2017612426,8796093,2017612426,8796093,8796093,2017614574,58720320.0,3670016.0,4039729200,542552180,920171661,256,128,220,1.19,67108,138412,8.0,7.159021,75.865809,53.030863 +1024,4096,128256,128,32,128,3,4.833504,1.049829,1497540205,271802844,86.476888,21399952651,18050426842,275427662,18050426842,275427662,275427662,13608405235,396053200.0,32833536.0,36241873567,5497703094,14409956727,8016,128,164,24.74,50331,171966,12.0,11.866872,95.859057,66.896463 +1024,4096,1,64,32,64,2,0.153536,1.055087,8834881,0,5.093275,13975756,8933531,2097,8933531,1048,2097,0,0.0,0.0,35253091,0,0,16,128,92,0.03,16777,104857,20.0,3.991608,3.919674,6.088024 +1024,8192,10240,128,32,256,3,0.782656,1.047707,228733046,17291404,82.754327,2613189099,2161727907,21990232,2161727907,21990232,21990232,1804251336,52510800.0,3932160.0,5775866650,793798657,781186771,320,256,218,2.96,75497,104857,8.0,8.001965,94.722959,51.697916 +8192,128,1024,128,32,128,3,0.0272,1.021026,2500342,577241,92.176214,65492547,35184372,17592186,35184372,17592186,17592186,30811558,917904.0,65536.0,79164837,8204511,31703301,512,128,164,1.58,50331,171966,12.0,10.226931,34.959922,42.739866 +8192,1024,128,64,32,256,4,0.02576,1.020876,17896457,325451,48.999562,61575603,43980465,2199023,43980465,2199023,2199023,36009005,1073152.0,81920.0,71468255,17085379,43980465,64,128,162,0.3,83886,171966,8.0,3.929513,36.91961,34.962418 +1024,8192,8192,128,32,128,4,0.6597759999999999,1.047157,249002339,17269795,82.040977,2825939230,2305843387,17592186,2305843387,17592186,17592186,2305850903,67109008.0,4194304.0,4620693320,618336861,1076558248,512,128,220,2.37,67108,138412,8.0,7.465593,89.938927,62.866808 +1024,8192,57344,128,32,128,3,4.355616,1.049975,1686570306,122060017,83.414221,19465585280,16140901562,123145302,16140901562,123145302,123145302,12137223243,353238848.0,29360128.0,32344852171,4933365866,12677924192,3584,128,164,11.06,50331,171966,12.0,11.719001,95.109853,66.403134 +1024,28672,8192,128,32,128,4,2.246528,1.049467,985278114,17286303,80.305796,10016486504,8070450781,17592186,8070450781,17592186,17592186,8070457223,234881168.0,14680064.0,16149908109,2227016778,3781289207,512,128,220,2.37,67108,138412,8.0,7.465995,92.245134,64.417251 +1024,8192,128256,128,32,128,3,9.630464,1.04046,4266578635,274132864,81.167765,44408272844,36100854759,275427662,36100854759,275427662,275427662,27146214092,790056208.0,65667072.0,72342728326,11026818505,28633316454,8016,128,164,24.74,50331,171966,12.0,11.869756,97.089292,66.89006 +1024,8192,1,64,32,64,2,0.30160000000000003,1.057017,17639430,13762,3.106613,27316730,17867063,2097,17867063,1048,2097,0,0.0,0.0,70437463,0,0,16,128,92,0.03,16777,104857,20.0,3.991963,3.942794,6.156956 +1024,16384,18432,128,32,128,3,2.85632,1.040205,1199291234,39560138,81.461899,12735532589,10376293094,39582418,10376293094,39582418,39582418,7792364533,226787168.0,18874368.0,20772853066,3162183630,9015129911,1152,128,164,3.56,50331,171966,12.0,11.134014,94.111493,64.905723 +16384,128,1024,128,32,128,3,0.046784,1.021131,4700707,16606088,93.136978,137373253,70368744,35184372,70368744,35184372,35184372,61612379,1836096.0,131072.0,158329674,16498479,61105003,1024,128,164,3.16,50331,171966,12.0,11.292969,40.646908,50.773999 +16384,1024,128,128,32,256,3,0.042912,1.027601,35485556,2425045,36.116753,83941746,52776558,4398046,52776558,4398046,4398046,44530220,1327072.0,98304.0,142936511,18903895,52643934,64,256,152,0.59,75497,104857,8.0,7.990849,44.035483,43.610425 +1024,16384,16384,128,32,256,3,2.553408,1.04966,738906307,32129173,82.773994,8334510739,6917529087,35184372,6917529087,35184372,35184372,5769107446,167903312.0,12582912.0,18464758042,2609435297,2565774807,512,256,218,4.74,75497,104857,8.0,7.998441,92.735751,50.075447 +1024,16384,106496,128,32,256,3,15.763296,1.053701,5023419519,225993528,83.527452,52959598901,44963938533,228698418,44963938533,228698418,228698418,37499200547,1091370976.0,81788928.0,120020930498,16948851483,16659416858,3328,256,218,30.81,75497,104857,8.0,7.999406,97.266869,52.522158 +1024,53248,16384,128,32,256,3,8.191552,1.05317,2459006215,32135616,82.876442,26753356655,22481969803,35184372,22481969803,35184372,35184372,18739474888,545390592.0,40894464.0,59969932569,8494344726,8353434365,512,256,218,4.74,75497,104857,8.0,7.999665,93.634246,50.184621 +1024,16384,128256,128,32,128,3,18.9544,1.05332,9522446080,275487121,79.420551,89680657675,72201709518,275427662,72201709518,275427662,275427662,54221841471,1578060528.0,131334144.0,144544437844,22083531056,55151889796,8016,128,164,24.74,50331,171966,12.0,11.869806,97.454889,67.949664 +1024,16384,1,64,32,64,2,0.5973120000000001,1.061986,35250138,2350420,7.858025,57477903,35734127,2097,35734127,1048,2097,0,0.0,0.0,140806207,0,0,16,128,92,0.03,16777,104857,20.0,4.031643,3.942572,6.173377 From 2c63cffddd44fe54d7f3739ade343970a3177acf Mon Sep 17 00:00:00 2001 From: gkarfakis19 Date: Fri, 5 Dec 2025 00:50:50 -0800 Subject: [PATCH 72/77] added uci train tests validation --- tools/util_finder.py | 349 ++++++++++++++++++ validation_scripts/imec_data/uci_train.csv | 15 + validation_scripts/nvidia_inf.py | 54 ++- validation_scripts/nvidia_train.py | 92 +++-- validation_scripts/uci_train.py | 317 ++++++++++++++++ .../hardware-config/a100_80GB.yaml | 6 +- .../a100_80GB_korthikanti.yaml | 2 +- .../model-config/Llama2-7B.yaml | 20 + 8 files changed, 810 insertions(+), 45 deletions(-) create mode 100644 tools/util_finder.py create mode 100644 validation_scripts/imec_data/uci_train.csv create mode 100644 validation_scripts/uci_train.py create mode 100644 validation_scripts/validation_configs/model-config/Llama2-7B.yaml diff --git a/tools/util_finder.py b/tools/util_finder.py new file mode 100644 index 00000000..8f965f24 --- /dev/null +++ b/tools/util_finder.py @@ -0,0 +1,349 @@ +#!/usr/bin/env python3 +""" +Utility tuner (core util, HBM bandwidth util) for training validation. + +Modeled after overlap_finder.py but searches in two dimensions using a +coordinate-wise golden-section search. Objective is average absolute percent +error across Korthi+Selene training validation suites (nvidia_train.py). + +Global knobs (edit below): + - DEVICES: which training devices to validate (defaults: Korthi + Selene) + - APPLY_BEST: write the best utils back to all hardware configs via mod_field + - RANGES/TOL/ITER: search bounds and golden search controls + - COORD_STEPS: number of coordinate-descent passes (keep small; each eval is slow) +""" + +from __future__ import annotations + +import io +import math +import sys +from contextlib import redirect_stdout +from dataclasses import dataclass +from pathlib import Path +from typing import Any, Callable, Dict, Iterable, List, Mapping, Optional, Sequence, Tuple + +PROJECT_ROOT = Path(__file__).resolve().parents[1] +if str(PROJECT_ROOT) not in sys.path: + sys.path.insert(0, str(PROJECT_ROOT)) + +from validation_scripts import nvidia_train # type: ignore +from validation_scripts.validation_helpers import ValidationSpec, run_validation_suite # type: ignore +from tools import mod_field + +# ======= Global config ======= +DEVICES: Sequence[str] = ("A100_korthi", "A100_selene") +APPLY_BEST = False + +# Search bounds for utilities (floats in (0,1]). Set to None to skip tuning that axis. +CORE_UTIL_RANGE: Optional[Tuple[float, float]] = None +HBM_UTIL_RANGE: Optional[Tuple[float, float]] = (0.25, 0.55) +GOLD_TOL = 0.02 # stop when interval width < GOLD_TOL +GOLD_MAX_ITER = 6 # cap per-axis golden iterations +COORD_STEPS = 2 # how many coordinate-descent passes (keep small) +# ============================= + +# Exclude specific heavy/slow validation cases by metadata match. +EXCLUDE_TESTS: Tuple[Dict[str, Any], ...] = ( + # Skip 3k-GPU Selene GPT-1T case (dp=6, tp=8, pp=64, mb=512, batch=3072). + {"device": "A100_selene", "model": "GPT 1T", "pp": 64}, +) + + +def _merge_dict(target: Dict[str, Any], updates: Mapping[str, Any]) -> Dict[str, Any]: + for key, val in updates.items(): + if isinstance(val, Mapping): + current = target.get(key, {}) + if not isinstance(current, dict): + current = {} + target[key] = _merge_dict(dict(current), val) + else: + target[key] = val + return target + + +def _with_utils(spec: ValidationSpec, core_util: Optional[float], hbm_util: Optional[float]) -> ValidationSpec: + base_overrides: Dict[str, Any] = {} + if spec.hardware_overrides and isinstance(spec.hardware_overrides, Mapping): + base_overrides = _merge_dict({}, spec.hardware_overrides) + util_overrides: Dict[str, Any] = {"tech_param": {}} + if core_util is not None: + util_overrides["tech_param"]["core"] = {"util": float(core_util)} + if hbm_util is not None: + util_overrides["tech_param"]["DRAM"] = {"util": float(hbm_util)} + if not util_overrides["tech_param"]: + return ValidationSpec( + label=spec.label, + model_overrides=spec.model_overrides, + hardware_overrides=base_overrides, + metadata=spec.metadata, + model_config_path=spec.model_config_path, + hardware_config_path=spec.hardware_config_path, + order=spec.order, + ) + base_overrides = _merge_dict(base_overrides, util_overrides) + return ValidationSpec( + label=spec.label, + model_overrides=spec.model_overrides, + hardware_overrides=base_overrides, + metadata=spec.metadata, + model_config_path=spec.model_config_path, + hardware_config_path=spec.hardware_config_path, + order=spec.order, + ) + + +def _avg_pct_error(errors: Iterable[float]) -> float: + vals = [e for e in errors if not math.isnan(e)] + if not vals: + return float("inf") + return sum(vals) / len(vals) + + +def _all_hardware_configs() -> List[Path]: + roots = [ + PROJECT_ROOT / "configs" / "hardware-config", + PROJECT_ROOT / "validation_scripts" / "validation_configs" / "hardware-config", + ] + paths: List[Path] = [] + for root in roots: + if root.exists(): + paths.extend(sorted(root.glob("*.yaml"))) + return paths + + +def _should_exclude(meta: Mapping[str, Any]) -> bool: + for rule in EXCLUDE_TESTS: + matches = True + for key, val in rule.items(): + if str(meta.get(key)) != str(val): + matches = False + break + if matches: + return True + return False + + +def _filter_specs( + specs: Sequence[ValidationSpec], + actual_lookup: Mapping[Tuple[str, int, int, int, int, int, int, bool, str], float], +) -> Tuple[List[ValidationSpec], Dict[Tuple[str, int, int, int, int, int, int, bool, str], float]]: + kept: List[ValidationSpec] = [] + kept_lookup: Dict[Tuple[str, int, int, int, int, int, int, bool, str], float] = {} + for spec in specs: + meta = spec.metadata or {} + if _should_exclude(meta): + continue + key = ( + meta.get("model"), + int(meta.get("batch")), + int(meta.get("mb")), + int(meta.get("dp")), + int(meta.get("tp")), + int(meta.get("pp")), + int(meta.get("cp")), + bool(meta.get("tp_sp")), + str(meta.get("recomputation")), + ) + if key in actual_lookup: + kept_lookup[key] = actual_lookup[key] + kept.append(spec) + return kept, kept_lookup + + +@dataclass +class SearchResult: + best_x: float + best_y: float + evaluations: List[Tuple[float, float]] + + +def golden_section_search( + fn: Callable[[float], float], + lo: float, + hi: float, + tol: float = GOLD_TOL, + max_iter: int = GOLD_MAX_ITER, + log_fn: Optional[Callable[[int, int, float, float], None]] = None, +) -> SearchResult: + phi = (1 + 5 ** 0.5) / 2 + inv_phi = 1 / phi + inv_phi_sq = inv_phi ** 2 + + a, b = lo, hi + h = b - a + if h <= tol: + mid = (a + b) / 2 + val = fn(mid) + return SearchResult(best_x=mid, best_y=val, evaluations=[(mid, val)]) + + n = int(math.ceil(math.log(tol / h) / math.log(inv_phi))) + c = a + inv_phi_sq * h + d = a + inv_phi * h + yc = fn(c) + yd = fn(d) + if log_fn: + log_fn(1, max_iter, c, yc) + log_fn(2, max_iter, d, yd) + evals: List[Tuple[float, float]] = [(c, yc), (d, yd)] + + iter_idx = 2 + for _ in range(min(n, max_iter)): + if yc < yd: + b = d + d, yd = c, yc + h = inv_phi * h + c = a + inv_phi_sq * h + yc = fn(c) + evals.append((c, yc)) + else: + a = c + c, yc = d, yd + h = inv_phi * h + d = a + inv_phi * h + yd = fn(d) + evals.append((d, yd)) + iter_idx += 1 + if log_fn: + log_fn(iter_idx, max_iter, evals[-1][0], evals[-1][1]) + if h < tol: + break + + best_x, best_y = min(evals, key=lambda t: t[1]) + return SearchResult(best_x=best_x, best_y=best_y, evaluations=evals) + + +class UtilTuner: + def __init__(self) -> None: + self._eval_cache: Dict[Tuple[float, float], float] = {} + + def eval_pair(self, core_util: float, hbm_util: float) -> float: + key = (round(core_util, 6), round(hbm_util, 6)) + if key in self._eval_cache: + print(f"[CACHE HIT] core.util={core_util:.4f}, DRAM.util={hbm_util:.4f} -> {self._eval_cache[key]:.3f}") + return self._eval_cache[key] + + all_errors: List[float] = [] + for device in DEVICES: + specs, actual_lookup, base_model_path, base_hw_path = nvidia_train.build_specs_for_device(device) + specs, actual_lookup = _filter_specs(specs, actual_lookup) + if not specs: + print(f"[SKIP DEVICE] {device}: no specs after filtering.") + continue + specs = [_with_utils(spec, core_util, hbm_util) for spec in specs] + print(f"[EVAL] device={device}, specs={len(specs)}, core.util={core_util:.4f}, DRAM.util={hbm_util:.4f}") + with io.StringIO() as buf, redirect_stdout(buf): + results = run_validation_suite( + specs, + base_model_config_path=base_model_path, + base_hardware_config_path=base_hw_path, + result_parser=nvidia_train.parse_training_time, + run_perf_path=nvidia_train.RUN_PERF, + ) + rows = nvidia_train.compute_pct_errors(results, actual_lookup) + all_errors.extend(float(row["pct_error"]) for row in rows) + print(f"[RESULTS] device={device}, pct_errors={ [float(row['pct_error']) for row in rows] }") + + avg_error = _avg_pct_error(all_errors) + self._eval_cache[key] = avg_error + print(f"[EVAL DONE] core.util={core_util:.4f}, DRAM.util={hbm_util:.4f}, avg_abs_pct_err={avg_error:.3f}") + return avg_error + + def optimize(self) -> Tuple[Optional[float], Optional[float], float]: + core_util = None if CORE_UTIL_RANGE is None else (CORE_UTIL_RANGE[0] + CORE_UTIL_RANGE[1]) / 2 + hbm_util = None if HBM_UTIL_RANGE is None else (HBM_UTIL_RANGE[0] + HBM_UTIL_RANGE[1]) / 2 + + if CORE_UTIL_RANGE is None and HBM_UTIL_RANGE is None: + raise ValueError("At least one of CORE_UTIL_RANGE or HBM_UTIL_RANGE must be set.") + + # If only one axis is active, do a single golden search on that axis. + if CORE_UTIL_RANGE is None and HBM_UTIL_RANGE is not None: + h_lo, h_hi = HBM_UTIL_RANGE + print("Tuning only DRAM.util") + res_hbm = golden_section_search( + lambda y: self.eval_pair(core_util or 1.0, y), + lo=h_lo, + hi=h_hi, + tol=GOLD_TOL, + max_iter=GOLD_MAX_ITER, + log_fn=lambda i, m, y, v: print(f" [{i}/{m}] hbm_util={y:.4f} -> {v:.3f}"), + ) + best_val = res_hbm.best_y + return core_util, res_hbm.best_x, best_val + + if HBM_UTIL_RANGE is None and CORE_UTIL_RANGE is not None: + c_lo, c_hi = CORE_UTIL_RANGE + print("Tuning only core.util") + res_core = golden_section_search( + lambda x: self.eval_pair(x, hbm_util or 1.0), + lo=c_lo, + hi=c_hi, + tol=GOLD_TOL, + max_iter=GOLD_MAX_ITER, + log_fn=lambda i, m, x, y: print(f" [{i}/{m}] core_util={x:.4f} -> {y:.3f}"), + ) + best_val = res_core.best_y + return res_core.best_x, hbm_util, best_val + + # Both axes active: coordinate descent. + core_lo, core_hi = CORE_UTIL_RANGE # type: ignore[assignment] + hbm_lo, hbm_hi = HBM_UTIL_RANGE # type: ignore[assignment] + core = float(core_util) + hbm = float(hbm_util) + best_val = self.eval_pair(core, hbm) + + for step in range(COORD_STEPS): + print(f"\n--- Coordinate step {step + 1}/{COORD_STEPS} ---") + # Optimize core util with HBM fixed + print(f"Optimizing core util (hbm_util={hbm:.4f})") + res_core = golden_section_search( + lambda x: self.eval_pair(x, hbm), + lo=core_lo, + hi=core_hi, + tol=GOLD_TOL, + max_iter=GOLD_MAX_ITER, + log_fn=lambda i, m, x, y: print(f" [{i}/{m}] core_util={x:.4f} -> {y:.3f}"), + ) + core = res_core.best_x + best_val = res_core.best_y + + # Optimize HBM util with core fixed + print(f"Optimizing HBM util (core_util={core:.4f})") + res_hbm = golden_section_search( + lambda y: self.eval_pair(core, y), + lo=hbm_lo, + hi=hbm_hi, + tol=GOLD_TOL, + max_iter=GOLD_MAX_ITER, + log_fn=lambda i, m, y, v: print(f" [{i}/{m}] hbm_util={y:.4f} -> {v:.3f}"), + ) + hbm = res_hbm.best_x + best_val = min(best_val, res_hbm.best_y) + + best_val = self.eval_pair(core, hbm) + return core, hbm, best_val + + +def main() -> None: + tuner = UtilTuner() + best_core, best_hbm, best_err = tuner.optimize() + print("\n=== Best utils ===") + print(f" core.util = {best_core if best_core is None else f'{best_core:.4f}'}") + print(f" DRAM.util = {best_hbm if best_hbm is None else f'{best_hbm:.4f}'}") + print(f" avg abs pct error = {best_err:.3f}") + + if APPLY_BEST: + core_val = None if best_core is None else round(best_core, 2) + hbm_val = None if best_hbm is None else round(best_hbm, 2) + print("\nApplying tuned utils to all hardware configs") + for cfg_path in _all_hardware_configs(): + if core_val is not None: + msg1 = mod_field.set_field(cfg_path, "tech_param.core.util", core_val, dry_run=False) + print(f" {msg1}") + if hbm_val is not None: + msg2 = mod_field.set_field(cfg_path, "tech_param.DRAM.util", hbm_val, dry_run=False) + print(f" {msg2}") + + +if __name__ == "__main__": + main() diff --git a/validation_scripts/imec_data/uci_train.csv b/validation_scripts/imec_data/uci_train.csv new file mode 100644 index 00000000..7abf0763 --- /dev/null +++ b/validation_scripts/imec_data/uci_train.csv @@ -0,0 +1,15 @@ +variant,TP,PP,CP,DP,Status,Error,Avg_Step_Time_s,Total_Throughput_tokens_s,Throughput_per_GPU_tokens_s +DDP,1,1,1,4,SUCCESS,,36.9333,14195.52,3548.88 +DDP,4,1,1,1,SUCCESS,,70.3667,7450.8,1862.7 +DDP,1,4,1,1,SUCCESS,,40.5667,12924.11,3231.03 +DDP,1,1,4,1,SUCCESS,,107.0,4899.89,1224.97 +DDP,2,2,1,1,SUCCESS,,42.9722,12200.63,3050.16 +DDP,2,1,2,1,SUCCESS,,63.2944,8283.32,2070.83 +DDP,1,2,2,1,SUCCESS,,43.9556,11927.68,2981.92 +DDP,2,1,1,2,SUCCESS,,41.7444,12559.47,3139.87 +DDP,1,2,1,2,SUCCESS,,38.2194,13717.83,3429.46 +DDP,1,1,2,2,SUCCESS,,43.025,12185.66,3046.41 +FSDP,1,1,1,4,SUCCESS,,56.2639,9318.37,2329.59 +FSDP,2,1,1,2,SUCCESS,,74.1639,7069.32,1767.33 +FSDP,1,2,1,2,INSUFFICIENT_DATA,AssertionError: FSDP does not support pipeline parallelism,,, +FSDP,1,1,2,2,SUCCESS,,104.2222,5030.48,1257.62 diff --git a/validation_scripts/nvidia_inf.py b/validation_scripts/nvidia_inf.py index 1d8ff2a7..eb73cfd3 100644 --- a/validation_scripts/nvidia_inf.py +++ b/validation_scripts/nvidia_inf.py @@ -36,6 +36,12 @@ "Llama 2-70B": "Llama2-70B_inf.yaml", } +MODEL_DISPLAY = { + "Llama 2-7B": "Llama 2-7B", + "Llama 2-13B": "Llama 2-13B", + "Llama 2-70B": "Llama 2-70B", +} + PROJECT_ROOT = os.path.dirname(os.path.dirname(os.path.abspath(__file__))) RUN_PERF = os.path.join(PROJECT_ROOT, "run_perf.py") VALIDATION_CONFIG_ROOT = os.path.join(PROJECT_ROOT, "validation_scripts", "validation_configs") @@ -230,6 +236,7 @@ def compute_pct_errors(results, actual_lookup: Dict[Tuple[str, int], float]): "inference_time_s": inf_time, "actual_inference_time_s": actual, "pct_error": pct_error, + "display_model": MODEL_DISPLAY.get(str(model), str(model)), "success": res.success, "error": res.error, } @@ -250,6 +257,7 @@ def plot_device(df: pd.DataFrame, device: str, outdir: Path) -> Path | None: models = [m for m in desired_order if m in unique_models] + [ m for m in unique_models if m not in desired_order ] + display_models = [MODEL_DISPLAY.get(m, m) for m in models] bar_width = 0.28 # width per individual bar tp_gap = 0.08 # gap between TP subgroups inside a model cluster @@ -295,12 +303,12 @@ def plot_device(df: pd.DataFrame, device: str, outdir: Path) -> Path | None: x_seconds, seconds_vals, bar_width, label="DeepFlow", color="#1f77b4" ) bars_actual = ax.bar( - x_actual, actual_vals, bar_width, label="Actual [1]", color="#0bbd37" + x_actual, actual_vals, bar_width, label="Actual", color="#0bbd37" ) # Primary ticks: model names at cluster centers ax.set_xticks(model_centers) - ax.set_xticklabels(models, fontsize=11) + ax.set_xticklabels(display_models, fontsize=11) ax.set_ylabel("Inference Latency (s)") ax.set_title( f"Validation of Inference Latency on Systems of {device} (network ignored)" @@ -329,11 +337,6 @@ def plot_device(df: pd.DataFrame, device: str, outdir: Path) -> Path | None: ax.grid(axis="y", linestyle="--", alpha=0.3) ax.margins(x=0.02, y=0.1) - # Add citation below the graph - citation = "[1] J. Kundu et al., “Performance Modeling and Workload Analysis of Distributed Large Language Model Training and Inference,” (2024)" - plt.subplots_adjust(bottom=0.1) - fig.text(0.5, 0.015, citation, ha="center", va="bottom", fontsize=8) - outpath = outdir / f"inf_{device}.png" fig.savefig(outpath, dpi=200, bbox_inches="tight") plt.close(fig) @@ -399,6 +402,13 @@ def run( out = plot_device(data, device, out_dir) if out is not None: outputs.append(out) + # Secondary bar plot matching the training style (combined devices). + bar_path = _plot_error_bars( + rows, + path=out_dir / "inf_errors_bar.png", + title="Inference validation (combined)", + ) + outputs.append(bar_path) if not outputs: print("No plots generated (no matching device rows).") else: @@ -411,5 +421,35 @@ def run( return {"avg_abs_error": avg_abs_error, "rows": rows} +def _plot_error_bars(rows: List[Dict[str, object]], path: Path, title: str) -> Path: + # Build labels like " xTP", drop device prefixes. + labels: List[str] = [] + errors: List[float] = [] + for row in rows: + model = row.get("display_model") or row.get("model") + tp = row.get("tp") + labels.append(f"{model} xTP{tp}") + errors.append(float(row.get("pct_error", float("nan")))) + + fig_w = max(6.0, 0.6 * len(labels)) + fig, ax = plt.subplots(figsize=(fig_w, 4)) + bars = ax.bar(range(len(errors)), errors, color="#1f77b4") + ax.set_xticks(range(len(labels))) + ax.set_xticklabels(labels, rotation=20, ha="right", fontsize=9) + ax.set_ylabel("Percent Error") + ax.set_title(title) + for rect, err in zip(bars, errors): + if math.isnan(err): + continue + height = rect.get_height() + ax.text(rect.get_x() + rect.get_width() / 2, height, f"{err:.1f}%", ha="center", va="bottom", fontsize=8) + ax.grid(axis="y", linestyle="--", alpha=0.3) + fig.tight_layout() + path.parent.mkdir(parents=True, exist_ok=True) + fig.savefig(path, dpi=200) + plt.close(fig) + return path + + if __name__ == "__main__": run(enable_plot=True, network_ignored=True) diff --git a/validation_scripts/nvidia_train.py b/validation_scripts/nvidia_train.py index 3c2e27fd..b75d3c73 100644 --- a/validation_scripts/nvidia_train.py +++ b/validation_scripts/nvidia_train.py @@ -304,39 +304,12 @@ def run( plot_path = None if enable_plot and rows: out_dir = Path(PROJECT_ROOT) / "output" / "validation" / "train" - out_dir.mkdir(parents=True, exist_ok=True) - labels = [] - errors = [] - colors = [] - color_map = {"full": "#ff8c00", "partial": "#1f77b4", "selective": "#1f77b4"} - for row in rows: - gpus = int(row["dp"]) * int(row["tp"]) * int(row["pp"]) * int(row["cp"]) - labels.append(f"{row['model']} x{gpus}") - errors.append(row["signed_pct_error"]) - colors.append(color_map.get(str(row.get("recomputation")).lower(), "#1f77b4")) - fig_w = max(6.0, 0.6 * len(labels)) - fig, ax = plt.subplots(figsize=(fig_w, 4)) - bars = ax.bar(range(len(errors)), errors, color=colors) - ax.set_xticks(range(len(labels))) - ax.set_xticklabels(labels, rotation=20, ha="right", fontsize=9) - ax.set_ylabel("Percent Error") - ax.set_title(f"Training validation ({device})") - # Legend - handles = [ - plt.Rectangle((0, 0), 1, 1, color="#ff8c00"), - plt.Rectangle((0, 0), 1, 1, color="#1f77b4"), - ] - ax.legend(handles, ["recompute=full", "recompute=partial/selective"]) - for rect, err in zip(bars, errors): - if math.isnan(err): - continue - height = rect.get_height() - ax.text(rect.get_x() + rect.get_width() / 2, height, f"{err:.1f}%", ha="center", va="bottom", fontsize=8) - ax.grid(axis="y", linestyle="--", alpha=0.3) - plot_path = out_dir / f"train_{device}.png" - fig.tight_layout() - fig.savefig(plot_path, dpi=200) - plt.close(fig) + plot_path = _plot_training_rows( + rows, + title=f"Training validation ({device})", + path=out_dir / f"train_{device}.png", + include_device=False, + ) avg_abs_error = sum(pct_errors) / len(pct_errors) if pct_errors else float("nan") if emit_logs: @@ -346,11 +319,62 @@ def run( return {"avg_abs_error": avg_abs_error, "rows": rows, "plot": plot_path} +def _plot_training_rows(rows, title: str, path: Path, include_device: bool = True) -> Path: + out_dir = path.parent + out_dir.mkdir(parents=True, exist_ok=True) + labels: List[str] = [] + errors: List[float] = [] + colors: List[str] = [] + color_map = {"full": "#ff8c00", "partial": "#1f77b4", "selective": "#1f77b4"} + for row in rows: + if str(row.get("model")) == "GPT 22B": + continue + gpus = int(row["dp"]) * int(row["tp"]) * int(row["pp"]) * int(row["cp"]) + prefix = "" + labels.append(f"{prefix}{row['model']} x{gpus}") + errors.append(row["signed_pct_error"]) + colors.append(color_map.get(str(row.get("recomputation")).lower(), "#1f77b4")) + + fig_w = max(6.0, 0.6 * len(labels)) + fig, ax = plt.subplots(figsize=(fig_w, 4)) + bars = ax.bar(range(len(errors)), errors, color=colors) + ax.set_xticks(range(len(labels))) + ax.set_xticklabels(labels, rotation=20, ha="right", fontsize=9) + ax.set_ylabel("Percent Error") + ax.set_title(title) + handles = [ + plt.Rectangle((0, 0), 1, 1, color="#ff8c00"), + plt.Rectangle((0, 0), 1, 1, color="#1f77b4"), + ] + ax.legend(handles, ["recompute=full", "recompute=partial/selective"]) + for rect, err in zip(bars, errors): + if math.isnan(err): + continue + height = rect.get_height() + ax.text(rect.get_x() + rect.get_width() / 2, height, f"{err:.1f}%", ha="center", va="bottom", fontsize=8) + ax.grid(axis="y", linestyle="--", alpha=0.3) + fig.tight_layout() + fig.savefig(path, dpi=200) + plt.close(fig) + return path + + if __name__ == "__main__": devices = [ "A100_korthi", # Uncomment to include korthi runs. "A100_selene", ] + combined_rows: List[Dict[str, object]] = [] for device in devices: print(f"=== Running {device} training validation ===") - run(device=device, emit_logs=True, show_progress=True) + result = run(device=device, emit_logs=True, show_progress=True) + combined_rows.extend(result.get("rows", [])) + + if combined_rows: + combined_path = _plot_training_rows( + combined_rows, + title="Training validation (combined)", + path=Path(PROJECT_ROOT) / "output" / "validation" / "train" / "train_combined.png", + include_device=False, + ) + print(f"Saved combined plot: {combined_path}") diff --git a/validation_scripts/uci_train.py b/validation_scripts/uci_train.py new file mode 100644 index 00000000..30067dd2 --- /dev/null +++ b/validation_scripts/uci_train.py @@ -0,0 +1,317 @@ +#!/usr/bin/env python3 +""" +UCI training validation harness combining DDP and FSDP grid-search measurements. + +Inspired by nvidia_train.py: builds ValidationSpecs from benchmark CSVs, +runs DeepFlow, compares training time to measured Avg_Step_Time_s, and +produces bar plots of percent error. +""" + +from __future__ import annotations + +import argparse +import math +import os +from pathlib import Path +from typing import Dict, Iterable, List, Mapping, Optional, Sequence, Tuple +import sys +import pandas as pd +import matplotlib.pyplot as plt + +try: + from .validation_helpers import ( + ValidationSpec, + run_validation_suite, + parse_training_time, + ) +except ImportError: + sys.path.append(str(Path(__file__).resolve().parents[1])) + from validation_scripts.validation_helpers import ( # type: ignore + ValidationSpec, + run_validation_suite, + parse_training_time, + ) + + +PROJECT_ROOT = Path(__file__).resolve().parents[1] +RUN_PERF = PROJECT_ROOT / "run_perf.py" +DEFAULT_HW_CONFIG = PROJECT_ROOT / "validation_scripts" / "validation_configs" / "hardware-config" / "a100_80GB.yaml" +DEFAULT_MODEL_CONFIG = PROJECT_ROOT / "validation_scripts" / "validation_configs" / "model-config" / "Llama2-7B.yaml" + +# Unified benchmark CSV with variant column (DDP/FSDP) +MERGED_CSV = PROJECT_ROOT / "validation_scripts" / "imec_data" / "uci_train.csv" + + +def _load_grid(csv_path: Path, default_variant: Optional[str] = None) -> pd.DataFrame: + df = pd.read_csv(csv_path) + if "variant" not in df.columns and default_variant is not None: + df.insert(0, "variant", default_variant) + # Normalize numeric fields + for col in ("TP", "PP", "CP", "DP"): + if col in df.columns: + df[col] = pd.to_numeric(df[col], errors="coerce").astype("Int64") + if "Avg_Step_Time_s" in df.columns: + df["Avg_Step_Time_s"] = pd.to_numeric(df["Avg_Step_Time_s"], errors="coerce") + # Keep only successes and drop unused columns + df = df[df["Status"].str.upper() == "SUCCESS"].copy() + for col in ("Status", "Error"): + if col in df.columns: + df.drop(columns=[col], inplace=True) + return df + + +def _iter_tests(df: pd.DataFrame): + for _, row in df.iterrows(): + actual = row.get("Avg_Step_Time_s") + if actual is None or pd.isna(actual): + continue + yield ( + str(row.get("variant")), + int(row["TP"]), + int(row["PP"]), + int(row["CP"]), + int(row["DP"]), + float(actual), + ) + + +def _build_spec( + variant: str, + tp: int, + pp: int, + cp: int, + dp: int, + idx: int, + model_config_path: str, + hardware_config_path: str, +) -> ValidationSpec: + label = f"{variant} TP={tp} PP={pp} CP={cp} DP={dp}" + hw_overrides = { + "parallelism": { + "dp": int(dp), + "tp": int(tp), + "tp_sp": False, + "cp": int(cp), + "lp": int(pp), + "mb": int(pp), # align microbatch count with pipeline stages + }, + "sw_param": { + # DDP -> zero_stage 0, FSDP -> zero_stage 3 + "dp_zero_stage": 0 if variant.upper() == "DDP" else 3, + }, + } + return ValidationSpec( + label=label, + model_overrides={ + "model_param": { + "seq_len": 4096, + "run_type": "training", + } + }, + hardware_overrides=hw_overrides, + model_config_path=model_config_path, + hardware_config_path=hardware_config_path, + metadata={ + "variant": variant, + "tp": int(tp), + "pp": int(pp), + "cp": int(cp), + "dp": int(dp), + }, + order=idx, + ) + + +def build_specs( + variants: Sequence[str], + model_config_path: str, + hardware_config_path: str, +) -> Tuple[List[ValidationSpec], Dict[Tuple[str, int, int, int, int], float]]: + specs: List[ValidationSpec] = [] + actual_lookup: Dict[Tuple[str, int, int, int, int], float] = {} + idx = 0 + df = _load_grid(MERGED_CSV) + df = df[df["variant"].str.lower().isin(set(variants))] + for variant, tp, pp, cp, dp, actual in _iter_tests(df): + spec = _build_spec(variant, tp, pp, cp, dp, idx, model_config_path, hardware_config_path) + specs.append(spec) + actual_lookup[(variant, tp, pp, cp, dp)] = actual + idx += 1 + + if not specs: + raise ValueError(f"No validation specs generated for variants={variants}") + return specs, actual_lookup + + +def compute_pct_errors(results, actual_lookup: Dict[Tuple[str, int, int, int, int], float]): + rows: List[Dict[str, object]] = [] + for res in results: + meta = res.spec.metadata or {} + variant = str(meta.get("variant")) + tp = int(meta.get("tp", 0)) + pp = int(meta.get("pp", 0)) + cp = int(meta.get("cp", 0)) + dp = int(meta.get("dp", 0)) + inf_time = float(res.metrics.get("training_time_s", float("nan"))) if res.success else float("nan") + actual = actual_lookup.get((variant, tp, pp, cp, dp), float("nan")) + + if math.isnan(inf_time) or actual == 0 or math.isnan(actual): + signed_pct_error = float("nan") + pct_error = float("nan") + else: + signed_pct_error = (inf_time - actual) / actual * 100.0 + pct_error = abs(signed_pct_error) + + rows.append( + { + "variant": variant, + "tp": tp, + "pp": pp, + "cp": cp, + "dp": dp, + "training_time_s": inf_time, + "actual_training_time_s": actual, + "signed_pct_error": signed_pct_error, + "pct_error": pct_error, + "success": res.success, + "error": res.error, + "raw_output": res.raw_output, + } + ) + return rows + + +def _plot_results(rows: List[Dict[str, object]], title: str, path: Path) -> Optional[Path]: + if not rows: + return None + labels: List[str] = [] + errors: List[float] = [] + colors: List[str] = [] + color_map = {"DDP": "#1f77b4", "FSDP": "#ff8c00"} + + for row in rows: + variant = str(row.get("variant")) + tp = int(row.get("tp", 0)) + pp = int(row.get("pp", 0)) + cp = int(row.get("cp", 0)) + dp = int(row.get("dp", 0)) + labels.append(f"{variant} TP{tp} PP{pp} CP{cp} DP{dp}") + errors.append(float(row.get("signed_pct_error", float("nan")))) + colors.append(color_map.get(variant.upper(), "#888888")) + + fig_w = max(6.0, 0.65 * len(labels)) + fig, ax = plt.subplots(figsize=(fig_w, 4)) + bars = ax.bar(range(len(errors)), errors, color=colors) + ax.set_xticks(range(len(labels))) + ax.set_xticklabels(labels, rotation=25, ha="right", fontsize=9) + ax.set_ylabel("Percent Error") + ax.set_title(title) + handles = [ + plt.Rectangle((0, 0), 1, 1, color="#1f77b4"), + plt.Rectangle((0, 0), 1, 1, color="#ff8c00"), + ] + ax.legend(handles, ["DDP", "FSDP"]) + for rect, err in zip(bars, errors): + if math.isnan(err): + continue + height = rect.get_height() + ax.text(rect.get_x() + rect.get_width() / 2, height, f"{err:.1f}%", ha="center", va="bottom", fontsize=8) + ax.grid(axis="y", linestyle="--", alpha=0.3) + fig.tight_layout() + path.parent.mkdir(parents=True, exist_ok=True) + fig.savefig(path, dpi=200) + plt.close(fig) + return path + + +def run( + *, + hardware_config: Optional[str] = None, + model_config: Optional[str] = None, + variants: Sequence[str] = ("ddp", "fsdp"), + enable_plot: bool = True, + show_progress: bool = False, + emit_logs: bool = True, +): + variants = tuple(v.lower() for v in variants) + hw_cfg = hardware_config or str(DEFAULT_HW_CONFIG) + model_cfg = model_config or str(DEFAULT_MODEL_CONFIG) + outputs: List[Dict[str, object]] = [] + for variant in variants: + specs, actual_lookup = build_specs((variant,), model_cfg, hw_cfg) + + validation_results = run_validation_suite( + specs, + base_model_config_path=model_cfg, + base_hardware_config_path=hw_cfg, + result_parser=parse_training_time, + run_perf_path=str(RUN_PERF), + show_progress=show_progress, + ) + + rows = compute_pct_errors(validation_results, actual_lookup) + pct_errors = [r["pct_error"] for r in rows if not math.isnan(r["pct_error"])] + + if emit_logs: + for row in rows: + block = [ + f"\n=== Result ({row['variant']}, TP={row['tp']}, PP={row['pp']}, CP={row['cp']}, DP={row['dp']}) ===" + ] + if row["success"] and not math.isnan(row["pct_error"]): + block.append(f" DeepFlow train time: {float(row['training_time_s']):.2f}s") + block.append(f" Actual train time: {float(row['actual_training_time_s']):.2f}s") + block.append(f" Percent Error: {float(row['signed_pct_error']):+.2f}%") + else: + block.append(f" DeepFlow run failed. {(row.get('error') or '')}".rstrip()) + if row.get("raw_output"): + block.append(" --- Raw output ---") + block.append(str(row["raw_output"]).strip()) + print("\n".join(block)) + + plot_path = None + if enable_plot and rows: + out_dir = PROJECT_ROOT / "output" / "validation" / "train" + plot_path = _plot_results(rows, f"Training validation ({variant.upper()} benchmarks)", out_dir / f"uci_train_{variant}.png") + if emit_logs and plot_path: + print(f"Saved plot: {plot_path}") + + avg_abs_error = sum(pct_errors) / len(pct_errors) if pct_errors else float("nan") + if emit_logs: + print(f"Average absolute percent error for {variant.upper()}: {avg_abs_error:.2f}%") + outputs.append({"variant": variant, "avg_abs_error": avg_abs_error, "rows": rows, "plot": plot_path}) + return outputs + + +def _parse_args() -> argparse.Namespace: + parser = argparse.ArgumentParser(description="UCI training validation (DDP + FSDP gridsearch).") + parser.add_argument( + "--hardware_config", + default=str(DEFAULT_HW_CONFIG), + help=f"Path to hardware config YAML (default: {DEFAULT_HW_CONFIG}).", + ) + parser.add_argument( + "--model_config", + default=str(DEFAULT_MODEL_CONFIG), + help=f"Path to model config YAML (default: {DEFAULT_MODEL_CONFIG}).", + ) + parser.add_argument( + "--variants", + nargs="+", + default=["ddp", "fsdp"], + choices=["ddp", "fsdp"], + help="Which benchmark variants to include (runs each separately).", + ) + parser.add_argument("--no-plot", dest="enable_plot", action="store_false", help="Disable plot generation.") + parser.add_argument("--show-progress", action="store_true", help="Show per-run progress.") + return parser.parse_args() + + +if __name__ == "__main__": + args = _parse_args() + run( + hardware_config=args.hardware_config, + model_config=args.model_config, + variants=args.variants, + enable_plot=args.enable_plot, + show_progress=args.show_progress, + ) diff --git a/validation_scripts/validation_configs/hardware-config/a100_80GB.yaml b/validation_scripts/validation_configs/hardware-config/a100_80GB.yaml index 0a53f4d8..c7fac8e4 100644 --- a/validation_scripts/validation_configs/hardware-config/a100_80GB.yaml +++ b/validation_scripts/validation_configs/hardware-config/a100_80GB.yaml @@ -79,11 +79,11 @@ memory_hierarchy: parallelism: dp: 2 # data parallelism. Acts as replica count during inference tp: 2 # tensor parallelism - tp_sp: true # tensor-sequence parallelism + tp_sp: false # tensor-sequence parallelism cp: 2 # context parallelism. Not supported for inference. lp: 4 # pipeline parallelism. Number of pipeline stages # mb is not a true parallelism parameter, does not affect num of GPUs. - mb: 4 # micro-batches (for best results set to >= lp). + mb: 1 # micro-batches (for best results set to >= lp). network: # Global faulty links apply to the active AstraSim dimension: [[src, dst, weight]] @@ -136,7 +136,7 @@ execution_backend: model: astra # analytical | astra astra: backend: astra - mode: full_astrasim_flattened # hybrid | full_astrasim_hierarchical | full_astrasim_flattened + mode: full_astrasim_hierarchical # hybrid | full_astrasim_hierarchical | full_astrasim_flattened collectives: # Default collective algorithms. Per-dimension collective_override entries can override these. # Allowed values per op: auto | ring | direct | halvingDoubling | doubleBinaryTree | mesh | mesh2d | hypercube | torus2d diff --git a/validation_scripts/validation_configs/hardware-config/a100_80GB_korthikanti.yaml b/validation_scripts/validation_configs/hardware-config/a100_80GB_korthikanti.yaml index a34e6fb8..e0b3b35a 100644 --- a/validation_scripts/validation_configs/hardware-config/a100_80GB_korthikanti.yaml +++ b/validation_scripts/validation_configs/hardware-config/a100_80GB_korthikanti.yaml @@ -42,7 +42,7 @@ tech_param: bandwidth: 1986 GB # (/s) dynamic_energy_per_bit: 4.5e-12 # (J) latency: 100e-9 # (s) - util: 0.7 + util: 0.4 SRAM-L2: size: 40 MB bandwidth: 7050 GB diff --git a/validation_scripts/validation_configs/model-config/Llama2-7B.yaml b/validation_scripts/validation_configs/model-config/Llama2-7B.yaml new file mode 100644 index 00000000..2d13a809 --- /dev/null +++ b/validation_scripts/validation_configs/model-config/Llama2-7B.yaml @@ -0,0 +1,20 @@ +model_param: + mode: "LLM" # Mode of operation + run_type: "training" # Type of run: "training" or "inference" + tied_embeddings: true # Share input and output embedding weights (WIP: affects memory estimation only) + model_type: "llama" # gpt assumes dense-style architecture with GELU; llama assumes llama-style architecture with SwiGLU + global_batch_size: 128 # Sequences per global batch (validation default) + gradient_accumulation_steps: 1 # Number of gradient accumulation micro-steps per optimizer update + seq_len: 2048 # Sequence length (tokens per sequence) + hidden_dim: 4096 # Transformer hidden dimension + attention: # Attention-related parameters + attention_type: "mha" # Supported attention types: "mha" or "gqa" ("mla" not implemented) + num_heads: 32 # Total number of attention heads + kv_heads: null # Number of key/value heads (used only when attention_type = "gqa") + use_flashattention: true # Enable FlashAttention optimization + attention_tile_size: 128 # Tile size used in attention computation + num_experts: 1 # Number of experts in MoE layers. Set to 1 to disable MoE. + top_k: 1 # Number of active experts per token in MoE + intermediate_size: 11008 # Output dimension of the first FFN layer (typically 4×hidden_dim) + vocab_size: 32000 # Vocabulary size + num_layers: 32 # Number of transformer layers From 90b562cce5ccd4cd687e106b1dffc1f39280becb Mon Sep 17 00:00:00 2001 From: NanoMembers Date: Fri, 5 Dec 2025 14:03:57 -0800 Subject: [PATCH 73/77] Disable SM util for flash attention --- time_calculation.py | 15 +++++++++------ 1 file changed, 9 insertions(+), 6 deletions(-) diff --git a/time_calculation.py b/time_calculation.py index 5ff5b6e9..03c84d55 100644 --- a/time_calculation.py +++ b/time_calculation.py @@ -834,12 +834,15 @@ def getGEMMTime(self, dim1, dim2, dim3, name="", if eff_sm > 0: mem_access_per_sm[1] = mem_access_per_sm[1] / eff_sm - sm_util = 1 - gemm_waves = self.memLayer[1].calc_waves_per_sm(dim1, dim2, dim3, gemm.l2_M, gemm.l2_K, gemm.l2_N) - if gemm_waves != -1: - full_waves = math.floor(gemm_waves) - partial_wave = gemm_waves - full_waves - sm_util = (full_waves + partial_wave * partial_wave) / gemm_waves + if flashattn_enable: + sm_util = 0.25 # assume 25% SM utilization for FlashAttention (Dao, 2023) + else: + sm_util = 1 + gemm_waves = self.memLayer[1].calc_waves_per_sm(dim1, dim2, dim3, gemm.l2_M, gemm.l2_K, gemm.l2_N) + if gemm_waves != -1: + full_waves = math.floor(gemm_waves) + partial_wave = gemm_waves - full_waves + sm_util = (full_waves + partial_wave * partial_wave) / gemm_waves GEMM_time = self.roofline(GEMM_flop, mem_access_per_sm, name, util=sm_util, flashattn_enable=flashattn_enable) if flashattn_enable or disable_overhead: From 980dad0210a0d69c16f39236e4b95622005642b9 Mon Sep 17 00:00:00 2001 From: gkarfakis19 Date: Fri, 5 Dec 2025 21:04:08 -0800 Subject: [PATCH 74/77] fixed mapping --- astrasim_lib/executor.py | 62 ++++++++++++++++++++++++---------------- 1 file changed, 38 insertions(+), 24 deletions(-) diff --git a/astrasim_lib/executor.py b/astrasim_lib/executor.py index 5bed77bc..dc2d5fd2 100644 --- a/astrasim_lib/executor.py +++ b/astrasim_lib/executor.py @@ -1303,29 +1303,6 @@ def analyze_for_collective(edge: Any) -> Tuple[Set[Any], Set[Any], Set[Any]]: pipeline_recv_cache: Dict[Tuple[Any, Any, int], int] = {} tag_counter = itertools.count(start=1) - # Build TP communicator groups (ids start at 1000) per label and dp - global _LAST_TP_GROUPS, _TP_LABEL_DP_TO_ID, _TP_MEMBERS_TO_GID - _LAST_TP_GROUPS = {} - _TP_LABEL_DP_TO_ID = {} - _TP_MEMBERS_TO_GID = {} - tp_gid_counter = _it_count(start=_TP_GROUP_BASE_ID) - - # Collect per-label TP members per dp index (use prebuilt mapping to avoid rescans) - # Step 9: Convert axis membership tokens into communicator IDs (reused across - # labels/dp replicas) so later we can write ``comm_groups.json``. - for label in sorted(label_group_tokens.keys()): - tokens = sorted(label_group_tokens[label]) - for axis, dp_idx, members_tuple in tokens: - members_list = list(members_tuple) - members_key = tuple(members_list) - gid = _TP_MEMBERS_TO_GID.get(members_key) - if gid is None: - gid = str(next(tp_gid_counter)) - _TP_MEMBERS_TO_GID[members_key] = gid - if gid not in _LAST_TP_GROUPS: - _LAST_TP_GROUPS[gid] = members_list - _TP_LABEL_DP_TO_ID[(label, dp_idx)] = gid - # Step 8 helper: materialise SEND/RECV control edges across stages whenever a # pipeline dependency exists between ``parent`` and ``target``. def ensure_pipeline(parent: Any, target: Any, dp_idx: int) -> int: @@ -1409,7 +1386,6 @@ def ensure_pipeline(parent: Any, target: Any, dp_idx: int) -> int: pipeline_recv_cache[key] = recv_id return recv_id - mapping_result = gmap.finalize_collection(collector) if mapping_result: stage_ids, stage_index = _remap_stages_for_mapping( @@ -1424,6 +1400,44 @@ def ensure_pipeline(parent: Any, target: Any, dp_idx: int) -> int: num_stages=num_stages, output_dir=output_dir, ) + # Rebuild communicator layout after remapping so pg_name and comm_groups.json + # reflect the SCOTCH permutation. + axis_groups = _build_axis_groups(axis_order, axis_sizes, stage_axis_coords, stage_to_ranks) + ( + tp_collective_labels, + label_to_edges, + label_group_tokens, + ) = _assign_collective_labels( + tp_collective_groups, + collective_info, + axis_order, + axis_sizes, + stage_axis_coords, + axis_groups, + stage_to_ranks, + dp_count, + ) + + # Build TP communicator groups (ids start at 1000) per label and dp using the + # (potentially remapped) label tokens. + global _LAST_TP_GROUPS, _TP_LABEL_DP_TO_ID, _TP_MEMBERS_TO_GID + _LAST_TP_GROUPS = {} + _TP_LABEL_DP_TO_ID = {} + _TP_MEMBERS_TO_GID = {} + tp_gid_counter = _it_count(start=_TP_GROUP_BASE_ID) + + for label in sorted(label_group_tokens.keys()): + tokens = sorted(label_group_tokens[label]) + for axis, dp_idx, members_tuple in tokens: + members_list = list(members_tuple) + members_key = tuple(members_list) + gid = _TP_MEMBERS_TO_GID.get(members_key) + if gid is None: + gid = str(next(tp_gid_counter)) + _TP_MEMBERS_TO_GID[members_key] = gid + if gid not in _LAST_TP_GROUPS: + _LAST_TP_GROUPS[gid] = members_list + _TP_LABEL_DP_TO_ID[(label, dp_idx)] = gid # Step 10: Emit ET nodes per stage (collectives first, then compute nodes), # iterating in the topological order computed above. From d4977ad53ffb622f9f4b2422d84a55d48146c07b Mon Sep 17 00:00:00 2001 From: gkarfakis19 Date: Mon, 15 Dec 2025 04:20:00 -0800 Subject: [PATCH 75/77] gitignore up --- .gitignore | 5 ++++- MLSynth | 1 + astra-sim | 2 +- symbolic_tensor_graph | 1 + 4 files changed, 7 insertions(+), 2 deletions(-) create mode 160000 MLSynth create mode 160000 symbolic_tensor_graph diff --git a/.gitignore b/.gitignore index 19c70b0c..fd069542 100644 --- a/.gitignore +++ b/.gitignore @@ -13,4 +13,7 @@ tools/fault_sweep_results* tools/*.png tools/*.tsv tools/*.json -tools/gmap_artifacts/ \ No newline at end of file +tools/gmap_artifacts/ +*.md +*.et +tools/comp/wrapper_outputs/* \ No newline at end of file diff --git a/MLSynth b/MLSynth new file mode 160000 index 00000000..665fce22 --- /dev/null +++ b/MLSynth @@ -0,0 +1 @@ +Subproject commit 665fce22c569d91ed8302365f4ca8e3fe56e2bb2 diff --git a/astra-sim b/astra-sim index 32bec13d..bded4d08 160000 --- a/astra-sim +++ b/astra-sim @@ -1 +1 @@ -Subproject commit 32bec13d1fee017c0184edba43454c1542fb1190 +Subproject commit bded4d08b27c5403ada56efcc47d07f991c93c06 diff --git a/symbolic_tensor_graph b/symbolic_tensor_graph new file mode 160000 index 00000000..3a7f49c0 --- /dev/null +++ b/symbolic_tensor_graph @@ -0,0 +1 @@ +Subproject commit 3a7f49c02ac3eb669e4eee597ef1a32f7db44d07 From 0acc83c7d1ba36688ff1d6da23443f4055e22aa3 Mon Sep 17 00:00:00 2001 From: gkarfakis19 Date: Mon, 15 Dec 2025 07:17:37 -0800 Subject: [PATCH 76/77] changed configs for 2d topo --- LLM_execution.py | 5 + astra-sim | 2 +- astrasim_lib/config_generation.py | 58 ++++++++--- astrasim_lib/gmap.py | 44 ++++++++- config.py | 157 ++++++++++++++++++++++++++---- 5 files changed, 233 insertions(+), 33 deletions(-) diff --git a/LLM_execution.py b/LLM_execution.py index 43e95150..65933d11 100644 --- a/LLM_execution.py +++ b/LLM_execution.py @@ -984,12 +984,17 @@ def _build_rank_layout_descriptor(self) -> Dict[str, Any]: size_value = getattr(dim, "size", None) if size_value is None: raise ValueError("optimize_2dmap requires an explicit dimension size.") + dims_value = getattr(dim, "size_2d", None) + if dims_value is not None: + dims_value = (int(dims_value[0]), int(dims_value[1])) optimize_cfg = { "dimension_index": idx, "topology": str(topo_type), "size": int(size_value), "parallelisms": tuple(getattr(dim, "parallelisms", ()) or ()), } + if dims_value: + optimize_cfg["dims"] = dims_value self._first_dim_optimize_cfg = optimize_cfg def _safe_int(value: Any, default: int = 1) -> int: diff --git a/astra-sim b/astra-sim index bded4d08..32bec13d 160000 --- a/astra-sim +++ b/astra-sim @@ -1 +1 @@ -Subproject commit bded4d08b27c5403ada56efcc47d07f991c93c06 +Subproject commit 32bec13d1fee017c0184edba43454c1542fb1190 diff --git a/astrasim_lib/config_generation.py b/astrasim_lib/config_generation.py index ce4e85b6..9e5ac393 100644 --- a/astrasim_lib/config_generation.py +++ b/astrasim_lib/config_generation.py @@ -107,23 +107,39 @@ def _expand_network_entries(entries: Sequence[Dict[str, Any]]) -> List[Dict[str, for entry in entries: topo = entry.get("topology") npus_value = entry.get("npus") + dims_value = entry.get("dims") bw_value = entry.get("bandwidth") lat_value = entry.get("latency") if isinstance(topo, str) and (topo.startswith("Torus") or topo.startswith("Mesh")): - dim_count = int(topo[-2]) if len(topo) >= 2 and topo[-2].isdigit() else 2 + dim_count: int + if dims_value is not None and isinstance(dims_value, (list, tuple)): + dim_count = len(dims_value) + elif isinstance(npus_value, (list, tuple)): + dim_count = len(npus_value) + elif len(topo) >= 2 and topo[-2].isdigit(): + dim_count = int(topo[-2]) + else: + dim_count = 1 base_name = "Ring" if topo.startswith("Torus") else "Mesh" + if dims_value is None and isinstance(npus_value, (list, tuple)) and len(npus_value) == dim_count: + dims_value = tuple(npus_value) for dim_idx in range(dim_count): - curr_npu = ( - npus_value[dim_idx] - if isinstance(npus_value, (list, tuple)) - else npus_value - ) - if dim_count == 2: - root = int(math.isqrt(int(curr_npu))) - if root * root != int(curr_npu): - raise ValueError(f"npus ({curr_npu}) must be a perfect square for 2D topology {topo}.") - curr_npu = root + if dims_value is not None: + if not isinstance(dims_value, (list, tuple)) or len(dims_value) != dim_count: + raise ValueError(f"dims for topology {topo} must match dimension count {dim_count}.") + curr_npu = dims_value[dim_idx] + else: + curr_npu = ( + npus_value[dim_idx] + if isinstance(npus_value, (list, tuple)) + else npus_value + ) + if dim_count == 2: + root = int(math.isqrt(int(curr_npu))) + if root * root != int(curr_npu): + raise ValueError(f"npus ({curr_npu}) must be a perfect square for 2D topology {topo}.") + curr_npu = root topology_name = "Mesh" if dim_count == 2 and curr_npu <= 2 else base_name expanded.append( { @@ -340,7 +356,21 @@ def generate_astrasim_configs_from_hw( for dim, axes_selected, product_size, _ in selected_dims: topo = _normalize_topology_name(dim.topology_type) - size = product_size + full_size = int(getattr(dim, "size", product_size) or product_size) + dims_value = getattr(dim, "size_2d", None) + dims_tuple: Optional[Tuple[int, int]] = None + use_shape = dims_value is not None and product_size == full_size + if use_shape: + dims_tuple = (int(dims_value[0]), int(dims_value[1])) # type: ignore[index] + shape_product = dims_tuple[0] * dims_tuple[1] + if shape_product != product_size: + raise ValueError( + f"Network dimension '{dim.label}' shape {dims_tuple} product {shape_product} " + f"does not match parallelism product {product_size}." + ) + size = shape_product + else: + size = product_size if transform_2d_to_1d and topo in ["Torus2D", "Mesh2D", "KingMesh2D"]: if topo == "Torus2D": topo = "Ring" @@ -348,6 +378,8 @@ def generate_astrasim_configs_from_hw( topo = "Mesh" if topo == "KingMesh2D": topo = "HyperCube" + # When flattening 2D, prefer the explicit shape product if available. + size = size if topo == "Ring" and size <= 2: topo = "FullyConnected" effective_bw = float(getattr(dim, "effective_bandwidth", dim.bandwidth)) @@ -360,6 +392,8 @@ def generate_astrasim_configs_from_hw( "bandwidth": round(_gbps_from_bps(effective_bw), 6), "latency": round(_ns_from_s(latency_s), 3), } + if dims_tuple is not None and topo in ["Torus2D", "Mesh2D", "KingMesh2D"]: + entry["dims"] = dims_tuple network_entries.append(entry) network_entries = _expand_network_entries(network_entries) diff --git a/astrasim_lib/gmap.py b/astrasim_lib/gmap.py index 1e97d494..b7d0470b 100644 --- a/astrasim_lib/gmap.py +++ b/astrasim_lib/gmap.py @@ -42,6 +42,7 @@ def __init__( output_dir: str, vertex_count: int, include_pipeline: bool, + dims: Optional[Tuple[int, int]] = None, ) -> None: if not subset_axes: raise ValueError("optimize_2dmap requires at least one axis in the first network dimension") @@ -54,6 +55,7 @@ def __init__( self.output_dir = output_dir self.vertex_count = int(vertex_count) self.include_pipeline = include_pipeline + self.dims: Optional[Tuple[int, int]] = tuple(int(v) for v in dims) if dims else None self._edge_weights: Dict[Tuple[int, int], float] = {} missing = [axis for axis in self.subset_axes if axis not in self.axis_sizes] @@ -312,12 +314,42 @@ def begin_collection( if str(axis).strip() ) subset_strides = build_subset_strides(subset_axes, axis_sizes) - vertex_count = 1 + dims_raw = optimize_cfg.get("dims") + if dims_raw is None and isinstance(optimize_cfg.get("size"), (tuple, list)): + dims_raw = optimize_cfg.get("size") + dims: Optional[Tuple[int, int]] = None + if dims_raw is not None: + if not isinstance(dims_raw, (tuple, list)) or len(dims_raw) != 2: + raise ValueError("optimize_2dmap dims must be a two-item tuple/list when provided.") + try: + dims = (int(dims_raw[0]), int(dims_raw[1])) + except (TypeError, ValueError) as exc: + raise ValueError("optimize_2dmap dims entries must be integers.") from exc + if dims[0] <= 0 or dims[1] <= 0: + raise ValueError("optimize_2dmap dims entries must be > 0.") + + subset_product = 1 for axis in subset_axes: if axis not in axis_sizes: raise ValueError(f"Axis '{axis}' from optimize_2dmap is missing in axis_sizes.") - vertex_count *= int(axis_sizes[axis]) - declared_size = int(optimize_cfg.get("size", vertex_count)) + subset_product *= int(axis_sizes[axis]) + + if dims: + vertex_count = int(dims[0]) * int(dims[1]) + if vertex_count != subset_product: + raise ValueError( + f"optimize_2dmap dims product {vertex_count} does not match product of subset axes {subset_product}." + ) + else: + vertex_count = subset_product + declared_size_raw = optimize_cfg.get("size", vertex_count) + if isinstance(declared_size_raw, (tuple, list)): + try: + declared_size = int(declared_size_raw[0]) * int(declared_size_raw[1]) + except Exception: + declared_size = vertex_count + else: + declared_size = int(declared_size_raw) if declared_size != vertex_count: raise ValueError( f"optimize_2dmap size {declared_size} does not match product of subset axes {vertex_count}." @@ -335,6 +367,7 @@ def begin_collection( output_dir=output_dir, vertex_count=vertex_count, include_pipeline=include_pipeline, + dims=dims, ) @@ -387,6 +420,11 @@ def finalize_collection(collector: Optional[GMapCollector]) -> Optional[MappingR def _infer_target_dims(collector: GMapCollector) -> Tuple[int, int]: + if collector.dims: + dim_x, dim_y = collector.dims + if dim_x <= 0 or dim_y <= 0: + raise ValueError("optimize_2dmap requires positive dims entries.") + return dim_x, dim_y total = int(collector.vertex_count) if total <= 0: raise ValueError("optimize_2dmap requires a positive vertex count.") diff --git a/config.py b/config.py index 4f43910a..27c64aa2 100644 --- a/config.py +++ b/config.py @@ -351,6 +351,7 @@ class NetworkDimensionLayout: bandwidth: float util: float latency: float + size_2d: Optional[Tuple[int, int]] = None collective_override: Dict[str, str] = field(default_factory=dict) parallelisms: Tuple[str, ...] = field(default_factory=tuple) energy_per_bit: float = 0.0 @@ -374,17 +375,34 @@ def from_raw( if "size" not in raw: raise ValueError(f"network dimension '{label}' is missing required field 'size'") size_raw = raw["size"] - size_is_auto = False - if isinstance(size_raw, str) and size_raw.strip().lower() == "auto": - size_is_auto = True - size = None + size_mode: str + size_value: Optional[int] = None + size_2d: Optional[Tuple[int, int]] = None + tuple_entries: Optional[Tuple[object, object]] = None + if isinstance(size_raw, str): + normalized_size = size_raw.strip() + if normalized_size.startswith("(") and normalized_size.endswith(")"): + tuple_entries = tuple( + part.strip() for part in normalized_size[1:-1].split(",") + ) # type: ignore[assignment] + size_mode = "tuple" + elif normalized_size.lower() == "auto": + size_mode = "auto_scalar" + else: + try: + size_value = int(size_raw) + except (TypeError, ValueError) as exc: + raise ValueError(f"network dimension '{label}' size must be an integer or tuple") from exc + if size_value < 1: + raise ValueError(f"network dimension '{label}' size must be >= 1") + size_mode = "scalar" + elif isinstance(size_raw, (list, tuple)): + if len(size_raw) != 2: + raise ValueError(f"network dimension '{label}' 2D size must have exactly two entries") + tuple_entries = tuple(size_raw) # type: ignore[assignment] + size_mode = "tuple" else: - try: - size = int(size_raw) - except (TypeError, ValueError) as exc: - raise ValueError(f"network dimension '{label}' size must be an integer") from exc - if size < 1: - raise ValueError(f"network dimension '{label}' size must be >= 1") + size_mode = "auto_scalar" topo_dict = raw.get("topology") if not isinstance(topo_dict, dict): @@ -469,21 +487,125 @@ def from_raw( normalized_parallelisms.append(normalized) alias_map[normalized] = name - computed_product = None - if size_is_auto: - computed_product = _compute_dimension_parallelism_product( + computed_product: Optional[int] = None + auto_product: Optional[int] = None + + if size_mode in {"auto_scalar", "tuple"}: + auto_product = _compute_dimension_parallelism_product( dimension_label=label, normalized_names=tuple(normalized_parallelisms), alias_map=alias_map, parallelism_params=parallelism_params, ) - size = computed_product - if size < 1: + if auto_product < 1: raise ValueError(f"network dimension '{label}' inferred size must be >= 1") + if size_mode == "auto_scalar": + size_value = auto_product + computed_product = auto_product + elif size_mode == "tuple": + assert tuple_entries is not None + resolved: List[Optional[int]] = [None, None] + auto_entries = 0 + for idx, entry in enumerate(tuple_entries): + if isinstance(entry, str): + normalized = entry.strip().lower() + if normalized == "auto": + auto_entries += 1 + if auto_entries > 1: + raise ValueError( + f"network dimension '{label}' 2D size may include at most one 'auto' entry" + ) + resolved[idx] = None + continue + if normalized not in parallelism_params: + alias = alias_map.get(normalized, normalized) + raise ValueError( + f"network dimension '{label}' 2D size references parallelism '{alias}' " + "which is not defined in parallelism" + ) + try: + factor = int(parallelism_params[normalized]) + except (TypeError, ValueError) as exc: + alias = alias_map.get(normalized, normalized) + raise ValueError( + f"network dimension '{label}' 2D size parallelism '{alias}' must be an integer" + ) from exc + if factor < 1: + alias = alias_map.get(normalized, normalized) + raise ValueError( + f"network dimension '{label}' 2D size parallelism '{alias}' must be >= 1" + ) + resolved[idx] = factor + else: + try: + factor = int(entry) + except (TypeError, ValueError) as exc: + raise ValueError( + f"network dimension '{label}' 2D size entries must be integers, parallelism names, or 'auto'" + ) from exc + if factor < 1: + raise ValueError(f"network dimension '{label}' 2D size entries must be >= 1") + resolved[idx] = factor + + known_product = 1 + for val in resolved: + if val: + known_product *= val + + if auto_entries: + if auto_product is None: + auto_product = _compute_dimension_parallelism_product( + dimension_label=label, + normalized_names=tuple(normalized_parallelisms), + alias_map=alias_map, + parallelism_params=parallelism_params, + ) + if known_product == 0: + raise ValueError(f"network dimension '{label}' 2D size known entries must be > 0") + if auto_product % known_product != 0: + raise ValueError( + f"network dimension '{label}' 2D size mismatch: auto product {auto_product} " + f"is not divisible by provided factors {tuple_entries}" + ) + auto_value = auto_product // known_product + if auto_value < 1: + raise ValueError( + f"network dimension '{label}' 2D size auto-resolved entry must be >= 1 (got {auto_value})" + ) + for idx in range(2): + if resolved[idx] is None: + resolved[idx] = auto_value + break + else: + if auto_product is None: + auto_product = _compute_dimension_parallelism_product( + dimension_label=label, + normalized_names=tuple(normalized_parallelisms), + alias_map=alias_map, + parallelism_params=parallelism_params, + ) + if known_product != auto_product: + raise ValueError( + f"network dimension '{label}' 2D size mismatch: provided shape {tuple_entries} " + f"product {known_product} does not match parallelism product {auto_product}" + ) + + if resolved[0] is None or resolved[1] is None: + raise ValueError(f"network dimension '{label}' 2D size could not be fully resolved") + + size_2d = (int(resolved[0]), int(resolved[1])) + size_value = int(size_2d[0]) * int(size_2d[1]) + computed_product = auto_product + else: + computed_product = None + + if size_value is None: + raise ValueError(f"network dimension '{label}' size could not be resolved") + _validate_dimension_parallelisms( dimension_label=label, - dimension_size=size, + dimension_size=int(size_value) if size_value is not None else 0, normalized_names=tuple(normalized_parallelisms), alias_map=alias_map, parallelism_params=parallelism_params, @@ -493,7 +615,8 @@ def from_raw( return cls( id=dim_id, label=label, - size=size, + size=int(size_value) if size_value is not None else 0, + size_2d=size_2d, topology_type=topo_type, bandwidth=bandwidth, util=util, From 7fcca220c67b93e871e7a5412179cc58dc574a7b Mon Sep 17 00:00:00 2001 From: gkarfakis19 Date: Mon, 15 Dec 2025 11:45:30 -0800 Subject: [PATCH 77/77] fix --- MLSynth | 1 - symbolic_tensor_graph | 1 - 2 files changed, 2 deletions(-) delete mode 160000 MLSynth delete mode 160000 symbolic_tensor_graph diff --git a/MLSynth b/MLSynth deleted file mode 160000 index 665fce22..00000000 --- a/MLSynth +++ /dev/null @@ -1 +0,0 @@ -Subproject commit 665fce22c569d91ed8302365f4ca8e3fe56e2bb2 diff --git a/symbolic_tensor_graph b/symbolic_tensor_graph deleted file mode 160000 index 3a7f49c0..00000000 --- a/symbolic_tensor_graph +++ /dev/null @@ -1 +0,0 @@ -Subproject commit 3a7f49c02ac3eb669e4eee597ef1a32f7db44d07