Unable to load a model trained via FSDP
Hello,I have trained a model using FSDP. More specifically, I used thisrun_clm.pyscript, with the option--fsdp "shard_grad_op auto_wrap".The training went fine and model was saved. However, while trying to load the model I get error:Loading checkpoint shards: 100%|███████████████████████████████████████| 3/3 [00:00<00:00, 19.55it/s]
Traceback (most recent call last):
  File "/home/tarun/memory-llm-paper/run_experiments_pythia.py", line 49, in <module>
    model = AutoModelForCausalLM.from_pretrained('/data/users/tarun/coref/models/output32/checkpoint-345')
  File "/home/tarun/miniconda3/envs/coref-1/lib/python3.9/site-packages/transformers/models/auto/auto_factory.py", line 564, in from_pretrained
    return model_class.from_pretrained(
  File "/home/tarun/miniconda3/envs/coref-1/lib/python3.9/site-packages/transformers/modeling_utils.py", line 4014, in from_pretrained
    ) = cls._load_pretrained_model(
  File "/home/tarun/miniconda3/envs/coref-1/lib/python3.9/site-packages/transformers/modeling_utils.py", line 4559, in _load_pretrained_model
    raise RuntimeError(f"Error(s) in loading state_dict for {model.__class__.__name__}:\n\t{error_msg}")
RuntimeError: Error(s) in loading state_dict for GPTNeoXForCausalLM:
        size mismatch for gpt_neox.embed_in.weight: copying a param with shape torch.Size([128778240]) from checkpoint, the shape in current model is torch.Size([50304, 2560]).
        size mismatch for gpt_neox.final_layer_norm.bias: copying a param with shape torch.Size([0]) from checkpoint, the shape in current model is torch.Size([2560]).
        size mismatch for embed_out.weight: copying a param with shape torch.Size([0]) from checkpoint, the shape in current model is torch.Size([50304, 2560]).
        You may consider adding `ignore_mismatched_sizes=True` in the model `from_pretrained` method.Now, if we specifically look at this line:size mismatch for gpt_neox.embed_in.weight: copying a param with shape torch.Size([128778240]) from checkpoint, the shape in current model is torch.Size([50304, 2560]).we can notice that 50304*2560 = 128778240. So it seems that at the end of FSDP training, the model’s params were stored in a flattened array. While trying to load it usingfrom_pretrainedthe library isn’t able to unflatten it automatically.This is my conda environment details if that helps:Conda env detailsname: coref-1channels:pytorchplotlynvidiaconda-forgedefaultsdependencies:_libgcc_mutex=0.1=conda_forge_openmp_mutex=4.5=2_gnu_sysroot_linux-64_curr_repodata_hack=3=haa98f57_10asttokens=2.4.1=pyhd8ed1ab_0attrs=23.1.0=py39h06a4308_0binutils_impl_linux-64=2.38=h2a08ee3_1binutils_linux-64=2.38.0=hc2dff05_0blas=1.0=openblasbzip2=1.0.8=h4bc722e_7c-ares=1.33.1=heb4867d_0ca-certificates=2024.7.4=hbcca054_0cmake=3.30.2=hf8c4bd3_0comm=0.2.2=pyhd8ed1ab_0cuda-cudart=12.4.127=0cuda-cupti=12.4.127=0cuda-libraries=12.4.0=0cuda-nvrtc=12.4.127=0cuda-nvtx=12.4.127=0cuda-opencl=12.4.127=0cuda-runtime=12.4.0=0cuda-version=11.8=hcce14f8_3cudatoolkit=11.8.0=h6a678d5_0cudnn=8.9.2.26=cuda11_0cupti=11.8.0=he078b1a_0debugpy=1.8.5=py39h98e3656_0decorator=5.1.1=pyhd8ed1ab_0exceptiongroup=1.2.2=pyhd8ed1ab_0executing=2.0.1=pyhd8ed1ab_0filelock=3.13.1=py39h06a4308_0fsspec=2024.6.1=py39h06a4308_0gcc_impl_linux-64=11.2.0=h1234567_1gcc_linux-64=11.2.0=h5c386dc_0gh=2.55.0=h76a2195_0gmp=6.2.1=h295c915_3gmpy2=2.1.2=py39heeb90bb_0gxx_impl_linux-64=11.2.0=h1234567_1gxx_linux-64=11.2.0=hc2dff05_0importlib-metadata=8.4.0=pyha770c72_0importlib_metadata=8.4.0=hd8ed1ab_0intel-openmp=2022.0.1=h06a4308_3633ipykernel=6.29.5=pyh3099207_0ipython=8.18.1=pyh707e725_3ipywidgets=8.1.5=pyhd8ed1ab_0jedi=0.19.1=pyhd8ed1ab_0jinja2=3.1.4=py39h06a4308_0jsonschema-specifications=2023.7.1=py39h06a4308_0jupyter_client=8.6.2=pyhd8ed1ab_0jupyter_core=5.7.2=py39hf3d152e_0jupyterlab_widgets=3.0.13=pyhd8ed1ab_0kernel-headers_linux-64=3.10.0=h57e8cba_10keyutils=1.6.1=h166bdaf_0krb5=1.21.3=h659f571_0ld_impl_linux-64=2.38=h1181459_1libabseil=20240116.2=cxx17_h6a678d5_0libblas=3.9.0=16_linux64_openblaslibcublas=12.4.2.65=0libcufft=11.2.0.44=0libcufile=1.9.1.3=0libcurand=10.3.5.147=0libcurl=8.9.1=hdb1bdb2_0libcusolver=11.6.0.99=0libcusparse=12.3.0.142=0libedit=3.1.20191231=he28a2e2_2libev=4.33=hd590300_2libexpat=2.6.2=h59595ed_0libffi=3.4.4=h6a678d5_1libgcc-devel_linux-64=11.2.0=h1234567_1libgcc-ng=14.1.0=h77fa898_0libgfortran-ng=11.2.0=h00389a5_1libgfortran5=11.2.0=h1234567_1libgomp=14.1.0=h77fa898_0liblapack=3.9.0=16_linux64_openblaslibmagma=2.8.0=hfdb99dd_0libmagma_sparse=2.8.0=h9ddd185_0libnghttp2=1.58.0=h47da74e_1libnpp=12.2.5.2=0libnsl=2.0.1=hd590300_0libnvfatbin=12.4.127=0libnvjitlink=12.4.99=0libnvjpeg=12.3.1.89=0libopenblas=0.3.21=h043d6bf_0libprotobuf=4.25.3=h08a7969_0libsodium=1.0.18=h36c2ea0_1libsqlite=3.46.0=hde9e2c9_0libssh2=1.11.0=h0841786_0libstdcxx-devel_linux-64=11.2.0=h1234567_1libstdcxx-ng=14.1.0=hc0a3c3a_0libuuid=2.38.1=h0b41bf4_0libuv=1.48.0=hd590300_0libxcrypt=4.4.36=hd590300_1libzlib=1.3.1=h4ab18f5_1llvm-openmp=15.0.7=h0cdce71_0magma=2.8.0=h4aca40b_0markupsafe=2.1.3=py39h5eee18b_0matplotlib-inline=0.1.7=pyhd8ed1ab_0mkl=2022.1.0=hc2b9512_224mpc=1.1.0=h10f8cd9_1mpfr=4.0.2=hb69a4c5_1mpmath=1.3.0=py39h06a4308_0nbformat=5.9.2=py39h06a4308_0ncurses=6.5=he02047a_1nest-asyncio=1.6.0=pyhd8ed1ab_0networkx=3.2.1=py39h06a4308_0numpy=1.26.4=py39heeff2f4_0numpy-base=1.26.4=py39h8a23956_0openssl=3.3.1=hb9d3cd8_3packaging=24.1=pyhd8ed1ab_0parso=0.8.4=pyhd8ed1ab_0pexpect=4.9.0=pyhd8ed1ab_0pickleshare=0.7.5=py_1003pip=24.2=py39h06a4308_0platformdirs=4.2.2=pyhd8ed1ab_0plotly=5.23.0=py_0prompt-toolkit=3.0.47=pyha770c72_0psutil=6.0.0=py39hd3abc70_0ptyprocess=0.7.0=pyhd3deb0d_0pure_eval=0.2.3=pyhd8ed1ab_0pybind11-abi=4=hd3eb1b0_1pygments=2.18.0=pyhd8ed1ab_0python=3.9.19=h0755675_0_cpythonpython-dateutil=2.9.0=pyhd8ed1ab_0python-fastjsonschema=2.16.2=py39h06a4308_0python_abi=3.9=5_cp39pytorch=2.3.0=gpu_cuda118py39h796af20_101pytorch-cuda=12.4=hc786d27_6pytorch-mutex=1.0=cudapyyaml=6.0.1=py39h5eee18b_0pyzmq=26.2.0=py39h4e4fb57_0readline=8.2=h5eee18b_0rhash=1.4.4=hd590300_0rpds-py=0.10.6=py39hb02cf49_0scipy=1.13.1=py39heeff2f4_0setuptools=72.1.0=py39h06a4308_0six=1.16.0=pyh6c4a22f_0sqlite=3.46.0=h6d4b2fc_0stack_data=0.6.2=pyhd8ed1ab_0sympy=1.12=py39h06a4308_0sysroot_linux-64=2.17=h57e8cba_10tenacity=8.2.3=py39h06a4308_0tk=8.6.13=noxft_h4845f30_101torchtriton=2.3.0=cuda123py39hdb19cb5_0tornado=6.4.1=py39hd3abc70_0traitlets=5.14.3=pyhd8ed1ab_0typing_extensions=4.11.0=py39h06a4308_0wcwidth=0.2.13=pyhd8ed1ab_0wheel=0.43.0=py39h06a4308_0widgetsnbextension=4.0.13=pyhd8ed1ab_0xz=5.4.6=h5eee18b_1yaml=0.2.5=h7b6447c_0zeromq=4.3.5=h75354e8_4zipp=3.20.1=pyhd8ed1ab_0zlib=1.3.1=h4ab18f5_1zstd=1.5.6=ha6fb4c9_0pip:accelerate==0.34.2aiohappyeyeballs==2.4.2aiohttp==3.10.8aiosignal==1.3.1async-timeout==4.0.3certifi==2024.8.30charset-normalizer==3.3.2contourpy==1.3.0cycler==0.12.1datasets==3.0.1dill==0.3.8evaluate==0.4.3fonttools==4.54.1frozenlist==1.4.1huggingface-hub==0.25.1idna==3.10importlib-resources==6.4.5joblib==1.4.2kiwisolver==1.4.7matplotlib==3.9.2multidict==6.1.0multiprocess==0.70.16pandas==2.2.3peft==0.13.0pillow==10.4.0protobuf==5.28.2pyarrow==17.0.0pyparsing==3.1.4pytz==2024.2regex==2024.9.11requests==2.32.3safetensors==0.4.5scikit-learn==1.5.2sentencepiece==0.2.0threadpoolctl==3.5.0tokenizers==0.20.0tqdm==4.66.5transformers==4.46.0.dev0tzdata==2024.2urllib3==2.2.3xxhash==3.5.0yarl==1.13.1prefix: /home/tarun/miniconda3/envs/coref-1Does anyone have any idea what could be causing this? I didn’t find much on surfing the internet.
2024-10-08T10:22:21Z
This is about all I could find.stackoverflow.comSize mismatch for embed_out.weight: copying a param with shape torch.Size([0]) from checkpoint - Huggingface PyTorchpytorch, huggingface-transformers, large-language-model, huggingfaceasked byMasthanon06:29PM - 05 Jul 24 UTChuggingface.codatabricks/dolly-v2-3b · PEFT/LoRA error when merging base and adapterI tried this example :
2024-10-12T12:09:24Z
I couldn’t figure out the underlying reason for the bug.But downgrading my torch version to below helped me get around this error:conda install pytorch==2.2.2 pytorch-cuda=12.1 -c pytorch -c nvidiaPreviously, I had torch 2.3.0 and cuda 12.4 (as can be seen in the conda environment details mentioned in original question.)
2024-10-12T12:52:39Z
PyTorch Forums – 12 May 19Size mismatch while loading state dictionaryvisionHi  While loading a pre-trained model, I am getting the same runtime error for every layer.  E.g. size mismatch for conv_batchnorm.weight: copying a param with shape torch.Size([1, 32]) from checkpoint, the shape in current model is...It seems to be resolved, but I thought this must be a torch bug, so I did a search and it seems to be normal on the torch forum…Well, it seems that’s just the way it is.