Update checkpointing directory -> using vLLM and from_pretrained

[felipemello1]

Co-authored-by: vancoyendall vancoykendall@gmail.com#### Context
What is the purpose of this PR? Is it to

• add a new feature
• fix a bug
• update tests and/or documentation
• other (please add here)

TLDR using vLLM and Huggingface .from_pretrained

1. train your model using torchtune main (nightlies). It will produce a folder like this:
   

2. Copy the contents of your latest epoch to base_model folder, which contains the checkpoint_dir original's content, without the model files (.pt, .bin, .safetensors):

cp /tmp/llama_3_2_1b/lora_single_device/epoch_2/* /tmp/llama_3_2_1b/lora_single_device/base_model

Making it look like this:

3. Now, you can use it with vLLM and Huggingface. There is one catch here: when using lora, we output MERGED weights AND the adapter. You should NOT use both at the same time. Either use ONLY the merged weights OR the BASE UNTRAINED MODEL + adapter.

3.1. Using huggingface .from_pretrained with BASE UNTRAINED MODEL + adapter

from peft import PeftModel
from transformers import AutoModelForCausalLM, AutoTokenizer

# Define the model and adapter paths
original_model_name = "meta-llama/Llama-3.2-1B-Instruct"
trained_model_path = "/tmp/torchtune/llama3_2_1B/lora_single_device/base_model"

model = AutoModelForCausalLM.from_pretrained(original_model_name)

# huggingface will look for adapter_model.safetensors and adapter_config.json
peft_model = PeftModel.from_pretrained(model, trained_model_path)

# Load the tokenizer
tokenizer = AutoTokenizer.from_pretrained(original_model_name)

# Function to generate text
def generate_text(model, tokenizer, prompt, max_length=50):
    inputs = tokenizer(prompt, return_tensors="pt")
    outputs = model.generate(**inputs, max_length=max_length)
    return tokenizer.decode(outputs[0], skip_special_tokens=True)

prompt = "Complete the sentence: 'Once upon a time...'"
print("Base model output:", generate_text(peft_model, tokenizer, prompt))

3.2. Using huggingface with FULLY TRAINED model

from transformers import AutoModelForCausalLM, AutoTokenizer

# Define the model and adapter paths
trained_model_path = "/tmp/torchtune/llama3_2_1B/full_single_device/base_model"

model = AutoModelForCausalLM.from_pretrained(
    pretrained_model_name_or_path=trained_model_path,
)

# Load the tokenizer
tokenizer = AutoTokenizer.from_pretrained(trained_model_path, safetensors=True)


# Function to generate text
def generate_text(model, tokenizer, prompt, max_length=50):
    inputs = tokenizer(prompt, return_tensors="pt")
    outputs = model.generate(**inputs, max_length=max_length)
    return tokenizer.decode(outputs[0], skip_special_tokens=True)


prompt = "Complete the sentence: 'Once upon a time...'"
print("Base model output:", generate_text(model, tokenizer, prompt))

3.3. using the MERGED TRAINED MODEL with vLLM

IMPORTANT: this will not work right away. Your output directory has 2 files:

• ft-model-00001-of-00001.safetensors
• adapter_model.safetensors

vLLM doesnt know what is a model and what is an adapter. When it tries to load the adapter, it will raise an error. Therefore, delete adapter_model.safetensors from the folder, and it will work

rm /tmp/torchtune/llama3_2_1B/lora_single_device/base_model/adapter_model.safetensors

Now you can run vLLM locally

from vllm import LLM, SamplingParams

def print_outputs(outputs):
    for output in outputs:
        prompt = output.prompt
        generated_text = output.outputs[0].text
        print(f"Prompt: {prompt!r}, Generated text: {generated_text!r}")
    print("-" * 80)

llm = LLM(
    model="/tmp/torchtune/llama3_2_1B/lora_single_device",
    load_format="safetensors",
    kv_cache_dtype="auto",
)
sampling_params = SamplingParams(max_tokens=16, temperature=0.5)

conversation = [
    {"role": "system", "content": "You are a helpful assistant"},
    {"role": "user", "content": "Hello"},
    {"role": "assistant", "content": "Hello! How can I assist you today?"},
    {
        "role": "user",
        "content": "Write an essay about the importance of higher education.",
    },
]
outputs = llm.chat(conversation, sampling_params=sampling_params, use_tqdm=False)
print_outputs(outputs)

Context

In torchtune's current state, if checkpoint_dir != outputdir, it breaks. Since the files are all mixed, saved as .pt and without the proper configs, its hard for users to readily use it with vllm/huggingface, resulting in issues such as #2048, #2025 and #2118.

This PR is NOT a major refactor. Everything is backwards compatible. The intention here is just to organize the output_dir and allow users to quickly use their models with HF and vLLM.

Changelog

1. The folder is automatically created/populated like described above

2. Initially, base_model has all the files from the checkpoint_dir, except those that end in .pt, .safetensors, .bin, etc

3. We save the original_model info in the adapter.config

4. naming got standardized

5. safetensors is the default for HF ckpt

6. Solved bugs when input_dir != output_dir

Next steps

Update docs

Unresolved issues

Left TODOs in the code, to be addressed in follow up PRs. It makes the code ugly, but we are due a refactor.

Test plan

• unit tests
• resumed pretrained
  
• ran vllm and huggingface

[pytorch-bot]

🔗 Helpful Links

🧪 See artifacts and rendered test results at hud.pytorch.org/pr/pytorch/torchtune/2074

• 📄 Preview Python docs built from this PR

Note: Links to docs will display an error until the docs builds have been completed.

✅ No Failures

As of commit 34c12ff with merge base 2b1ee6d ():
💚 Looks good so far! There are no failures yet. 💚

This comment was automatically generated by Dr. CI and updates every 15 minutes.

[felipemello1]

related: #2026

[felipemello1]

kaggle donwload

[felipemello1]

replaced it with "copy_files", so we save every file, and not only config

[felipemello1]

hf download

[felipemello1]

kaggle donwload

[felipemello1]

moved down

[felipemello1]

i don't really like this. The other options are 1) to actually fix the issue, which is knowing save_adapter_only, or 2) silently let it happen, which is dangerous if the adapter + model was trained (e.g. embeddings), and the user forgot to change the file names.

Not sure if there is a 3rd

[ebsmothers]

I'm not sure I fully follow this, but for (1) couldn't it be done by just e.g. saving save_adapter_weights_only as part of the recipe state and pulling it from there?

[felipemello1]

are you ok with saving it as part of the recipe_state? that would work.

I'm not sure I fully follow this,

1. User trains a model that has both lora + finetuning (e.g. vision model)
2. User resumes from ckpt, and forgets to update the ckpt files
3. Since we are looking at the ckpt_dir, we will get the untrained model, which is a silent bug. Looking at ckpt_dir is only safe IF save_adapter_only=True. Else, we always have to look at output_dir, which will raise "file not found" error if the user forgot to update fnames

[felipemello1]

removed this. when lora=True, it will always get it from ckpt_dir (which is our current state anyway). Need a follow up PR to address it.

[ebsmothers]

Seems like there are a few follow-ups on this PR.. can we file an issue so we can track the different todos in a single place?

[felipemello1]

no more .pt. Lets do .bin.

[joecummings]

@ebsmothers How was the PEFT model figuring this out before?

[ebsmothers]

Oh it wasn't. That's why we loaded PEFT models in two steps:

model = AutoModelForCausalLM.from_pretrained(model_id)
peft_model = PeftModel.from_pretrained(model, checkpoint_dir)

instead of

AutoModelForCausalLM.from_pretrained(checkpoint_dir)

I had a hacky version of this in #2026 but it was pointed out by @pbontrager that this shouldn't be present for models that are gonna get pushed to the hub (in that case we would want the hub model ID here, not a local path). (Edited) Looks like that is addressed here though

[joecummings]

Can you make ignore suffixes kwarg only?

[felipemello1]

I don't follow. Can you give an example?

[joecummings]

def copy_files(
	input,
	output,
	*,
	ignore_suffixes
)

[ebsmothers]

Call me crazy but do we need a full Python function reinventing recursive copy here? E.g.

os.system(
f"rsync -av --ignore-existing {" ".join([f"--exclude *{}" for ignore_suffix in ignore_suffixes])} {input_dir} {output_dir}"
)

[SalmanMohammadi]

I just use shutil.copy_tree edit: though I may be missing some of the nuances of this function

[joecummings]

Yeah wait actually is there a reason why can't use copytree? I'd rather that be maintained by core Python.

[felipemello1]

beyond the suffixes, i also have to ignore .cache and .git__

idk, the function gives us some flexibility and its readable. But it seems to be 3x1.

Can we leave it for the ckpt refactoring?

[ebsmothers]

Hmm I do like copy_tree and feel like it should be workable. But won't block on it, just include it in the follow-up task (as I mentioned above)

[ebsmothers]

remove

[ebsmothers]

Thanks for getting all of this sorted out! I know this was not an easy one. Once CI is green I think we're good to merge here. One request: can you update the summary to lead with the example usage within vLLM and PEFT? In case anyone is coming to this PR later I think that's what they will want to see. We should also add a section in the readme giving this process explicitly for better visibility.

[nouranali]

How would this method affect GPU consumption, having to load two models? While using only 1?