o j9:j~@sUddlmZddlZddlZddlmZmZddlmZddl m Z m Z m Z m Z mZmZddlZddlmZddlmZddlmZdd lmZmZerYdd lmZdd lmZeeZGd d d e Z e!e!e"dfedBfZ#de$d<Gddde%Z&edddGdddZ'edddGddde'Z(e'e(BZ)de$d<eddGdddZ*edddGdddZ+edddGdd d Z,dxd&d'Z-dyd+d,Z.Gd-d.d.Z/Gd/d0d0eZ0d1d2dzd5d6Z1d7d8Z2 9d{d|d@dAZ3 9d{d}dCdDZ4eddGdEdFdFZ5d~dHdIZ6e ddJddPdQZ7e ddUdQZ7d1dJddWdQZ7dddXdYZ8e ddd\d]Z9e ddd_d]Z9 1dddad]Z9d1d1dbddjdkZ:dd1dbddndoZ;ddvdwZs z_TensorMeta.from_tensordevicetorch.device | strcCst||}||j|S)zReconstruct a tensor on ``device`` from this metadata. Args: device: Target device for the tensor. Returns: An empty strided tensor on ``device``. )_make_tensor_from_metarequires_grad_r0)rr7trrr to_tensorVs z_TensorMeta.to_tensorother list[str]cCs||krgSg}|j|jkr|d|jd|j|j|jkr.|d|jd|j|j|jkrA|d|jd|j|S)zReturn field-by-field differences with ``other``. Args: other: Metadata to compare against. Returns: List of human-readable difference strings (empty if equal). zshape mismatch:  vs zstride mismatch: zdtype mismatch: )r+appendr-r.rr=diffsrrrget_diffcs    z_TensorMeta.get_diffN)r1r2rr))r7r8rr2r=r)rr>) r!r"r#r$__annotations__ staticmethodr6r<rCrrrrr)0s    r)c@seZdZUdZedddZded<eddZd ed <eddZd ed <eddZ d ed<eddZ ded<e d%ddZ e d&ddZd'ddZd(d#d$ZdS)) _DTensorMetaaDTensor metadata extending :class:`_TensorMeta` with distribution info. Inherited fields (shape, stride, etc.) are LOCAL shard attributes. Additional fields capture global shape and placement information needed to reconstruct a :class:`DTensor` via ``DTensor.from_local()``. The :class:`DeviceMesh` is **not** stored (not serializable for P2P); it is looked up from :class:`_MeshCache` using ``(mesh_dim_names, mesh_layout)`` as the key. cCs tgSr)torchSizerrrrs z_DTensorMeta.)default_factoryr* global_shaper)defaultr, global_strideztuple[Placement, ...] placementsrrNrrdtensorrrc CsJ|j}t|jj|j|j|j|j||jj|j rt |j nd|j d S)zCreate metadata from a DTensor. Args: dtensor: The DTensor to extract metadata from. Returns: Metadata capturing both local and global attributes. r) r+r-r.r0rLrNrOrr) device_meshrG _local_tensorr+r-r.r0_specrOrtuple_layout)rPrQrrr from_dtensors z_DTensorMeta.from_dtensorr%cCs |j|jfS)zcCs||krgSt||}t|trr|j|jkr$|d|jd|j|j|jkr7|d|jd|j|j|jkrJ|d|jd|j|j|jkr]|d|jd|j|j |j krp|d|j d|j |S|d|S)zReturn field-by-field differences, including DTensor-specific fields. Args: other: Metadata to compare against. Returns: List of human-readable difference strings (empty if equal). zglobal_shape mismatch: r?zglobal_stride mismatch: zplacements mismatch: zmesh_dim_names mismatch: zmesh_layout mismatch: z!type: _DTensorMeta vs _TensorMeta) r)rCr4rGrLr@rNrOrrrArrrrCs6         z_DTensorMeta.get_diff)rPrrrG)rr%)r7r8rYrrrrD)r!r"r#r$rrLrErNrOrrrFrVpropertyrXr]rCrrrrrG{s"     rG TensorMeta)r(c@s`eZdZUdZdZded<dZded<dZded<dZded<dd d Z dd dZ dddZ dS) _StageMetazPConsolidated tensor metadata for a pipeline stage's forward and backward passes.Ntuple[TensorMeta, ...] | Noneinputsoutputs$tuple[TensorMeta | None, ...] | None input_grads output_gradsrr/cCs"tdd|j|j|j|jfDS)z)Check if any metadata field is populated.css|]}|duVqdSrr).0vrrr s  z%_StageMeta.has_any..)anyrbrcrerfrWrrrhas_anysz_StageMeta.has_anycCs2|j|jfD]}|rtdd|DrdSqdS)z3Check if any input/output metadata is DTensor type.css|] }|rt|tVqdSr)r4rGrgmrrrrisz*_StageMeta.has_dtensors..TF)rbrcrj)rmetasrrr has_dtensorss z_StageMeta.has_dtensorscCs|jduo |jduS)z-Check if forward metadata is fully populated.N)rbrcrWrrris_complete_for_forwardsz"_StageMeta.is_complete_for_forward)rr/) r!r"r#r$rbrErcrerfrkrorprrrrr`s       r`c@eZdZUdZded<dS)_StageForwardMetazLForward metadata transmitted from stage *i* to stage *i+1* during inference.tuple[TensorMeta, ...] forward_metasNr!r"r#r$rErrrrrr!s  rrc@rq)_StageBackwardMetauBackward metadata transmitted from stage *i* to stage *i-1* during inference. Gradient placements may differ from forward activations (e.g., ``Replicate`` → ``Partial``). tuple[TensorMeta | None, ...]backward_metasNrurrrrrv(s  rvmetar7r8rr2cCstj|j|j|j|dS)zCreate a tensor from metadata. Args: meta: Metadata with shape, stride, and dtype. device: Target device for the tensor. Returns: Empty tensor preserving the exact memory layout. )sizer-r.r7)rH empty_stridedr+r-r.)ryr7rrrr95s  r9 tensor_metasrstuple[_TensorMeta | None, ...]cCstdd|DS)zDerive gradient metadata from tensor metadata. Returns metadata with the same shape/stride/dtype but ``requires_grad=False``. Entries where the source has ``requires_grad=False`` become ``None``. css0|]}|jrt|j|j|jddndVqdS)Fr3N)r0r)r+r-r.rlrrrriRs z%_derive_grad_metas..)rT)r|rrr_derive_grad_metasJsr~c@sNeZdZdZddddZdd d ZdddZdddZd ddZd!ddZ dS)" _MeshCachezCache for :class:`DeviceMesh` objects keyed by ``(mesh_dim_names, mesh_layout)``. Assumes all pipeline stages share the same rank tensor (true for TorchTitan-style frameworks where meshes derive from a common world). N get_mesh_cbGetMeshCallback | NonerNonecCsi|_||_dSr)_cache _get_mesh_cb)rrrrr__init__as z_MeshCache.__init__keyr%rcCst||jvr |j|S|\}}|jdurtd|d|d|||}|dur3td|d|d||j|<|S)aReturn a cached mesh, or create one via the callback. Args: key: Cache key ``(mesh_dim_names, mesh_layout)``. Returns: The ``DeviceMesh``. Raises: PipeliningMetadataError: If not cached and no callback provided. Nz+Mesh not found in cache for mesh_dim_names=z, mesh_layout=z`, and no get_mesh callback provided. Provide a get_mesh callback or use DTensors in static mode.z>Mesh lookup failed: callback returned None for mesh_dim_names=z6. Ensure all stages use meshes from the same universe.)rrr&)rrrrrYrrrget_meshes(    z_MeshCache.get_meshrYcCs||j|<dS)zAdd a mesh to the cache.Nr)rrrYrrrputsz_MeshCache.puttensorstuple[torch.Tensor | None, ...]cCsT|D]%}t|tr'|j}|jrt|jnd}|j}||f}||jvr'||j|<qdS)zJExtract and cache meshes from any :class:`DTensor` instances in *tensors*.rN)r4rrQrrTrUr)rrr1rY dim_namesrrrrrupdate_from_tensorss   z_MeshCache.update_from_tensorsr/cCs ||jvSrr)rrrrr __contains__ z_MeshCache.__contains__intcCs t|jSr)lenrrWrrr__len__rz_MeshCache.__len__r)rrrr)rr%rr)rr%rYrrr)rrrr)rr%rr/)rr) r!r"r#r$rrrrrrrrrrrZs   "  rc@s&eZdZdZdZdZed d d Zd S) InferenceModeaPipeline-level metadata inference mode, determined collectively across all PP ranks. The mode is set by the schedule (not individual stages) because ``has_backward`` is only known at schedule creation time and all stages must agree to avoid P2P hangs. .. attribute:: STATIC All stages have sufficient metadata; runtime inference is skipped. .. attribute:: DYNAMIC At least one stage requires runtime metadata inference. staticdynamicryr`stage_has_backwardr/rcCs<|sdS|s dS|sdS|jdus|jdurdSdS)a'Determine whether dynamic metadata inference is needed for a stage. Args: meta: Stage metadata from user-provided args. stage_has_backward: Whether a backward pass will be performed. Returns: ``True`` if dynamic inference is needed. TFN)rprorerf)clsryrrrr needs_dynamics zInferenceMode.needs_dynamicN)ryr`rr/rr/)r!r"r#r$STATICDYNAMIC classmethodrrrrrrs rFdetachrr/cCs4t|\}}|rdd|D}t||}||fS|S)a;Flatten ``args`` into a list, optionally detaching tensors. Args: args: Nested arguments to flatten. detach: If ``True``, detach tensors while preserving ``requires_grad``. Returns: ``(new_args, flat_detached_args)`` when ``detach=True``; ``flat_args`` list otherwise. cSs,g|]}t|tjr||jn|qSr)r4rHTensorrr:r0)rgarrr s  z flatten_args..r)argsr flat_argstreespec flat_detachednew_argsrrr flatten_argss  rcCs t|ddS)zHFlatten and detach. Deprecated: use ``flatten_args(args, detach=True)``.Tr)r)rrrrflatten_args_detachs rlooppp_sizer num_stagesstylestrdict[int, int]cCsi}|dkrt|D]}||||<q |S|dkrS||dkr*td|d|dd}t|D] }|||<|d|dkr?q0||ddkrL|d7}q0|d8}q0|Std |d ) z Compute the stage id to rank mapping for either a looped or V-style schedule. Most commonly num_stages == pp_size * 2, but this function can be used to compute the mapping for any number of stages per rank. rrhrz num_stages z% must be evenly divisible by pp_size z for V scheduleszStyle z is not supported.)range ValueError)rrrmapping stage_index rank_indexrrrgenerate_stage_to_rank_mappings(      rdict[int, list[int]]cCsZt|||}i}|D]\}}||vrg||<|||q |D]}|q$|S)a Compute the rank to stage id mapping for either a looped or V-style schedule. This function inverts the stage_to_rank_mapping to get which stages are assigned to each rank. Returns a dictionary mapping rank -> list of stage indices assigned to that rank. )ritemsr@valuessort)rrr stage_to_rankrank_to_stagesstage_idrankstagesrrrgenerate_rank_to_stage_mappings   rc@s*eZdZUdZded<ded<ded<dS) PipeInfoz> Captures information for a pipeline (`Pipe` object). zfx.Graphgraphrrr/has_loss_and_backwardNrurrrrr6s  rr1cCst|tr t|St|S)aExtract metadata from a tensor. Handles both plain Tensor and DTensor correctly: DTensors are dispatched to ``_DTensorMeta.from_dtensor`` which captures local shard attributes plus global shape/placement info, while plain tensors use ``_TensorMeta.from_tensor``. Args: tensor: A plain tensor or DTensor. Returns: ``_TensorMeta`` for plain tensors, ``_DTensorMeta`` for DTensors. )r4rrGrVr)r6r5rrrextract_tensor_metaFs   r) allow_nonertuple[torch.Tensor, ...] | NonerLiteral[False]racCrrrrrrrrextract_tensor_metasZr&tuple[torch.Tensor | None, ...] | None Literal[True]rdcCrrrrrrrrbrAtuple[torch.Tensor | None, ...] | tuple[torch.Tensor, ...] | NonecCs`|durdSg}d}|D]}t|tjr|t|q d}|dq |s,|r,tdt|S)aExtract metadata from a tuple of tensors. Args: tensors: Tuple of tensors (may include ``None`` when ``allow_none=True``). allow_none: If ``True``, preserve ``None`` elements (for gradients). Returns: Tuple of ``TensorMeta``, or ``None`` if ``tensors`` is ``None``. Raises: PipeliningMetadataError: If ``None`` found and ``allow_none=False``. NFTz_None values are not allowed in tensor metadata tuples. Use allow_none=True for optional values.)r4rHrr@rr&rT)rrmetas_with_nonehas_noner;rrrrjs  cCs&|r|n|}t|tr|S|S)uConvert a DTensor to its local shard, or return a plain tensor as-is. When ``detach=True``, the tensor is detached before conversion — this applies to both DTensors and plain tensors. Args: tensor: A tensor that may be a DTensor. detach: If ``True``, detach before ``to_local()`` to avoid redistribution during backward. Returns: The local tensor component. )rr4rto_local)r1rmaybe_detached_tensorrrrto_local_if_dtensors rrCtorch.Tensor | tuple[torch.Tensor, ...] | list[torch.Tensor] | NonecCrrrrrrrrvalidate_and_normalize_to_tupler rQtorch.Tensor | tuple[torch.Tensor | None, ...] | list[torch.Tensor | None] | NonecCrrrrrrrrstorch.Tensor | tuple[torch.Tensor, ...] | tuple[torch.Tensor | None, ...] | list[torch.Tensor] | list[torch.Tensor | None] | NonecCs|durdSt|tjr|fSt|ttfrMt|D]'\}}|dur-|s,td|dqt|tjsAtd|dt|jdqt|trKt|S|Stdt|jd)aNormalize ``args`` to a tuple and validate that all elements are tensors. Args: args: A single tensor, tuple/list of tensors, or ``None``. allow_none: If ``True``, permit ``None`` elements (for gradients). Returns: Tuple of tensors, or ``None`` if ``args`` is ``None``. Raises: PipeliningMetadataError: On non-tensor values (or ``None`` when ``allow_none=False``). Nz Stage arg[zF] is None. Stage args must be tensors. Use kwargs for optional values.z ] has type zC. All stage args must be tensors. Use kwargs for non-tensor inputs.zcCst|tjr t|}n|}t|t|urEdt|jdt|jg}|r1t|d|d|rCtj|d|ddt dd|S| |}|rm|rZt|dd ||rmtj|d d |dt dd|S) al Compare expected metadata against actual tensor or metadata. This is the unified validation/comparison function that uses get_diff() from metadata classes. Works with both plain tensors and DTensors. For plain tensors: compares shape/stride/dtype/requires_grad. For DTensors: compares all properties including global shape and placements. Args: desc: Description for error/warning messages. expected: Expected tensor metadata (_TensorMeta or _DTensorMeta). actual: Actual tensor or metadata to compare against. raise_on_mismatch: If True, raise PipeliningMetadataError on mismatch. warn_on_mismatch: If True, issue a warning on mismatch. Returns: List of differences (empty if metadata matches). Raises: PipeliningMetadataError: If raise_on_mismatch=True and differences exist. ztype: expected , got : rz: Metadata type mismatch. z.. Using dynamically inferred metadata instead.r stacklevelz; z: Metadata mismatch. ) r4rHrrrr!r&warningswarn UserWarningrCjoin)rrrrr actual_meta type_diffrBrrrvalidate_metadatas4   rrw,tuple[torch.Tensor | TensorMeta | None, ...]c Cs"t|t|kr)|dt|dt|}|rt||r&tj|tdd|gSg}tt||ddD]Z\}\}} |durC| durCq4|dusK| durz|d|d |durWd nd d | dur`d nd }|rjt||rttj|tdd||q4t|d|d || ||d} | | q4|S)a2Validate metadata for a tuple of tensors element-wise. Args: desc: Description prefix for error/warning messages. expected: Tuple of expected metadata (may include ``None`` for grads). actual: Tuple of actual tensors or metadata to compare against. raise_on_mismatch: If ``True``, raise on the first mismatch. warn_on_mismatch: If ``True``, issue warnings for mismatches. Returns: Aggregated list of difference strings. Raises: PipeliningMetadataError: If lengths differ or on mismatch. z : expected z tensors, got rrTstrictN[z ]: expected rmetadatar]r) rr&rrrrzipr@rextend) rrrrrmsg all_diffsrexpactrBrrrvalidate_tensors_metadata-s>  rrtuple[torch.Tensor, ...]gradsris_inputrc CsZ|rdnd}|d}|d}t|t|kr0td|d|dt|d|dt|d tt||d d D]q\}\}} |js`| d ur`td|d|d |d|d |dt| jd |jr| d urtjd|d|d |d|d |d t ddt |t r|jr| d urt | t std|d|d |d|d |dt| jd q9d S)u/ Validate the args↔grads contract for static mode. Enforces four rules for each (arg, grad) pair: 1. len(args) must equal len(grads). 2. If arg.requires_grad is False, grad must be None. 3. If arg.requires_grad is True and grad is None, emit a warning (this is legal at pipeline boundaries but may indicate a bug). 4. If arg is a DTensor with requires_grad=True and grad is not None, grad must also be a DTensor. Args: stage_index: The stage index for error messages. args: Tuple of forward tensors. grads: Tuple of gradient tensors (can include None). is_input: True for input_args/input_grads, False for output_args/output_grads. Raises: PipeliningMetadataError: If any hard rule (1, 2, or 4) is violated. inputoutput_args_gradszStage rz length (z) does not match zo). Each forward tensor must have a corresponding gradient entry (use None for tensors that don't require grad).TrNrz] has requires_grad=False, but z] is not None (zE). Non-differentiable tensors must have None as their gradient entry.z] has requires_grad=True, but zT] is None. This is legal at pipeline boundaries but may indicate a missing gradient.rrz,] is a DTensor with requires_grad=True, but z] is za, expected DTensor or None. DTensor gradients may have different placements than forward tensors.) rr&rrr0rr!rrrr4r) rrrrkind args_name grads_namerrgradrrr'validate_static_arg_grad_correspondencefsd    r)ryr)r7r8rr2)r|rsrr})rr/)r)rrrrrrrr)rrrrrrrr)r1r2rr_)rrrrrra)rrrrrrd)rrrr/rrd)F)r1r2rr/rr2).)rrrrrr)rrrrrr)rrrr/rr) rrrr_rrrr/rr/rr>) rrrrwrrrr/rr/rr>) rrrrrrrr/rr)= __future__rloggingr dataclassesrrenumrtypingrrrr r r rHr torch.distributed._mesh_layoutr torch.distributed.tensorrtorch.utils._pytreerrtorch.distributed.device_meshr(torch.distributed.tensor.placement_typesr getLoggerr!loggerrrTrr%rE RuntimeErrorr&r)rGr_r`rrrvr9r~rrrrrrrrrrrrrrrrrrs           J   H8 #     $ :J9