torch.jit.trace can trace shared weights, no need to clone weights when tracing
Author: LSC527
Pull request: no
Labels: name: bug

### System Info

[source code of "tie or clone weights"](https://github.com/huggingface/transformers/blob/8a61fe023430115bb61ec328a29d35571f4fc2c4/src/transformers/modeling_utils.py#L1137)
[document](https://huggingface.co/docs/transformers/v4.20.1/en/serialization#torchscript-flag-and-tied-weights)
I did a experiment and results showed that `torch.jit.trace` can trace shared weights and use the `TorchScript` for training. Correct me if I was wrong, thx!

### Who can help?

_No response_

### Information

- [ ] The official example scripts
- [ ] My own modified scripts

### Tasks

- [ ] An officially supported task in the `examples` folder (such as GLUE/SQuAD, ...)
- [ ] My own task or dataset (give details below)

### Reproduction

```
import torch
import torch.nn as nn

batch_size = 32
seq_len = 32
emb_size = 128
vocab_size = 32

class Model(nn.Module):
    def __init__(self):
        super().__init__()
        self.emb1 = nn.Embedding(seq_len, emb_size)
        self.emb2 = nn.Embedding(seq_len, emb_size)
            
    def forward(self, x):
        y1 = self.emb1(x)
        y2 = self.emb2(x)
        return (y1, y2)

model = Model()
model.emb2.weight = model.emb1.weight
model.eval()

with torch.no_grad():
    example_input = torch.randint(vocab_size, [batch_size, seq_len])
    example_output = model(example_input)
    s = example_output[0].size()
    weight_before_train = model.emb1.weight.clone()
    print(weight_before_train)

# origin model
model.train()
loss_fn = nn.L1Loss()
optimizer = torch.optim.SGD(model.parameters(), 0.1)
for _ in range(100):
    optimizer.zero_grad()
    inputs = torch.randint(vocab_size, [batch_size, seq_len])
    targets = (torch.randn(s), torch.randn(s))
    outputs = model(inputs)
    assert torch.allclose(outputs[0], outputs[1])
    loss = loss_fn(targets[0], outputs[0]) + loss_fn(targets[1], outputs[1])
    loss.backward()
    optimizer.step()
    
model.eval()
with torch.no_grad():
    weight_after_train = model.emb1.weight.clone()
    print(weight_after_train)
    
assert torch.equal(model.emb1.weight, model.emb2.weight)
assert not torch.allclose(weight_before_train, weight_after_train)
    

# traced model
traced = torch.jit.trace(model, example_input)
traced.eval()
with torch.no_grad():
    weight_before_train = traced.emb1.weight.clone()
    print(weight_before_train)
traced.train()
loss_fn = nn.L1Loss()
optimizer = torch.optim.SGD(traced.parameters(), 0.1)

for _ in range(100):
    optimizer.zero_grad()
    inputs = torch.randint(vocab_size, [batch_size, seq_len])
    targets = (torch.randn(s), torch.randn(s))
    outputs = traced(inputs)
    assert torch.allclose(outputs[0], outputs[1])
    loss = loss_fn(targets[0], outputs[0]) + loss_fn(targets[1], outputs[1])
    loss.backward()
    optimizer.step()
    
traced.eval()
with torch.no_grad():
    weight_after_train = traced.emb1.weight.clone()
    print(weight_after_train)
    
assert torch.equal(traced.emb1.weight, traced.emb2.weight)
assert not torch.allclose(weight_before_train, weight_after_train)
```

### Expected behavior

shared weights can be traced

Comments:
This issue has been automatically marked as stale because it has not had recent activity. If you think this still needs to be addressed please comment on this thread.

Please note that issues that do not follow the [contributing guidelines](https://github.com/huggingface/transformers/blob/main/CONTRIBUTING.md) are likely to be ignored.