[fix] remove "with_suffix"

[felipemello1]

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)

1. HF requires the file "model.safetensors.index.json" to find models.
2. In torchtune i had it defined as:

SAFETENSOR_INDEX_FNAME = "model.safetensors.index"
TORCH_INDEX_FNAME = "pytorch_model.bin.index"

if safetensors:
	index_name = SAFETENSOR_INDEX_FNAME
else:
	index_name = TORCH_INDEX_FNAME

Path.join(my/path, index_name).with_suffix("json")

with_suffix actually replace .index with .json, generating the file "model.safetensors.json" instead of "model.safetensors.index.json"

Thanks @mostafaelhoushi for catching that.

To run with huggingface

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)


# 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))

Test

HF script above

Also:

tune run full_finetune_single_device --config llama3_2/1B_full_single_device epochs=3 max_steps_per_epoch=3

produces the file "model.safetensors.index.json" in every epoch

[pytorch-bot]

🔗 Helpful Links

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

• 📄 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 3e0aae6 with merge base 5370e0d ():
💚 Looks good so far! There are no failures yet. 💚

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