fsdp Error report

[liyongqi67]

Thanks for this wonderful project. I used the following script to train the model.

torchrun --nnodes=1 --nproc_per_node=2 /home/share/yongqi/project/AutoregressiveImageRetrieval/code/open_flamingo/open_flamingo/train/finetuning.py \
  --lm_path anas-awadalla/mpt-1b-redpajama-200b \
  --tokenizer_path anas-awadalla/mpt-1b-redpajama-200b \
  --cross_attn_every_n_layers 1 \
  --dataset_resampled \
  --batch_size_mmc4 1 \
  --train_num_samples_mmc4 150000\
  --workers=2 \
  --run_name OpenFlamingo-3B-vitl-mpt1b \
  --num_epochs 2 \
  --warmup_steps  2000 \
  --mmc4_textsim_threshold 0.01 \
  --laion_shards  \
  --mmc4_shards \
  --logging_steps 1 \
  --mmc4_max_num_images 1 \
  --precision fp16 \
  --fsdp

However, if I set the fsdp flag, it will report an error as follows:

AttributeError: 'MosaicGPT' object has no attribute 'set_output_embeddings'

location: flamingo.py line 294.

If I remove this flag, there is no error. Do you have any idea about this?

[i-gao]

Ah, sorry about that! The issue is from this line of the FSDP wrapping function. The MPT models are still missing some standard HF Transformers functions.

Would it work for your use case to comment out the aforementioned line from our codebase? The output embedding weight will then not be sharded. Alternatively, we can add a hack to get around this similar to this part of the code for MPT-1B.

[liyongqi67]

//Thanks for your quick reply!

1. If I only comment out the corresponding line "self.lang_encoder.set_output_embeddings( wrap(wrap(self.lang_encoder.get_output_embeddings())) )", it will report another error:

  File "/home/yongqi/miniconda3/envs/openflamingo/lib/python3.9/site-packages/torch/nn/modules/module.py", line 1501, in _call_impl
    loss_mmc4 = model(
  File "/home/yongqi/miniconda3/envs/openflamingo/lib/python3.9/site-packages/torch/nn/modules/module.py", line 1501, in _call_impl
    return forward_call(*args, **kwargs)
  File "/home/share/yongqi/project/AutoregressiveImageRetrieval/code/open_flamingo/open_flamingo/src/flamingo.py", line 111, in forward
    return forward_call(*args, **kwargs)
  File "/home/share/yongqi/project/AutoregressiveImageRetrieval/code/open_flamingo/open_flamingo/src/flamingo.py", line 111, in forward
    output = self.lang_encoder(
  File "/home/yongqi/miniconda3/envs/openflamingo/lib/python3.9/site-packages/torch/nn/modules/module.py", line 1501, in _call_impl
    output = self.lang_encoder(
  File "/home/yongqi/miniconda3/envs/openflamingo/lib/python3.9/site-packages/torch/nn/modules/module.py", line 1501, in _call_impl
    return forward_call(*args, **kwargs)
  File "/home/share/yongqi/project/AutoregressiveImageRetrieval/code/open_flamingo/open_flamingo/src/flamingo_lm.py", line 157, in forward
    return forward_call(*args, **kwargs)
  File "/home/share/yongqi/project/AutoregressiveImageRetrieval/code/open_flamingo/open_flamingo/src/flamingo_lm.py", line 157, in forward
    return super().forward(**kwargs)  # Call the other parent's forward method
    return super().forward(**kwargs)  # Call the other parent's forward method  File "/home/yongqi/.cache/huggingface/modules/transformers_modules/anas-awadalla/mpt-1b-redpajama-200b/bfa38d4f431e091fe599d7b4cdb62972532f3c7c/mosaic_gpt.py", line 366, in forward

  File "/home/yongqi/.cache/huggingface/modules/transformers_modules/anas-awadalla/mpt-1b-redpajama-200b/bfa38d4f431e091fe599d7b4cdb62972532f3c7c/mosaic_gpt.py", line 366, in forward
        logits = F.linear(x, self.transformer.wte.weight, None)
RuntimeErrorRuntimeError: : size mismatch, got 15, 15x2048,51486720

It is very strange.
2. If I want to set the set_output_embeddings() function in open_flamingo. How to do it?
I notice there is a corresponding code in https://huggingface.co/mosaicml/mpt-7b/blob/main/modeling_mpt.py. I copy this code into open_flamingosrc/factory.py?

def get_input_embeddings(self):
  return self.transformer.wte
def set_input_embeddings(self, value):
  self.transformer.wte = value
def get_output_embeddings(self):
   return self.transformer.wte
def set_output_embeddings(self, new_embeddings):
  self.transformer.wte = new_embeddings

But it still reports the above mismatch error.
Could you update how to set the get_output_embeddings() and set_output_embeddings()?

[hungvo304ml]

Have you solved the issue? I have the same problem when training with fsdp.