[BugFix] Model State Reload with Quantized Stubs in SparseAutoModelForCausalLM

[rahul-tuli]

Description

Identified a bug in the main branch where the model state fails to reload when using quantized stubs with SparseAutoModelForCausalLM.from_pretrained(...). The issue was due to the reload_model_state method expecting weight files in a local directory, not accounting for remotely hosted model directories.

Solution

Propose downloading the model directory before invoking reload_model_state to ensure weight files are available locally for model state reload.

Testing

Tested with the following script, confirming the fix resolves the issue:

from sparseml.transformers import SparseAutoModelForCausalLM

model_path = "mgoin/llama2.c-stories15M-quant-pt"
m1 = SparseAutoModelForCausalLM.from_pretrained(model_path)

Observations

Before the Fix:

Model state fails to reload due to missing local weight files, as shown in warnings and errors in the logs.

..
..
2024-04-05 17:33:33 sparseml.core.recipe.recipe INFO     Loading recipe from file /home/rahul/.cache/huggingface/hub/models--mgoin--llama2.c-stories15M-quant-pt/snapshots/aa70fc9dc46615b68f935fb5405ae7875b88b716/recipe.yaml
manager stage: Model structure initialized
2024-04-05 17:33:34 sparseml.pytorch.model_load.helpers INFO     Applied an unstaged recipe to the model at mgoin/llama2.c-stories15M-quant-pt
2024-04-05 17:33:34 sparseml.pytorch.model_load.helpers WARNING  Model state was not reloaded for SparseML: could not find model weights for mgoin/llama2.c-stories15M-quant-pt

After the Fix:

..
..
..
2024-04-05 21:17:16 sparseml.pytorch.model_load.helpers INFO     Reloaded model state after SparseML recipe structure modifications from /nm/drive0/rahul/.cache/huggingface/hub/models--mgoin--llama2.c-stories15M-quant-pt/snapshots/aa70fc9dc46615b68f935fb5405ae7875b88b716

Successfully reloaded model state with the fix.

---

• To see the specific tasks where the Asana app for GitHub is being used, see below:
  • https://app.asana.com/0/0/1204866054040952

[mgoin]

The snapshot_download can pull down a lot of extra files since it downloads the whole folder. This might interact weirdly with the resolve_recipe call with specifically tries to download the recipe. This works, but I would like to be a bit more selective with the download

[bfineran]

LGTM pending @mgoin's comment

[rahul-tuli]

The snapshot_download can pull down a lot of extra files since it downloads the whole folder. This might interact weirdly with the resolve_recipe call with specifically tries to download the recipe. This works, but I would like to be a bit more selective with the download

Addressed in latest commit @mgoin

[mgoin]

Thanks a lot! It's additional complexity but I think good to have

[rahul-tuli]

Thanks a lot! It's additional complexity but I think good to have

It was a great callout. Really appreciate it.