Diff Coverage

            device_type == "cpu"
            and platform.platform_type == PlatformType.PYTORCH
            and str(platform.get_backend()) == "fake"
    ):
        return platform.get_dry_run_device("cpu", 0)
    if device_type not in ("npu", "cuda"):
        raise AssertionError(
            f"hyper_parallel.fully_shard support device in [torch.npu, torch.cuda], "
            f"but got '{device_type}'"

        Raises:
            NotImplementedError: When the selected backend has no fake process
                group implementation.
        """
        raise NotImplementedError(
            "Dry-run process groups are not supported by this platform"
        )

    @staticmethod

        Raises:
            NotImplementedError: When the platform has no fake device support.
        """
        raise NotImplementedError(
            "Dry-run fake devices are not supported by this platform"
        )

    @staticmethod

        rank = int(os.environ["RANK"])
        world_size = int(os.environ["WORLD_SIZE"])
        local_rank = int(os.environ["LOCAL_RANK"])
        if world_size < 1:
            raise ValueError(f"WORLD_SIZE must be >= 1, got {world_size}")
        if rank < 0 or rank >= world_size:
            raise ValueError(
                f"RANK must be in [0, {world_size}), got {rank}"
            )
        if local_rank < 0:
            raise ValueError(f"LOCAL_RANK must be >= 0, got {local_rank}")
        return cls(rank=rank, world_size=world_size, local_rank=local_rank)


def _report_limitations(metadata: Optional[Dict[str, Any]] = None) -> list[str]:

def _category_name(category: Any) -> str:
    """Return a stable JSON key for a MemTracker category."""
    if isinstance(category, str):
        return category
    value = getattr(category, "value", None)
    return str(value if value is not None else category)


def _normalize_snapshot(snapshot: Dict[Any, Dict[Any, int]]) -> Dict[str, Dict[str, int]]:
    """Convert device/category objects in a MemTracker snapshot to JSON keys."""

    """Aggregate category bytes for devices matching ``device_type``."""
    breakdown: Dict[str, int] = {}
    for device, categories in snapshot.items():
        if device.split(":", maxsplit=1)[0] != device_type:
            continue
        for category, value in categories.items():
            if category == "Total":
                continue
            breakdown[category] = breakdown.get(category, 0) + int(value)

def _serialize_module(module_stats: Any, device_type: str) -> Dict[str, Any]:
    """Serialize one MemTracker module-stat object without private enum types."""
    snapshots = {}
    for state, snapshot_list in getattr(module_stats, "snapshots", {}).items():
        snapshots[_category_name(state)] = [
            _normalize_snapshot(snapshot)
            for snapshot in snapshot_list
        ]
    return {

    """Capture and intern operator allocation stacks for one trace mode."""

    def __init__(self) -> None:
        """Initialize an empty cache keyed by complete filtered frame tuples."""
        self._cache: Dict[tuple[tuple[str, int, str], ...], Dict[str, Any]] = {}

    @staticmethod
    def _is_internal_frame(filename: str, function_name: str) -> bool:
        """Return whether a frame was excluded by the previous stack collector."""
        normalized_path = filename.replace("\\", "/")
        if normalized_path == __file__ and function_name in (
                "__torch_dispatch__",
                "capture",
                "_record_storage_resize",
                "resize_",

                "capture",
                "_record_storage_resize",
                "resize_",
        ):
            return True
        return normalized_path.endswith(("/torch/_compile.py", "/torch/_dynamo/eval_frame.py"))

    def capture(self) -> Dict[str, Any]:
        """Return cached stack text and leaf location for the current caller."""
        frame = sys._getframe()  # pylint: disable=W0212
        frames = []
        while frame is not None and len(frames) < 64:
            code = frame.f_code
            frames.append((code.co_filename, frame.f_lineno, code.co_name))
            frame = frame.f_back
        frames.reverse()
        stack_key = tuple(
            frame_info
            for frame_info in frames
            if not self._is_internal_frame(frame_info[0], frame_info[2])
        )
        cached = self._cache.get(stack_key)
        if cached is not None:
            return cached
        python_stack = "|".join(
            f"File:{filename};Line:{line_num};Function:{function_name}"
            for filename, line_num, function_name in stack_key
        )
        leaf_frame = stack_key[-1] if stack_key else ("", 0, "")
        captured = {
            "python_stack": python_stack,
            "file_name": leaf_frame[0],
            "line_num": leaf_frame[1],
        }
        self._cache[stack_key] = captured
        return captured


def _create_indexed_mem_tracker(mem_tracker_type: Any) -> Any:
    """Create a MemTracker whose module peak lookup follows active FQNs."""

            self._next_module_registration_order = 0

        def _remove_from_fqn(self, module_fqn: str, module: Any) -> None:
            """Remove a module from one FQN bucket and prune an empty bucket."""
            bucket = self._module_stats_by_fqn.get(module_fqn)
            if bucket is None:
                return
            bucket.pop(module, None)
            if not bucket:
                self._module_stats_by_fqn.pop(module_fqn, None)

        def _rebind_module_stats(self, module: Any) -> None:
            """Bind a module's latest stats FQN while preserving first-seen order."""
            module_stats = self.memory_tracking.get(module)
            if module_stats is None:
                return
            new_fqn = str(module_stats.mod_fqn)
            old_fqn = self._module_fqn_by_module.get(module)
            if old_fqn == new_fqn:
                return
            if old_fqn is not None:
                self._remove_from_fqn(old_fqn, module)
            if module not in self._module_registration_order:
                self._module_registration_order[module] = self._next_module_registration_order
                self._next_module_registration_order += 1
            self._module_stats_by_fqn.setdefault(new_fqn, {})[module] = module_stats
            self._module_fqn_by_module[module] = new_fqn

        def _pre_fw_hook(self, module: Any, inputs: Any) -> None:
            """Run the original hook, then index its final module FQN."""
            super()._pre_fw_hook(module, inputs)
            self._rebind_module_stats(module)

        def reset_mod_stats(self) -> None:
            """Clear module statistics and all indexes derived from them."""
            super().reset_mod_stats()
            self._module_stats_by_fqn.clear()
            self._module_fqn_by_module.clear()
            self._module_registration_order.clear()
            self._next_module_registration_order = 0

        def _update_peak_stats(self, peak_state: Any) -> None:
            """Update active module peaks in registration order and global peak."""
            active_modules = {}
            for module_fqn in self._mod_tracker.parents:
                for module, module_stats in self._module_stats_by_fqn.get(str(module_fqn), {}).items():
                    active_modules[module] = module_stats
            ordered_modules = sorted(
                active_modules.items(),
                key=lambda item: self._module_registration_order[item[0]],
            )
            current_snapshot = self._curr_mem_snap
            for _, module_stats in ordered_modules:
                if peak_state not in module_stats.snapshots:
                    continue
                for device, device_snapshot in current_snapshot.items():
                    if module_stats.local_peak.get(device, 0) < device_snapshot["Total"]:
                        module_stats.local_peak[device] = device_snapshot["Total"]
                        module_stats.snapshots[peak_state][-1][device] = copy.deepcopy(device_snapshot)

            for device, device_snapshot in current_snapshot.items():
                if self._peak_mem.get(device, 0) < device_snapshot["Total"]:
                    self._peak_mem[device] = device_snapshot["Total"]
                    self._peak_mem_snap[device] = copy.deepcopy(device_snapshot)

        def get_device_memory_totals(self, device_name: str) -> tuple[int, int]:
            """Read current and peak totals for one device without copying snapshots."""
            current_total = 0
            for device, device_snapshot in self._curr_mem_snap.items():
                if str(device) == device_name:
                    current_total = int(device_snapshot.get("Total", 0))
                    break
            recorded_peak_total = 0
            for device, device_snapshot in self._peak_mem_snap.items():
                if str(device) == device_name:
                    recorded_peak_total = int(device_snapshot.get("Total", 0))
                    break
            return current_total, max(current_total, recorded_peak_total)

    return _IndexedMemTracker()

def _operator_phase(tracker: Any) -> str:
    """Resolve the current training phase from MemTracker state."""
    if getattr(tracker, "_in_opt", False):
        return "optimizer"
    from hyper_parallel.core.activation_checkpoint.recompute_state import (  # pylint: disable=C0415
        is_recomputing,
    )
    if is_recomputing():
        return "recompute"
    module_tracker = getattr(tracker, "_mod_tracker", None)
    if getattr(module_tracker, "is_bw", False):
        return "backward"
    return "forward"


def _active_module_fqn(tracker: Any) -> str:
    """Return the deepest active module as optional operator context."""
    module_tracker = getattr(tracker, "_mod_tracker", None)
    parents = getattr(module_tracker, "parents", ())
    candidates = [str(parent) for parent in parents if str(parent) not in ("", "Global")]
    return max(candidates, key=lambda fqn: (fqn.count("."), len(fqn)), default="")


def _create_operator_trace_mode(tracker: Any, device_type: str) -> Any:
    """Create a lazy TorchDispatchMode that records logical memory blocks."""

def _create_operator_trace_mode(tracker: Any, device_type: str) -> Any:
    """Create a lazy TorchDispatchMode that records logical memory blocks."""
    import torch  # pylint: disable=C0415
    from torch.distributed._tools.mem_tracker import get_untyped_storages  # pylint: disable=C0415
    from torch.utils._python_dispatch import TorchDispatchMode  # pylint: disable=C0415
    from torch.utils._pytree import tree_flatten  # pylint: disable=C0415

    def _storage_map(values: Any) -> Dict[int, Dict[str, Any]]:
        """Return storage identifiers, logical devices, and byte sizes."""
        flat_values, _ = tree_flatten(values)
        storages = {}
        for value in flat_values:
            if not isinstance(value, torch.Tensor):
                continue
            try:
                tensor_storages = get_untyped_storages(value)
            except (RuntimeError, TypeError):
                continue
            for storage in tensor_storages:
                storage_key = int(getattr(storage, "_cdata", id(storage)))
                storages[storage_key] = {
                    "device": str(value.device),
                    "size": int(storage.size()),
                    "storage": storage,
                }
        return storages

    def _lookup_tracked_storage(storage: Any) -> Optional[Dict[str, Any]]:
        """Read MemTracker metadata for one storage without scanning all entries."""
        entry = tracker._WINFO.get(storage)
        if entry is None:
            return None
        winfo, storage_ref = entry
        return {
            "device": str(winfo.device),
            "size": int(winfo.mem_consumed),
            "type": _category_name(winfo.reftype),
            "storage_ref": storage_ref,

            "type": _category_name(winfo.reftype),
            "storage_ref": storage_ref,
        }

    def _tracked_output_info(
            storage_keys: set[int],
            output_storages: Dict[int, Dict[str, Any]],
    ) -> Dict[int, Dict[str, Any]]:
        """Read MemTracker metadata only for selected operator outputs."""
        tracked = {}
        for storage_key in storage_keys:
            tracked_storage = _lookup_tracked_storage(
                output_storages[storage_key]["storage"]
            )
            if tracked_storage is not None:
                tracked[storage_key] = tracked_storage
        return tracked

    class _OperatorTraceMode(TorchDispatchMode):
        """Track producer, consumer, and lifetime data for fake storages."""

        def __init__(self) -> None:
            """Initialize ordered logical task and storage state."""
            super().__init__()
            self.memory_blocks: list[Dict[str, Any]] = []
            self._active_blocks: Dict[int, Dict[str, Any]] = {}
            self._active_storage_refs: Dict[int, Any] = {}
            self._active_storage_sizes: Dict[int, int] = {}
            self._released_storage_keys: set[int] = set()
            self._next_task_index = 0
            self._next_lifecycle_event_index = 0
            self._next_virtual_address = _VIRTUAL_ADDRESS_BASE
            self._previous_storage_resize = None
            self._stack_cache = _PythonStackCache()

        def __enter__(self) -> Any:
            """Enter dispatch mode and observe MemTracker's storage resize hook."""
            super().__enter__()
            try:
                self._install_storage_resize_hook()
            except Exception:
                super().__exit__(None, None, None)
                raise
            return self

        def __exit__(self, *args: Any) -> Any:
            """Restore MemTracker's storage resize hook before leaving dispatch mode."""
            self._restore_storage_resize_hook()
            return super().__exit__(*args)

        def _install_storage_resize_hook(self) -> None:
            """Wrap the active storage resize method to record capacity lifetimes."""
            previous_resize = torch.UntypedStorage.resize_
            self._previous_storage_resize = previous_resize

            @functools.wraps(previous_resize)
            def resize_(storage: Any, size: int) -> Any:
                """Resize storage through MemTracker, then record its new lifetime."""
                old_size = int(storage.size())
                result = previous_resize(storage, size)
                new_size = int(storage.size())
                if old_size != new_size:
                    self._record_storage_resize(storage, old_size, new_size)
                return result

            torch.UntypedStorage.resize_ = resize_  # type: ignore[method-assign, assignment]

        def _restore_storage_resize_hook(self) -> None:
            """Restore the resize method that was active when this mode entered."""
            if self._previous_storage_resize is not None:
                torch.UntypedStorage.resize_ = self._previous_storage_resize  # type: ignore[method-assign, assignment]
            self._previous_storage_resize = None

        def _close_blocks(self, storage_keys: set[int], end_index: int) -> None:
            """Close active blocks and discard their lifetime bookkeeping."""
            for storage_key in sorted(storage_keys):
                block = self._active_blocks.pop(storage_key, None)
                if block is None:
                    continue
                self._active_storage_refs.pop(storage_key, None)
                self._active_storage_sizes.pop(storage_key, None)
                self._released_storage_keys.discard(storage_key)
                block["end_time_stamp"] = end_index
                block["_end_lifecycle_event_index"] = self._next_lifecycle_event_index
                self._next_lifecycle_event_index += 1

        def _close_released_blocks(self, end_index: int) -> None:
            """Close blocks whose storage weak refs disappeared."""
            released_keys = self._released_storage_keys
            self._released_storage_keys = set()
            self._close_blocks(released_keys, end_index)

        def _allocate_virtual_address(self, size: int) -> str:
            """Return a deterministic aligned address in a reserved fake range."""
            address = self._next_virtual_address
            aligned_size = max(
                _VIRTUAL_ADDRESS_ALIGNMENT,
                (
                    (size + _VIRTUAL_ADDRESS_ALIGNMENT - 1)
                    // _VIRTUAL_ADDRESS_ALIGNMENT

                    (size + _VIRTUAL_ADDRESS_ALIGNMENT - 1)
                    // _VIRTUAL_ADDRESS_ALIGNMENT
                ) * _VIRTUAL_ADDRESS_ALIGNMENT,
            )
            self._next_virtual_address += aligned_size
            return f"0x{address:x}"

        def _assign_plot_coordinates(self) -> None:
            """Assign profiler-style fractional coordinates to lifecycle events.

            Logical task indices intentionally remain unchanged. Release and
            allocation events sharing one task are ordered by their observation

            Logical task indices intentionally remain unchanged. Release and
            allocation events sharing one task are ordered by their observation
            sequence and spread inside that task's ``[task, task + 1)`` range.
            """
            events_by_task: Dict[int, list[tuple[int, Dict[str, Any], str]]] = {}
            for block in self.memory_blocks:
                start_task = int(block["start_time_stamp"])
                events_by_task.setdefault(start_task, []).append((
                    int(block.pop("_start_lifecycle_event_index")),
                    block,
                    "start_plot_time_stamp",
                ))
                end_event_index = block.pop("_end_lifecycle_event_index", None)
                if end_event_index is None:
                    continue
                end_task = int(block["end_time_stamp"])
                events_by_task.setdefault(end_task, []).append((
                    int(end_event_index),
                    block,
                    "end_plot_time_stamp",
                ))

                    block,
                    "end_plot_time_stamp",
                ))

            for task_index, task_events in events_by_task.items():
                ordered_events = sorted(task_events, key=lambda event: event[0])
                denominator = Decimal(len(ordered_events) + 1)
                for event_offset, (_, block, field) in enumerate(ordered_events, start=1):
                    coordinate = Decimal(task_index) + Decimal(event_offset) / denominator
                    block[field] = format(coordinate, ".12f")

        def finalize(self) -> None:
            """Close released blocks and mark remaining storages persistent."""
            self._close_released_blocks(self._next_task_index)
            for storage_key, block in self._active_blocks.items():
                storage_ref = self._active_storage_refs.get(storage_key)
                storage = storage_ref() if storage_ref is not None else None
                tracked_storage = (
                    _lookup_tracked_storage(storage)
                    if storage is not None
                    else None
                )
                if tracked_storage is None or tracked_storage["size"] <= 0:
                    block["end_time_stamp"] = self._next_task_index
                    block["_end_lifecycle_event_index"] = self._next_lifecycle_event_index
                    self._next_lifecycle_event_index += 1
                    continue
                block["type"] = tracked_storage["type"]
                block["end_time_stamp"] = _PERSISTENT_END_INDEX
                block["is_persistent"] = 1
            self._active_blocks.clear()
            self._active_storage_refs.clear()
            self._active_storage_sizes.clear()
            self._released_storage_keys.clear()
            self._assign_plot_coordinates()

        def _classify_output_keys(
                self,
                input_storages: Dict[int, Dict[str, Any]],
                output_storages: Dict[int, Dict[str, Any]],
                active_input_keys: set[int],

                output_storages: Dict[int, Dict[str, Any]],
                active_input_keys: set[int],
        ) -> tuple[set[int], set[int]]:
            """Return resized inputs and newly allocated target-device outputs."""
            resized_output_keys = {
                storage_key
                for storage_key in active_input_keys & set(output_storages)
                if output_storages[storage_key]["size"] != self._active_storage_sizes[storage_key]
            }
            new_output_keys = {
                storage_key
                for storage_key, storage in output_storages.items()
                if (
                    storage_key not in input_storages

                    and storage["device"].split(":", maxsplit=1)[0] == device_type
                    and storage["size"] > 0
                )
            }
            new_output_keys.update(
                storage_key
                for storage_key in resized_output_keys
                if (
                    output_storages[storage_key]["device"].split(":", maxsplit=1)[0] == device_type

                    output_storages[storage_key]["device"].split(":", maxsplit=1)[0] == device_type
                    and output_storages[storage_key]["size"] > 0
                )
            )
            return resized_output_keys, new_output_keys

        def _record_input_users(self, active_input_keys: set[int], task_index: int) -> None:
            """Append the task index to every active input storage's user list."""
            for storage_key in active_input_keys:
                user_tasks = self._active_blocks[storage_key]["user_tasks"]
                if not user_tasks or user_tasks[-1] != task_index:
                    user_tasks.append(task_index)

        def _track_new_outputs(
                self,
                output_keys: set[int],
                output_storages: Dict[int, Dict[str, Any]],
                tracked_storages: Dict[int, Dict[str, Any]],

                tracked_storages: Dict[int, Dict[str, Any]],
                context: Dict[str, Any],
        ) -> None:
            """Create logical memory blocks and weak lifetime refs for outputs."""
            for storage_key in sorted(output_keys):
                output_storage = output_storages[storage_key]
                tracked_storage = tracked_storages.get(storage_key, {})
                device = output_storage["device"]
                size = int(tracked_storage.get("size", output_storage["size"]))
                current_total, peak_total = context["device_totals"].get(device, (0, 0))
                block = {
                    "start_time_stamp": context["task_index"],
                    "end_time_stamp": _PERSISTENT_END_INDEX,
                    "_start_lifecycle_event_index": self._next_lifecycle_event_index,
                    "device_addr": self._allocate_virtual_address(size),

                    "python_stack": context["stack"]["python_stack"],
                    "is_persistent": 0,
                    "is_small": int(size < _SMALL_ALLOCATION_BYTES),
                }
                self._next_lifecycle_event_index += 1
                self.memory_blocks.append(block)
                self._active_blocks[storage_key] = block
                self._active_storage_sizes[storage_key] = output_storage["size"]
                storage_ref = tracked_storage.get("storage_ref")
                storage = storage_ref() if storage_ref is not None else None
                if storage is None:
                    storage = output_storage["storage"]
                self._active_storage_refs[storage_key] = weakref.ref(
                    storage,
                    lambda _, key=storage_key: self._released_storage_keys.add(key),
                )

        def _record_storage_resize(self, storage: Any, old_size: int, new_size: int) -> None:
            """Record one tracked storage capacity transition as a block lifetime."""
            storage_key = int(getattr(storage, "_cdata", id(storage)))
            tracked_storage = _lookup_tracked_storage(storage)
            if tracked_storage is None:
                return
            device = tracked_storage["device"]
            if device.split(":", maxsplit=1)[0] != device_type:
                return

            if old_size > 0 or storage_key in self._active_blocks:
                self._close_blocks({storage_key}, self._next_task_index)
            if new_size <= 0:
                return

            task_index = self._next_task_index
            self._next_task_index += 1
            output_storages = {
                storage_key: {
                    "device": device,
                    "size": new_size,
                    "storage": storage,

                    "size": new_size,
                    "storage": storage,
                },
            }
            tracked_storages = {storage_key: tracked_storage}
            self._track_new_outputs(
                {storage_key},
                output_storages,
                tracked_storages,
                {

                    "module_fqn": _active_module_fqn(tracker),
                },
            )

        def __torch_dispatch__(self, func, types, args=(), kwargs=None):
            del types
            self._close_released_blocks(self._next_task_index)
            input_storages = _storage_map((args, kwargs))
            active_input_keys = {
                storage_key
                for storage_key in input_storages
                if storage_key in self._active_blocks
            }
            result = func(*args, **(kwargs or {}))
            output_storages = _storage_map(result)
            resized_output_keys, new_output_keys = self._classify_output_keys(
                input_storages,
                output_storages,
                active_input_keys,
            )
            self._close_released_blocks(self._next_task_index)
            if not active_input_keys and not new_output_keys and not resized_output_keys:
                return result
            task_index = self._next_task_index
            self._next_task_index += 1
            self._record_input_users(active_input_keys, task_index)
            self._close_blocks(resized_output_keys, task_index)
            if not new_output_keys:
                return result
            tracked_storages = _tracked_output_info(
                new_output_keys,
                output_storages,
            )
            output_devices = {
                output_storages[storage_key]["device"]
                for storage_key in new_output_keys
            }
            device_totals = {
                device: tracker.get_device_memory_totals(device)
                for device in output_devices
            }
            self._track_new_outputs(
                new_output_keys,
                output_storages,
                tracked_storages,
                {

                    "operator_name": str(func),
                    "module_fqn": _active_module_fqn(tracker),
                },
            )
            return result

    return _OperatorTraceMode()


def build_memory_report(
        tracker: Any,

    peak = _normalize_snapshot(tracker.get_tracker_snapshot("peak"))
    current = _normalize_snapshot(tracker.get_tracker_snapshot("current"))
    peak_bytes = _snapshot_total(peak, device_type)
    if peak_bytes <= 0:
        raise ValueError(
            f"MemTracker observed no {device_type!r} tensor memory; "
            f"tracked devices are {sorted(peak)}"
        )

    Raises:
        ValueError: If ``report`` does not describe a successful run.
    """
    if report.get("status") != "ok":
        raise ValueError("CSV output requires a successful HyperDryRun report")
    rows = []
    for block in report.get("memory_blocks", []):
        user_tasks = list(block.get("user_tasks", []))
        row = {

            writer = csv.DictWriter(report_file, fieldnames=_CSV_FIELDS)
            writer.writeheader()
            writer.writerows(rows)
        os.replace(temporary_path, destination)
    except Exception:
        if os.path.exists(temporary_path):
            os.unlink(temporary_path)
        raise
    return str(destination)


class TorchDryRunRunner:

    def _get_runtime(self) -> DryRunRuntime:
        """Return the cached torchrun identity for this worker."""
        if self._runtime is None:
            self._runtime = DryRunRuntime.from_torchrun_env()
        return self._runtime

    def _prepare_args(self, torch_module: Any) -> Any:
        """Validate user input and return an isolated dry-run config copy."""

        if build_dry_run_batch_fn is None:
            micro_batch_size = int(base.args.train.micro_batch_size)
            sequence_length = int(base.args.data.max_seq_len)
            if micro_batch_size < 1:
                raise ValueError(
                    f"train.micro_batch_size must be >= 1, got {micro_batch_size}"
                )
            if sequence_length < 2:
                raise ValueError(
                    "data.max_seq_len must be >= 2 for causal-LM dry-run, "
                    f"got {sequence_length}"
                )
            shape = (micro_batch_size, sequence_length)

            target_device_type: str,
            batch_shapes: Dict[str, list[int]],
    ) -> Dict[str, Any]:
        """Build reproducibility metadata for a successful report."""
        dims = base.parallel_dims
        runtime = self._get_runtime()
        parallel = {
            name: int(getattr(dims, name))
            for name in ("dp_replicate", "dp_shard", "tp", "cp", "ep", "pp", "etp")
        }
        return {
            "model": base.args.model.name,
            "torch_version": torch_module.__version__,
            "rank": runtime.rank,
            "world_size": runtime.world_size,

    @staticmethod
    def _initialize_flat_buffers(base: Any) -> None:
        """Materialize enabled FSDP flat shards before tracker registration."""
        for hsdp_state in base._iter_hsdp_states():
            param_group = getattr(hsdp_state, "param_group", None)
            if param_group is not None and param_group.enable_zero_copy:
                param_group._init_flat_param_buffer()

    def _execute_fake_step(
            self,
            torch_module: Any,

            self._stage = "memory_tracking"
            tracker = _create_indexed_mem_tracker(mem_tracker_type)
            tracker.track_external(base.model, base.optimizer, batch)
            operator_trace = _create_operator_trace_mode(tracker, base.device.type)
            with tracker, operator_trace:
                if hasattr(base.model, "set_requires_gradient_sync"):
                    base.model.set_requires_gradient_sync(True)
                if hasattr(base.model, "set_is_last_backward"):
                    base.model.set_is_last_backward(True)
                global_tokens = token_count * base.parallel_dims.dp_size * base.parallel_dims.cp
                base.forward_backward_step(batch, token_count, global_tokens)
                base._run_post_fsdp_grad_reduce()
                base.optimizer.step()
                base.optimizer.zero_grad(set_to_none=True)
            operator_trace.finalize()

            self._stage = "report_generation"
            return build_memory_report(
                tracker=tracker,
                metadata=self._metadata(
                    torch_module,
                    base,

                f"rank_{runtime.rank}",
                f"rank_{runtime.rank}_memory.csv",
            )
            if self.args.train.backend != "torch":
                raise NotImplementedError(
                    "HyperDryRun currently supports train.backend='torch' only, "
                    f"got {self.args.train.backend!r}"
                )

            # ordinary meta tensors. DeviceMesh owns real CPU rank metadata;
            # deepcopying that metadata inside FakeTensorMode would incorrectly
            # mix fake meta storage with its real CPU storage.
            if base.spec.parallelize_fn is None:
                raise ValueError(
                    f"Model spec {base.spec.name!r} must define parallelize_fn "
                    "for HyperDryRun"
                )
            base.model = base.spec.parallelize_fn(base.model, base.mesh, base.args)

                (skips ``_build_layout``; see ``DTensor.from_local_with_layout``).
        """
        if isinstance(local_tensor, DTensorBase):
            # Copy from existing DTensorBase — use alias_placements to preserve multi-axis ordering
            if getattr(local_tensor, "_is_fake_wrapper", False):
                copy_placements = (
                    local_tensor.layout.alias_placements
                    if local_tensor.layout
                    else local_tensor.placements
                )
                return cls(
                    local_tensor._local_tensor,
                    local_tensor.device_mesh,
                    copy_placements,
                    local_tensor.layout,

        # path below.
        # pylint: disable=C0415
        from torch._subclasses.fake_tensor import FakeTensor
        if isinstance(local_tensor, FakeTensor):
            t = Tensor._make_wrapper_subclass(
                cls,
                local_tensor.size(),
                strides=local_tensor.stride(),
                storage_offset=local_tensor.storage_offset(),

                layout=local_tensor.layout,
                device=torch.device("meta"),
                requires_grad=local_tensor.requires_grad,
            )
            t._is_fake_wrapper = True
        else:
            t = Tensor._make_subclass(cls, local_tensor, local_tensor.requires_grad)
        t.__init_data__(local_tensor, device_mesh, placements, layout)
        return t

            # Wrapper subclasses must enter at TorchDispatch so autograd sees
            # their meta C++ device. Calling the op on unwrapped FakeTensors
            # here would make autograd inspect the outer wrapper's logical
            # CUDA/NPU device on a CPU-only build.
            return super().__torch_function__(func, types, args, kwargs)
        # pylint: disable=C0415
        from hyper_parallel.core.shard._op_dispatch import _OP_DISPATCHER
        out = _OP_DISPATCHER.dispatch(func, args, kwargs)
        return out

    @classmethod
    def __torch_dispatch__(cls, func, types, args=(), kwargs=None):
        """Dispatch fake-wrapper DTensors through HyperParallel layout rules."""
        # pylint: disable=C0415
        from hyper_parallel.core.shard._op_dispatch import _OP_DISPATCHER
        return _OP_DISPATCHER.dispatch(func, args, kwargs or {})

    def __tensor_flatten__(self):
        """Expose local storage so MemTracker can inspect wrapper DTensors."""
        context = (self._device_mesh, self._alias_placements(), self._layout)
        return ["_local_tensor"], context

    @classmethod
    def __tensor_unflatten__(cls, inner_tensors, context, outer_size, outer_stride):
        """Rebuild a traceable wrapper DTensor from its local tensor."""
        del outer_size, outer_stride
        device_mesh, placements, layout = context
        return cls(
            inner_tensors["_local_tensor"],
            device_mesh=device_mesh,
            placements=placements,
            layout=layout,

        Returns:
            Optional[Tensor]: The gradient tensor, or None if no gradient is set.
        """
        if getattr(self, "_is_fake_wrapper", False):
            return Tensor.grad.__get__(self, type(self))  # pylint: disable=C2801
        return self._local_tensor.grad

    @grad.setter
    def grad(self, value: Optional[Tensor]) -> None:

        Args:
            value (Optional[Tensor]): The gradient tensor to set, or None to clear.
        """
        if getattr(self, "_is_fake_wrapper", False):
            Tensor.grad.__set__(self, value)  # pylint: disable=C2801
            return
        self._local_tensor.grad = value

    @property
    def requires_grad(self) -> bool:

    def data(self, value: Tensor) -> None:
        """Set the underlying tensor data, extracting the local shard if a DTensor is given."""
        local_value = value.to_local() if isinstance(value, DTensorBase) else value
        if getattr(self, "_is_fake_wrapper", False):
            if self.dtype != local_value.dtype:
                raise ValueError(
                    "Fake-wrapper DTensor data replacement must preserve dtype, "
                    f"got {local_value.dtype} for {self.dtype}"
                )
            # Wrapper subclasses have no mutable storage of their own. Rebase

            # Wrapper subclasses have no mutable storage of their own. Rebase
            # the traceable inner FakeTensor, which is the storage exposed to
            # dispatch and MemTracker. FSDP legitimately rebases a full-shape
            # wrapper onto a local shard when it creates its flat buffer.
            self._local_tensor = local_value
            return
        # Tensor.data.__set__ on a Tensor subclass otherwise enters __torch_function__
        # and only rebinds _local_tensor through DTensor dispatch.
        with getattr(torch, "_C").DisableTorchFunctionSubclass():
            Tensor.data.__set__(self, local_value)

            return False
        param_data = self.hsdp_params[0]._sharded_param_data
        from torch._subclasses.fake_tensor import FakeTensor
        if isinstance(param_data, FakeTensor):
            return param_data.untyped_storage() is self._flat_param_buffer.untyped_storage()
        return param_data.data_ptr() == self._flat_param_buffer.data_ptr()

    def unshard(self, async_op: bool = False):
        """Trigger fused all-gather to reconstruct full parameters from shards.

    )
    # PyTorch may expose a renamed PrivateUse1 backend by either spelling,
    # depending on when the profiler event was materialized.
    if device_name.lower() in ("privateuse1", "privateuseone"):
        return "npu"
    return device_name


def _tensor_identity(metadata: Any) -> Optional[tuple[str, int]]:

    """Return an allocation identity, falling back to its storage pointer."""
    allocation_id = getattr(metadata, "allocation_id", None)
    if allocation_id is not None:
        return "allocation", int(allocation_id)
    pointer = getattr(metadata, "storage_data_ptr", None)
    if pointer is None:
        pointer = getattr(metadata, "ptr", None)
    return None if pointer is None else ("pointer", int(pointer))


def _flatten_tensor_inputs(value: Any) -> Iterable[Any]:
    """Yield tensor metadata objects from nested profiler operator inputs."""
    if isinstance(value, (list, tuple)):
        for item in value:
            yield from _flatten_tensor_inputs(item)
    elif getattr(value, "storage_data_ptr", None) is not None:
        yield value


def _python_frames(event: Any) -> list[Any]:
    """Collect Python callsite frames from root to the event."""

    frames = []
    parent = event
    while parent is not None:
        if _event_type_name(parent) == "PyCall":
            callsite = getattr(parent.typed[1], "callsite", None)
            if callsite is not None:
                frames.append(callsite)
        parent = getattr(parent, "parent", None)
    frames.reverse()
    deduplicated = []
    for frame in frames:
        key = (
            str(getattr(frame, "file_name", "")),
            int(getattr(frame, "line_number", 0)),
            str(getattr(frame, "function_name", "")),
        )
        if not deduplicated or deduplicated[-1][0] != key:
            deduplicated.append((key, frame))
    return [frame for _, frame in deduplicated]


def _module_name(event: Any) -> str:

    """Return the nearest profiled module class name."""
    parent = event
    while parent is not None:
        if _event_type_name(parent) == "PyCall":
            module = getattr(parent.typed[1], "module", None)
            if module is not None:
                return str(getattr(module, "cls_name", ""))
        parent = getattr(parent, "parent", None)
    return ""

        if (
            _event_type_name(parent) == "PyCall"
            and getattr(parent.typed[1], "optimizer", None) is not None
        ):
            return "optimizer"
        parent = getattr(parent, "parent", None)
    scope = str(getattr(event.typed[1], "scope", "")).lower()
    if "backward" in scope or any(
        "backward" in name or "autograd" in name for name in names

    scope = str(getattr(event.typed[1], "scope", "")).lower()
    if "backward" in scope or any(
        "backward" in name or "autograd" in name for name in names
    ):
        return "backward"
    if any("recompute" in name or "checkpointfunction" in name for name in names):
        return "recompute"
    return "forward"


def _stack_context(event: Any) -> Dict[str, Any]: