diff --git a/deepinfer/model_executor/models/deepseek.py b/deepinfer/model_executor/models/deepseek.py index 1946b18..040bf46 100644 --- a/deepinfer/model_executor/models/deepseek.py +++ b/deepinfer/model_executor/models/deepseek.py @@ -185,7 +185,11 @@ class DeepseekModel(nn.Module): getattr(input_metadata, "images_seq_mask", None), ) + from deepinfer.worker.model_runner import _rank_timeline + + _rank_timeline("engram0") ngram_embed_hooks = self.engram_runtime.fetch(hashes) + _rank_timeline("engram1") for i in range(self.num_layers): if input_metadata.is_prompt and i == self.yoco_boundary_layer_idx: diff --git a/deepinfer/neuron_ops/vendor/deep_ep/buffers/engram.py b/deepinfer/neuron_ops/vendor/deep_ep/buffers/engram.py index 1555921..cad3083 100644 --- a/deepinfer/neuron_ops/vendor/deep_ep/buffers/engram.py +++ b/deepinfer/neuron_ops/vendor/deep_ep/buffers/engram.py @@ -53,6 +53,33 @@ def _accept_legacy_storage_size(init): +def _prefault(tables: List[torch.Tensor]) -> None: + """Populate this process's page tables for the shared tables (MADV_POPULATE_READ, Linux >= 5.14). + + Every step each rank gathers a few hundred random rows of tables of ~100 GB; mapped lazily, each row is a page + fault on the same shmem file in 64 processes at once: the gather took 15-22 ms a layer (p90 25-38 ms, outliers + 150 ms), and the slowest rank paced every collective of the step (issue #29). Populated, it is ~0.1 ms.""" + import ctypes + import time + + libc = ctypes.CDLL(None, use_errno=True) + libc.madvise.argtypes = (ctypes.c_void_p, ctypes.c_size_t, ctypes.c_int) + madv_populate_read = 22 + page = os.sysconf("SC_PAGE_SIZE") + started = time.perf_counter() + for t in tables: + start = t.data_ptr() // page * page + length = t.data_ptr() + t.numel() * t.element_size() - start + if libc.madvise(start, length, madv_populate_read) != 0: + import logging + + logging.getLogger(__name__).warning("Engram prefault: madvise failed (errno %d)", ctypes.get_errno()) + return + import logging + + logging.getLogger(__name__).info("Engram prefault: %d tables in %.1f s", len(tables), time.perf_counter() - started) + + def _take_rows_to(table: torch.Tensor, index: torch.Tensor, device: torch.device) -> torch.Tensor: """Gather rows where the table lives (Engram tables default to host memory, which the GPU reads through UVA), then move the result to the compute device.""" @@ -241,19 +268,28 @@ class EngramBuffer(BufferBase): self._sf_shards.append((sf_table, sf.dtype)) self._shared_tables = True comm.barrier(self) + if os.environ.get("DEEPINFER_ENGRAM_PREFAULT", "1") == "1": + _prefault([*self._storages, *(t for t, _ in self._sf_shards)]) def _fetch_layer_shared(self, layer_idx: int, indices: torch.Tensor) -> torch.Tensor | tuple: + from deepinfer.worker.model_runner import _rank_timeline + device = indices.device num_tokens, num_entries_per_token = indices.shape flat = indices.reshape(-1).to(device="cpu", dtype=torch.int64) + _rank_timeline(f"eg_d2h{layer_idx}") flat = flat.clamp(0, self._storages[layer_idx].size(0) - 1) rows = self._storages[layer_idx].index_select(0, flat).view(self.dtype) + _rank_timeline(f"eg_gather{layer_idx}") fetched = rows.reshape(num_tokens, num_entries_per_token * self.hidden).to(device) + _rank_timeline(f"eg_h2d{layer_idx}") if not self._sf_shards: return fetched sf_table, sf_dtype = self._sf_shards[layer_idx] sf_rows = sf_table.index_select(0, flat).view(sf_dtype) - return fetched, sf_rows.reshape(num_tokens, num_entries_per_token * self.num_sf_packs).to(device) + out = fetched, sf_rows.reshape(num_tokens, num_entries_per_token * self.num_sf_packs).to(device) + _rank_timeline(f"eg_sf{layer_idx}") + return out # --- Fetch ---------------------------------------------------------------------- diff --git a/deepinfer/platform.py b/deepinfer/platform.py index 6bfa363..6d32273 100644 --- a/deepinfer/platform.py +++ b/deepinfer/platform.py @@ -250,6 +250,10 @@ class PlatformNeuron(Platform): os.environ["CUDA_VISIBLE_DEVICES"] = cores local_rank = int(os.getenv("DEEPINFER_NEURON_LOCAL_RANK", "0")) os.environ["NEURON_RT_VISIBLE_CORES"] = cores.split(",")[local_rank] + # The runtime's low-latency completion busy-polls the notification queues (~0.7 of a CPU per process); + # with 64 ranks on one host the pollers crowd the cores at every step boundary, the main threads get + # preempted, and the slowest rank paces every collective. Off: decode step 580 -> 522 ms (issue #29). + os.environ.setdefault("NEURON_RT_XU_WORKER_ENABLE_LOW_LATENCY_COMPLETION", "0") # torch-neuronx numbers a process's torch device by LOCAL_RANK out of LOCAL_WORLD_SIZE when the first # "neuron" process group is created (distributed/backend.py, else it derives them from the hosts of all # ranks). With one core per process every rank must be device 0 of a 1-device world: `cuda:N` maps to diff --git a/deepinfer/worker/model_runner.py b/deepinfer/worker/model_runner.py index e0ea990..9347850 100644 --- a/deepinfer/worker/model_runner.py +++ b/deepinfer/worker/model_runner.py @@ -54,6 +54,25 @@ class MainModelOutput: graph_size: int = 0 + +_RANK_TL: list = [] + + +def _rank_timeline(phase: str) -> None: + """DEEPINFER_NEURON_RANK_TIMELINE=1: wall-clock (time.time, comparable across the host's processes) of a step's + host phases on every rank, appended to /tmp/rank_tl_.txt every 200 entries (startup stagger, issue #29).""" + if os.environ.get("DEEPINFER_NEURON_RANK_TIMELINE") != "1": + return + import resource + + ru = resource.getrusage(resource.RUSAGE_THREAD) + _RANK_TL.append(f"{phase} {time.time():.6f} {time.thread_time():.6f} {ru.ru_nvcsw} {ru.ru_nivcsw} {ru.ru_minflt} {ru.ru_majflt}") + if len(_RANK_TL) >= 200: + rank = dist.get_rank() if dist.is_initialized() else 0 + with open(f"/tmp/rank_tl_{rank}.txt", "a") as f: + f.write("\n".join(_RANK_TL) + "\n") + _RANK_TL.clear() + class ModelRunner: def __init__( self, @@ -387,7 +406,9 @@ class ModelRunner: IMPORTANT: This function must NOT contain any cuda synchronize. """ + _rank_timeline("enter") input_metadata, total_q, run_cuda_graph, cuda_graph_used, graph_size = self._prepare_input(seq_metadata_list, metadata_zos, total_q) + _rank_timeline("prepared") with ( cpu_timer.record("launch_forward", record_name=f"launch_forward_N{input_metadata.num_reqs}_B{total_q}"), @@ -398,6 +419,7 @@ class ModelRunner: model_output = self.cuda_graph_runner.replay(cuda_graph_used) else: model_output = self.model(kv_caches=kv_caches, input_metadata=input_metadata) + _rank_timeline("dispatched") return MainModelOutput( model_output=model_output, diff --git a/deepinfer/worker/worker.py b/deepinfer/worker/worker.py index fdf47ad..a3fa772 100644 --- a/deepinfer/worker/worker.py +++ b/deepinfer/worker/worker.py @@ -186,6 +186,9 @@ class Worker(WorkerBase): else: self.prev_step_state = curr_step_state + from deepinfer.worker.model_runner import _rank_timeline + + _rank_timeline("step_end") return request_outputs def _correct_seqlens_qk(