Model support for `stable_diffusion_text_encoder` with Torch_XLA2

[ManfeiBai]

Fix the model test for stable_diffusion_text_encoder.py

1. setup env according to Run a model under torch_xla2
2. Run model test under run_torchbench/ with python models/your_target_model_name.py
3. Fix the failure.

Please refer to
this guide as guide to fix:

Also refer to these PRs:

• Add documentation on fixing a torchbench model on torchxla2 #7583

[barney-s]

Running text encoder:

> python  models/stable_diffusion_text_encoder.py

/usr/local/google/home/barni/miniconda3/envs/torch_xla2/lib/python3.10/site-packages/torch/distributed/distributed_c10d.py:337: UserWarning: Device capability of jax unspecified, assuming `cpu` and `cuda`. Please specify it via the `devices` argument of `register_backend`.
  warnings.warn(
Loading pipeline components...:  33%|█████████████████████████████████████▎                                                                          | 2/6 [00:00<00:00,  5.34it/s]/usr/local/google/home/barni/miniconda3/envs/torch_xla2/lib/python3.10/site-packages/transformers/tokenization_utils_base.py:1601: FutureWarning: `clean_up_tokenization_spaces` was not set. It will be set to `True` by default. This behavior will be depracted in transformers v4.45, and will be then set to `False` by default. For more details check this issue: https://github.com/huggingface/transformers/issues/31884
  warnings.warn(
Loading pipeline components...: 100%|████████████████████████████████████████████████████████████████████████████████████████████████████████████████| 6/6 [00:00<00:00, 12.29it/s]
WARNING:root:Duplicate op registration for aten.trunc
Traceback (most recent call last):
  File "/usr/local/google/home/barni/workspace/pytorch-tpu/run_torchbench/models/stable_diffusion_text_encoder.py", line 61, in <module>
    sys.exit(main())
  File "/usr/local/google/home/barni/workspace/pytorch-tpu/run_torchbench/models/stable_diffusion_text_encoder.py", line 39, in main
    xla2_ans = model(*example)
  File "/usr/local/google/home/barni/miniconda3/envs/torch_xla2/lib/python3.10/site-packages/torch/nn/modules/module.py", line 1736, in _wrapped_call_impl
    return self._call_impl(*args, **kwargs)
  File "/usr/local/google/home/barni/miniconda3/envs/torch_xla2/lib/python3.10/site-packages/torch/nn/modules/module.py", line 1747, in _call_impl
    return forward_call(*args, **kwargs)
  File "/usr/local/google/home/barni/miniconda3/envs/torch_xla2/lib/python3.10/site-packages/transformers/models/clip/modeling_clip.py", line 986, in forward
    return self.text_model(
  File "/usr/local/google/home/barni/miniconda3/envs/torch_xla2/lib/python3.10/site-packages/torch/nn/modules/module.py", line 1736, in _wrapped_call_impl
    return self._call_impl(*args, **kwargs)
  File "/usr/local/google/home/barni/miniconda3/envs/torch_xla2/lib/python3.10/site-packages/torch/nn/modules/module.py", line 1747, in _call_impl
    return forward_call(*args, **kwargs)
  File "/usr/local/google/home/barni/miniconda3/envs/torch_xla2/lib/python3.10/site-packages/transformers/models/clip/modeling_clip.py", line 877, in forward
    hidden_states = self.embeddings(input_ids=input_ids, position_ids=position_ids)
  File "/usr/local/google/home/barni/miniconda3/envs/torch_xla2/lib/python3.10/site-packages/torch/nn/modules/module.py", line 1736, in _wrapped_call_impl
    return self._call_impl(*args, **kwargs)
  File "/usr/local/google/home/barni/miniconda3/envs/torch_xla2/lib/python3.10/site-packages/torch/nn/modules/module.py", line 1747, in _call_impl
    return forward_call(*args, **kwargs)
  File "/usr/local/google/home/barni/miniconda3/envs/torch_xla2/lib/python3.10/site-packages/transformers/models/clip/modeling_clip.py", line 227, in forward
    position_embeddings = self.position_embedding(position_ids)
  File "/usr/local/google/home/barni/miniconda3/envs/torch_xla2/lib/python3.10/site-packages/torch/nn/modules/module.py", line 1736, in _wrapped_call_impl
    return self._call_impl(*args, **kwargs)
  File "/usr/local/google/home/barni/miniconda3/envs/torch_xla2/lib/python3.10/site-packages/torch/nn/modules/module.py", line 1747, in _call_impl
    return forward_call(*args, **kwargs)
  File "/usr/local/google/home/barni/miniconda3/envs/torch_xla2/lib/python3.10/site-packages/torch/nn/modules/sparse.py", line 190, in forward
    return F.embedding(
  File "/usr/local/google/home/barni/miniconda3/envs/torch_xla2/lib/python3.10/site-packages/torch/nn/functional.py", line 2516, in embedding
    return handle_torch_function(
  File "/usr/local/google/home/barni/miniconda3/envs/torch_xla2/lib/python3.10/site-packages/torch/overrides.py", line 1718, in handle_torch_function
    result = mode.__torch_function__(public_api, types, args, kwargs)
  File "/usr/local/google/home/barni/workspace/pytorch/xla/experimental/torch_xla2/torch_xla2/tensor.py", line 215, in __torch_function__
    return func(*args, **(kwargs or {}))
  File "/usr/local/google/home/barni/miniconda3/envs/torch_xla2/lib/python3.10/site-packages/torch/nn/functional.py", line 2551, in embedding
    return torch.embedding(weight, input, padding_idx, scale_grad_by_freq, sparse)
  File "/usr/local/google/home/barni/workspace/pytorch/xla/experimental/torch_xla2/torch_xla2/tensor.py", line 230, in __torch_dispatch__
    return self.env.dispatch(func, types, args, kwargs)
  File "/usr/local/google/home/barni/workspace/pytorch/xla/experimental/torch_xla2/torch_xla2/tensor.py", line 423, in dispatch
    args, kwargs = self.t2j_iso((args, kwargs))
  File "/usr/local/google/home/barni/workspace/pytorch/xla/experimental/torch_xla2/torch_xla2/tensor.py", line 480, in t2j_iso
    return torch_pytree.tree_map_only(torch.Tensor, to_jax, torchtensors)
  File "/usr/local/google/home/barni/miniconda3/envs/torch_xla2/lib/python3.10/site-packages/torch/utils/_pytree.py", line 1136, in tree_map_only
    return tree_map(map_only(__type_or_types_or_pred)(func), tree, is_leaf=is_leaf)
  File "/usr/local/google/home/barni/miniconda3/envs/torch_xla2/lib/python3.10/site-packages/torch/utils/_pytree.py", line 964, in tree_map
    return treespec.unflatten(map(func, *flat_args))
  File "/usr/local/google/home/barni/miniconda3/envs/torch_xla2/lib/python3.10/site-packages/torch/utils/_pytree.py", line 803, in unflatten
    leaves = list(leaves)
  File "/usr/local/google/home/barni/miniconda3/envs/torch_xla2/lib/python3.10/site-packages/torch/utils/_pytree.py", line 1082, in wrapped
    return func(x)
  File "/usr/local/google/home/barni/workspace/pytorch/xla/experimental/torch_xla2/torch_xla2/tensor.py", line 478, in to_jax
    assert isinstance(x, XLATensor2), f'Expect a XLATensor2 but got {type(x)}; usually this means there is a mixed math between XLATensor and torch.Tensor'
AssertionError: Expect a XLATensor2 but got <class 'torch.Tensor'>; usually this means there is a mixed math between XLATensor and torch.Tensor
barni@barni ~/workspace/pytorch-tpu/run_torchbench

[barney-s]

With the attached fixes:

 % JAX_PLATFORMS=cpu python models/stable_diffusion_text_encoder.py   

/usr/local/google/home/barni/miniconda3/envs/diffusion-models-2/lib/python3.10/site-packages/torch/distributed/distributed_c10d.py:337: UserWarning: Device capability of jax unspecified, assuming `cpu` and `cuda`. Please specify it via the `devices` argument of `register_backend`.
  warnings.warn(
Loading pipeline components...:  50%|██████████████████████████████████████████████████████▌                                                      | 3/6 [00:00<00:00,  7.44it/s]/usr/local/google/home/barni/miniconda3/envs/diffusion-models-2/lib/python3.10/site-packages/transformers/tokenization_utils_base.py:1601: FutureWarning: `clean_up_tokenization_spaces` was not set. It will be set to `True` by default. This behavior will be depracted in transformers v4.45, and will be then set to `False` by default. For more details check this issue: https://github.com/huggingface/transformers/issues/31884
  warnings.warn(
Loading pipeline components...: 100%|█████████████████████████████████████████████████████████████████████████████████████████████████████████████| 6/6 [00:00<00:00, 11.52it/s]
WARNING:root:Duplicate op registration for aten.trunc
[XLATensor2(<class 'jaxlib.xla_extension.ArrayImpl'> [[11415   957  6492  1000 16295  1753 15015 18010 30834 11090]])]
Eager mode time 2.0050325570628047
Eager max abs vs expected,last_hidden_state tensor(1.1921e-05, grad_fn=<MaxBackward1>)
Eager max abs vs expected,pooler_output tensor(8.8215e-06, grad_fn=<MaxBackward1>)
Jitted mode time 0.946164971916005
Jitted max abs vs expected,last_hidden_state tensor(1.1921e-05, grad_fn=<MaxBackward1>)
Jitted max abs vs expected,pooler_output tensor(8.8215e-06, grad_fn=<MaxBackward1>)
barni@barni ~/workspace/pytorch-tpu/run_torchbench
 % 

[barney-s]

please assign this to me @ManfeiBai