Dynamo logs generated by torch.distributed

torch.compile
1

How should I interpret the following dynamo log for distributed calls.

_dispatch.py:152] Dispatching op_call: Op(op=aten.native_layer_norm.default, args_schema=Spec(S(1) on (2, 16, 16)), [16], Spec(R on (16,)), Spec(R on (16,)), 1e-05 @ mesh: (2,))
_dispatch.py:156] output_sharding for aten.native_layer_norm.default: OutputSharding(output_spec=(DTensorSpec(mesh=DeviceMesh((tp=2), device: ‘xpu’, stride: (1,)), placements=(Shard(dim=1),), tensor_meta=TensorMeta(shape=torch.Size([2, 16, 16]), stride=(256, 16, 1), dtype=torch.float32)), DTensorSpec(mesh=DeviceMesh((tp=2), device: ‘xpu’, stride: (1,)), placements=(Shard(dim=1),), tensor_meta=TensorMeta(shape=torch.Size([2, 16, 1]), stride=(16, 1, 1), dtype=torch.float32)), DTensorSpec(mesh=DeviceMesh((tp=2), device: ‘xpu’, stride: (1,)), placements=(Shard(dim=1),), tensor_meta=TensorMeta(shape=torch.Size([2, 16, 1]), stride=(16, 1, 1), dtype=torch.float32))), redistribute_schema=None, needs_redistribute=False)
_redistribute.py:202] redistribute from S(1) to R on mesh dim 0

_dispatch.py:152] Dispatching op_call: Op(op=aten.clone.default, args_schema=Spec(S(1) on (2, 4, 16, 4)) @ mesh: (2,))
_dispatch.py:156] output_sharding for aten.clone.default: OutputSharding(output_spec=DTensorSpec(mesh=DeviceMesh((tp=2), device: ‘xpu’, stride: (1,)), placements=(Shard(dim=1),), tensor_meta=TensorMeta(shape=torch.Size([2, 4, 16, 4]), stride=(256, 64, 4, 1), dtype=torch.float32)), redistribute_schema=None, needs_redistribute=False)
_dispatch.py:152] Dispatching op_call: Op(op=aten._unsafe_view.default, args_schema=Spec(S(1) on (2, 4, 16, 4)), [8, 16, 4] @ mesh: (2,))
rank2 propagate_shape_and_sharding input_src_placements (Shard(dim=1),) global_input_shape (2, 4, 16, 4) rule (Flatten(input_dims=(InputDim(input_dim=0), InputDim(input_dim=1))), InputDim(input_dim=2), InputDim(input_dim=3)) mesh_sizes (2,) strict_view True

2

Here’s an example that might help for layer_norm

  1. You are dispatching to a specific PyTorch operator aten.native_layer_norm.default

  2. On a specific input shape ([2, 16, 16])

  3. That’s sharded (S(1)) on dimension 1 - meaning each of your 2 devices has half of the 16 columns

  4. PyTorch needs to redistribute from sharded to replicated (S(1) to R) because layer norm needs to see the full tensor to compute proper statistics

4

Thanks , is there a good documentation somewhere ? Also quick questions on the following,

  1. Does “S(1) to R on mesh dim 0” imply replicate the dim 1 shard along dim 0 ?
  2. What does [16], Spec(R on (16,)), Spec(R on (16,)) mean ?
  3. What do 1e-05 @ mesh: (2,) and [8, 16, 4] @ mesh: (2,) mean ?
5

AI explanation for

“_dispatch.py:152] Dispatching op_call: Op(op=aten.native_layer_norm.default, args_schema=Spec(S(1) on (2, 16, 16)), [16], Spec(R on (16,)), Spec(R on (16,)), 1e-05 @ mesh: (2,))” .

Pretty good !

Here’s a breakdown of the key components:

  • Dispatching op_call:

    This is the core action being performed—PyTorch is dispatching the operation to a specific device and backend for execution.

  • Op(op=aten.native_layer_norm.default:

    This identifies the specific PyTorch operator being invoked. aten.native_layer_norm.default is the PyTorch internal (or “native”) implementation for Layer Normalization.

  • args_schema:

    This describes the expected input arguments for the operation.

  • Spec(S(1) on (2, 16, 16)):

    This specifies the input tensor’s layout.

    • S(1) means it’s a sharded tensor.

    • The sharding occurs on dimension 1 (the second dimension, which has a size of 16).

    • on (2, 16, 16) refers to the input tensor’s shape.

  • [16]:

    This likely refers to the size of the weight and bias parameters for Layer Normalization, which are often the same size as the normalized dimension.

  • Spec(R on (16,)):

    This describes the layout of the weight and bias parameters.

    • R means the parameters are replicated across the devices.

    • on (16,) indicates they are 1-dimensional tensors of size 16.

  • 1e-05:

    This is the small constant (epsilon) added to the variance during the Layer Normalization calculation to prevent division by zero.

  • @ mesh: (2,):

    This shows the operation is being dispatched to a compute mesh with two devices or nodes.

In essence, this message explains that PyTorch is preparing to run a native Layer Normalization on a sharded input tensor and replicated parameters across two devices in a distributed system

6

AI explanation for -

OutputSharding(output_spec=(DTensorSpec(mesh=DeviceMesh((tp=2), device: ‘xpu’, stride: (1,)), placements=(Shard(dim=1),), tensor_meta=TensorMeta(shape=torch.Size([2, 16, 16]), stride=(256, 16, 1), dtype=torch.float32)), DTensorSpec(mesh=DeviceMesh((tp=2), device: ‘xpu’, stride: (1,)), placements=(Shard(dim=1),), tensor_meta=TensorMeta(shape=torch.Size([2, 16, 1]), stride=(16, 1, 1), dtype=torch.float32)), DTensorSpec(mesh=DeviceMesh((tp=2), device: ‘xpu’, stride: (1,)), placements=(Shard(dim=1),), tensor_meta=TensorMeta(shape=torch.Size([2, 16, 1]), stride=(16, 1, 1), dtype=torch.float32))), redistribute_schema=None, needs_redistribute=False)

The provided information describes the OutputSharding of a distributed tensor operation, likely within a PyTorch distributed training context, specifically utilizing DTensor for distributed tensor management.

Here’s a breakdown of the components:

  • OutputSharding(...):

    This indicates that the object describes how the output of a computation is sharded across devices.

  • output_spec=(...):

    This tuple defines the DTensorSpec for each of the output tensors. In this case, there are three output tensors.

  • DTensorSpec(...):

    This specifies the distributed properties of a single tensor.

    • mesh=DeviceMesh((tp=2), device: 'xpu', stride: (1,)): This defines the device mesh, which is a logical topology of devices.

      • tp=2: Indicates a tensor parallelism (TP) dimension with a size of 2, meaning the tensor is sharded across two devices for tensor parallelism.

      • device: 'xpu': Specifies that the devices are of type ‘xpu’ (likely Intel Gaudi AI accelerators).

      • stride: (1,): Refers to the stride within the device mesh, often related to how ranks are ordered.

    • placements=(Shard(dim=1),): This indicates how the tensor is sharded. Shard(dim=1) means the tensor is sharded along its second dimension (index 1).

    • tensor_meta=TensorMeta(...): This provides metadata about the local shard of the tensor on a single device.

      • shape=torch.Size([2, 16, 16]), shape=torch.Size([2, 16, 1]), shape=torch.Size([2, 16, 1]): These are the local shapes of the three sharded output tensors on each device. The original unsharded tensor would have a shape that combines these local shards along the sharded dimension.

      • stride=(...): This describes the memory layout of the local tensor shard.

      • dtype=torch.float32: Specifies the data type of the tensor.

  • redistribute_schema=None:

    This indicates that no explicit redistribution schema is applied after the operation.

  • needs_redistribute=False:

    This confirms that no redistribution of the output tensors is required after the operation, implying the current sharding is already the desired final state.

In summary, this OutputSharding object describes that the operation produces three output tensors, each sharded along its second dimension across two XPU devices, and no further redistribution is needed for these outputs.

7

AI explanation for

_redistribute.py:202] redistribute from S(1) to R on mesh dim 0

When you see the message redistribute from S(1) to R on mesh dim 0, it signifies the following sequence of events:

  1. Input state: The framework was working with a tensor that was sharded along its second axis (dimension 1). Each device in the mesh held a different piece of the tensor.

  2. Required action: A subsequent operation required a fully replicated tensor. For example, a global reduction operation or a specific mathematical function might need the complete tensor on every device to proceed.

  3. The redistribution: The framework automatically performed a “gather” collective communication operation. It collected all the shards from every device, reconstructed the full tensor, and then broadcast that full tensor back out to all the devices.

  4. Device mesh axis: The communication to perform this gather-and-broadcast was coordinated along mesh dim 0, meaning it used the devices organized along the first dimension of the logical device grid.

  5. Output state: Each device now holds a complete, identical copy of the tensor.