Shared weights are duplicated in checkpoints

[fansiawang]

🐛 Bug

When I trained a simple transformer model by setting share_all_embeddings=True, share_decoder_input_output_embed=True, I got a same model size by setting these two parameters to false. In addition, if I train the model on CPU, then I will get different model size by setting these two parameters to false, the model size of shared one maybe half of the not shared one. I think, sharing embedding not only to reduce the parameter number, but also to reduce the model size.

I try the previous fairseq(0.8.0), I can get different model size by setting these two parameters to different value. I don't quite understand why this implementation has to be modified.

The reason why I got the same size whether sharing embedding or not is that the GPU version will change the address of embedding tensor after convert all state to CPU.

I print the address of embedding tensor in fairseq/checkpoint_utils.py.

280    print("************* before convert all state to CPU *********")
281     for k, v in model_state_dict.items():
282     ¦   if "encoder.embed_tokens.weight" in k or "decoder.output_projection.weight" in k:
283     ¦   ¦   print(k, "#########", v.data_ptr())
284      
285     # convert all state to CPU
286     state_dict = utils.move_to_cpu(state_dict)
287      
288     print("************* after convert all state to CPU *********")
289     for k, v in state_dict["model"].items():
290     ¦   if "encoder.embed_tokens.weight" in k or "decoder.output_projection.weight" in k:
291     ¦   ¦   print(k, "#########", v.data_ptr())

The result is:

************* before convert all state to CPU *********
encoder.embed_tokens.weight ######### 139854906130432
decoder.output_projection.weight ######### 139854906130432

************* after convert all state to CPU *********
encoder.embed_tokens.weight ######### 139849528733760
decoder.output_projection.weight ######### 139840124805184

The address changed after coverting all state to CPU.

To Reproduce

• training script:

  
  data_dir=data/small
  model_dir=model/small
  seed=456
  max_tokens=1000

python train.py $data_dir \ --dropout 0.1 \ --clip-norm 0.1 \ --max-tokens $max_tokens \ --seed $seed \ --num-workers 4 \ --optimizer adafactor \ --criterion label_smoothed_cross_entropy \ --label-smoothing 0.1 \ --weight-decay 0.0 \ --lr 0.0003 \ --lr-scheduler inverse_sqrt \ --warmup-init-lr 1e-07 \ --warmup-updates 4000 \ --arch transformer \ --save-dir $model_dir \ --update-freq 8 \ --me 100 \ --log-interval 1 \ --save-interval-updates 1 \ --keep-interval-updates 1 \ --no-progress-bar \ --empty-cache-freq 64 \ --ddp-backend=no_c10d

- train on CPU or GPU by setting the following envirment :
`CPU: export CUDA_VISIBLE_DEVICES=""`
`GPU: export CUDA_VISIBLE_DEVICES="0"`

- set the shared parameter in `fairseq/models/transformer.py`

args.share_decoder_input_output_embed = getattr(args, "share_decoder_input_output_embed", True) args.share_all_embeddings = getattr(args, "share_all_embeddings", True)

### Expected behavior

CPU shared model: 659M
CPU not shared model: 1.6G
GPU shared model: 1.6G
GPU not shared model: 1.6G

But I think the GPU shared model size should be 659M.

### Environment

 - fairseq Version (e.g., 1.0 or master): lastest(master, 2020.5.12)
 - PyTorch Version (e.g., 1.0): 1.4.0
 - OS (e.g., Linux): CentOS 7.2
 - How you installed fairseq (`pip`, source):
 - Build command you used (if compiling from source): pip install --editable .
 - Python version: 3.7.3
 - CUDA/cuDNN version: 10.0
 - GPU models and configuration: 2080
 - Any other relevant information:

[villmow]

Does this bug cause all Transformers trained with --share-all-embeddings to currently not share embeddings after initialization?

[fansiawang]

When you train with --share-all-embeddings, the model will share all embeddings during training on GPU. But when saving checkpoint, it will save three same parameters into the model. I think it is better to save only one parameter because we share the embedding.