Repository: pytorch
Issue: `Torch.distributed.tensor.parallel.Style.Rowwiseparallel` when send tensor, the use of the underlying operator is `aten.to.dtype_layout`, lack of distributed strategy. But after adding a policy, `AsyncCollectiveTensor` no placements error 
Labels: oncall: distributed

### 🐛 Describe the bug

```model = Transformer.from_model_args(simple_llama2_config).to("cuda")

#init model weights
model.init_weights()

#parallelize the first embedding and the last linear out projection
model = parallelize_module(
    model,
    tp_mesh,
    {
        "tok_embeddings": RowwiseParallel(       # **Here's the problem**
            input_layouts=Replicate(),
            output_layouts=Shard(1),  
        ),
        "norm": SequenceParallel(),
        "output": ColwiseParallel(
            input_layouts=Shard(1),
            output_layouts=Replicate()
        ),
    }
)
```
### first error:
```Traceback (most recent call last):
  File ~/nfs/project/vision_model/GeoChat/mutiSateChatInOrbit/run_demo.py:50 in main
    eval_instance.eval_model()
  File ~/nfs/project/vision_model/GeoChat/mutiSateChatInOrbit/dist_batch_geochat_vqa.py:348 in eval_model
    output_ids = self.model.generate(final_input_tensors,
  File ~/nfs/anaconda3/envs/geochat/lib/python3.10/site-packages/torch/utils/_contextlib.py:116 in decorate_context
    return func(*args, **kwargs)
  File ~/nfs/anaconda3/envs/geochat/lib/python3.10/site-packages/transformers/generation/utils.py:1538 in generate
    return self.greedy_search(
  File ~/nfs/anaconda3/envs/geochat/lib/python3.10/site-packages/transformers/generation/utils.py:2362 in greedy_search
    outputs = self(
  File ~/nfs/anaconda3/envs/geochat/lib/python3.10/site-packages/torch/nn/modules/module.py:1736 in _wrapped_call_impl
    return self._call_impl(*args, **kwargs)
  File ~/nfs/anaconda3/envs/geochat/lib/python3.10/site-packages/torch/nn/modules/module.py:1747 in _call_impl
    return forward_call(*args, **kwargs)
  File ~/nfs/anaconda3/envs/geochat/lib/python3.10/site-packages/accelerate/hooks.py:165 in new_forward
    output = old_forward(*args, **kwargs)
  File ~/nfs/project/vision_model/GeoChat/geochat/model/language_model/geochat_llama.py:81 in forward
    input_ids, attention_mask, past_key_values, inputs_embeds, labels = self.prepare_inputs_labels_for_multimodal(input_ids, attention_mask, past_key_values, labels, images)
  File ~/nfs/project/vision_model/GeoChat/geochat/model/geochat_arch.py:187 in prepare_inputs_labels_for_multimodal
    cur_new_input_embeds.append(self.get_model().embed_tokens(cur_input_ids[:image_token_start]))  # 获取文本1的embedding
  File ~/nfs/anaconda3/envs/geochat/lib/python3.10/site-packages/torch/nn/modules/module.py:1736 in _wrapped_call_impl
    return self._call_impl(*args, **kwargs)
  File ~/nfs/anaconda3/envs/geochat/lib/python3.10/site-packages/torch/nn/modules/module.py:1844 in _call_impl
    return inner()
  File ~/nfs/anaconda3/envs/geochat/lib/python3.10/site-packages/torch/nn/modules/module.py:1790 in inner
    result = forward_call(*args, **kwargs)
  File ~/nfs/anaconda3/envs/geochat/lib/python3.10/site-packages/accelerate/hooks.py:160 in new_forward
    args, kwargs = module._hf_hook.pre_forward(module, *args, **kwargs)
  File ~/nfs/anaconda3/envs/geochat/lib/python3.10/site-packages/accelerate/hooks.py:290 in pre_forward
    return send_to_device(args, self.execution_device), send_to_device(
  File ~/nfs/anaconda3/envs/geochat/lib/python3.10/site-packages/accelerate/utils/operations.py:151 in send_to_device
    return honor_type(
  File ~/nfs/anaconda3/envs/geochat/lib/python3.10/site-packages/accelerate/utils/operations.py:83 in honor_type
    return type(obj)(generator)
  File ~/nfs/anaconda3/envs/geochat/lib/python3.10/site-packages/accelerate/utils/operations.py:152 in <genexpr>
    tensor, (send_to_device(t, device, non_blocking=non_blocking, skip_keys=skip_keys) for t in tensor)
  File ~/nfs/anaconda3/envs/geochat/lib/python3.10/site-packages/accelerate/utils/operations.py:167 in send_to_device
    return tensor.to(device, non_blocking=non_blocking)
  File ~/nfs/anaconda3/envs/geochat/lib/python3.10/site-packages/torch/_compile.py:32 in inner
    return disable_fn(*args, **kwargs)
  File ~/nfs/anaconda3/envs/geochat/lib/python3.10/site-packages/torch/_dynamo/eval_frame.py:632 in _fn
    return fn(*args, **kwargs)
  File ~/nfs/anaconda3/envs/geochat/lib/python3.10/site-packages/torch/distributed/tensor/_api.py:341 in __torch_dispatch__
    return DTensor._op_dispatcher.dispatch(
  File ~/nfs/anaconda3/envs/geochat/lib/python3.10/site-packages/torch/distributed/tensor/_dispatch.py:169 in dispatch
    self.sharding_propagator.propagate(op_info)
  File ~/nfs/anaconda3/envs/geochat/lib/python3.10/site-packages/torch/distributed/tensor/_sharding_prop.py:201 in propagate
    OutputSharding, self.propagate_op_sharding(op_info.schema)
  File ~/nfs/anaconda3/envs/geochat/lib/python3.10/site-packages/torch/distributed/tensor/_sharding_prop.py:46 in __call__
    return self.cache(*args, **kwargs)
  File ~/nfs/anaconda3/envs/geochat/lib/python3.10/site-packages/torch/distributed/tensor/_sharding_prop.py:450 in propagate_op_sharding_non_cached
    raise NotImplementedError(
NotImplementedError: Operator aten.to.dtype_layout does not have a sharding strategy registered.
```
### The band-aid modification error 01
#### alteration
`torch.distributed.tensor.ops._pointwise_ops` ，添加`dtype_layout`的策略

```
linear_pointwise_ops = [
    aten.div.Scalar,  # this op is linear on the first argument, and the second argument is scalar, so it fits as a linear op.
    aten.div_.Scalar,  # this op is linear on the first argument, and the second argument is scalar, so it fits as a linear op.
    aten.to.dtype,
    # aten.to.dtype_layout,  # here is*
    aten.add.Tensor,
    aten.add_.Tensor,
]
```

#### error
```
  File ~/nfs/project/vision_model/GeoChat/geochat/model/language_model/geochat_llama.py:81 in forward
    input_ids, attention_mask, past_key_values, inputs_embeds, labels = self.prepare_inputs_labels_for_multimodal(input_ids, attention_mask, past_key_values, labels, images)
  File ~/nfs/project/vision_model/GeoChat/geochat/model/geochat_arch.py:187 in prepare_inputs_labels_for_multimodal
    cur_new_input_embeds.append(self.get_model().embed_tokens(cur_input_ids[:image_token_start]))  # 获取文本1的embedding
  File ~/nfs/anaconda3/envs/geochat/lib/python3.10/site-packages/torch/nn/modules/module.py:1736 in _wrapped_call_impl
    return self._call_impl(*args, **kwargs)
  File ~/nfs/anaconda3/envs/geochat/lib/python3.10/site-packages/torch/nn/modules/module.py:1844 in _call_impl
    return inner()
  File ~/nfs/anaconda3/envs/geochat/lib/python3.10/site-packages/torch/nn/modules/module.py:1803 in inner
    hook_result = hook(self, args, result)
  File ~/nfs/anaconda3/envs/geochat/lib/python3.10/site-packages/torch/distributed/tensor/_api.py:896 in <lambda>
    lambda mod, inputs, outputs: output_fn(mod, outputs, device_mesh)
  File ~/nfs/anaconda3/envs/geochat/lib/python3.10/site-packages/torch/distributed/tensor/parallel/style.py:297 in _prepare_output_fn
    if outputs.placements != output_layouts:
AttributeError: 'AsyncCollectiveTensor' object has no attribute 'placements'
```
### The band-aid modification error 02
#### alteration
process ` asynccollectivetensor` type: ` Torch.distributed.tensor.parallel.Style.Rowwiseparallel: 297`
```
class RowwiseParallel(ParallelStyle):
    ……
    @staticmethod
    def _prepare_output_fn(output_layouts, use_local_output, mod, outputs, device_mesh):
        # Rowwise sharding produces partial output, depending on output layouts:
        # 1. to replicate -> allreduce
        # 2. to shard -> reduce_scatter
        # ----------------------------------------------------------------创可贴 1
        # from torch.distributed._functional_collectives  import AsyncCollectiveTensor
        # if isinstance(outputs, AsyncCollectiveTensor):
        # #     # 等待异步操作完成
        #     outputs.wait()
        #     new_outputs = outputs.clone()
        #     print("new_outputs",new_outputs)
        #     outputs = distribute_tensor(outputs, device_mesh, output_layouts)
        #     # outputs = distributed_outputs.clone()
        #     outputs = outputs.redistribute(placements=output_layouts, async_op=False)
        # ----------------------------------------------------------------创可贴 2
        # if isinstance(outputs, torch.distributed._functional_collectives.AsyncCollectiveTensor):
        #     result = outputs.trigger_wait()
        #     new_outputs = result.clone().detach()
       
        #     print("new_outputs",new_outputs)
        #     new_outputs = new_outputs.redistribute(placements=output_layouts, async_op=False)
        # ----------------------------------------------------------------创可贴3
        # def handle_async_tensor(async_tensor):
        #     # 如果是 AsyncCollectiveTensor，则等待操作完成并返回结果
        #     if isinstance(async_tensor, torch.distributed._functional_collectives.AsyncCollectiveTensor):
        #         result = async_tensor.trigger_wait()
        #         return result.clone().detach()
        #     # 否则直接返回原对象
        #     return async_tensor

        # def prepare_output(outputs):
        #     from torch.utils._pytree import tree_map, tree_flatten, tree_unflatten
        #     flat_outputs, spec = tree_flatten(outputs)
        #     processed_flat_outputs = [handle_async_tensor(item) for item in flat_outputs]
        #     # 重新构造原来的结构
        #     processed_outputs = tree_unflatten(processed_flat_outputs, spec)
            
        #     # 如果 outputs 是一个张量，直接克隆；如果是更复杂的结构，递归处理
        #     if torch.is_tensor(processed_outputs):
        #         new_outputs = processed_outputs.clone()
        #     else:
        #         # 递归处理非张量类型的输出
        #         new_outputs = tree_map(lambda x: x.clone() if torch.is_tensor(x) else x, processed_outputs)
            
        #     return new_outputs
        
        # new_outputs = prepare_output(outputs)
        # print("new_outputs",new_outputs)
        # ----------------------------------------------------------------原代码
        if outputs.placements != output_layouts:
            outputs = outputs.redistribute(placements=output_layouts, async_op=True)
        # back to local tensor if use_local_output is True
        return outputs.to_local() if use_local_output else outputs
```
#### but always: error
```
  File ~/nfs/anaconda3/envs/geochat/lib/python3.10/site-packages/torch/distributed/tensor/parallel/style.py:264 in _prepare_output_fn
    result = outputs.trigger_wait()
  File ~/nfs/anaconda3/envs/geochat/lib/python3.10/site-packages/torch/distributed/_functional_collectives.py:611 in trigger_wait
    out = wait_tensor(self.elem)
  File ~/nfs/anaconda3/envs/geochat/lib/python3.10/site-packages/torch/distributed/_functional_collectives.py:140 in wait_tensor
    return torch.ops._c10d_functional.wait_tensor(tensor)  # type: ignore[attr-defined]
  File ~/nfs/anaconda3/envs/geochat/lib/python3.10/site-packages/torch/_ops.py:1116 in __call__
    new_op = self._op(*args, **(kwargs or {}))
RuntimeError: Inplace update to inference tensor outside InferenceMode is not allowed.You can make a clone to get a normal tensor before doing inplace update.See https://github.com/pytorch/rfcs/pull/17 for more details.
```

### Versions

cuda11.8
python == 3.10.14, 
torch==2.5.1+cu118

cc @H-Huang @awgu @kwen2501 @wanchaol @fegin @fduwjj @wz337 @wconstab @d4l3k @c-p-i-o