Add Arctic modeling/config/tokenization files

#3
by tomaarsen HF staff - opened

Hello!

This PR should add trust_remote_code=True support for pure transformers, without having to rely on a fork.

Requirements

deepspeed>=0.14.2
transformers
huggingface_hub
hf_transfer

Usage

import os
# enable hf_transfer for faster ckpt download
os.environ["HF_HUB_ENABLE_HF_TRANSFER"] = "1"

from dataclasses import dataclass
import torch
from transformers import AutoModelForCausalLM, AutoTokenizer

tokenizer = AutoTokenizer.from_pretrained(
    "Snowflake/snowflake-arctic-instruct",
    trust_remote_code=True,
    revision="refs/pr/3",
)

@dataclass
class ArcticQuantizationConfig:
    q_bits: int = 8
    rounding: str = "nearest"
    mantissa_bits: int = 3
    group_size: int = 512

quant_config = ArcticQuantizationConfig(q_bits=8)

model = AutoModelForCausalLM.from_pretrained(
    "Snowflake/snowflake-arctic-instruct",
    trust_remote_code=True,
    low_cpu_mem_usage=True,
    device_map="auto",
    ds_quantization_config=quant_config,
    max_memory={i: "150GiB" for i in range(8)},
    torch_dtype=torch.bfloat16,
    revision="refs/pr/3",
)

messages = [{"role": "user", "content": "What is 1 + 1 "}]
input_ids = tokenizer.apply_chat_template(messages, add_generation_prompt=True, return_tensors="pt").to("cuda")

outputs = model.generate(input_ids=input_ids, max_new_tokens=20)
print(tokenizer.decode(outputs[0]))

This is still untested. Also note that it specifically pulls from the refs/pr/3 branch, i.e. this PR. On my local device, this is able to correctly pull the tokenizer - I haven't tried pulling the model yet as there's no 8xH100s available on AWS for me right now. Instead, I can run the code up until the model downloads to verify that 1) the tokenizer works, 2) the model complains if deepspeed and flash_attn are not installed, 3) the download starts correctly: https://gist.github.com/tomaarsen/3bfaa73c7876dce2b96b4590a4ce2ddc

  • Tom Aarsen

Hello @tomaarsen I tried adding the trust_remote_code=True, but it still doesn't work for me.
Attached is the screenshot.
Screenshot 2024-04-26 at 00.12.18.png

Hello!

I think your transformers version might be too old. You can try upgrading it: pip install -U transformers.

  • Tom Aarsen
Snowflake org

Thank you @tomaarsen for all your help here!

jeffra changed pull request status to merged

Hello @tomaarsen @jeffra
I've updated to transformers==4.40.2.
This is the tokenizer code :

AutoTokenizer.from_pretrained("snowflake/snowflake-arctic-instruct", trust_remote_code=True)

I still get the same error as @neha-ori
Error:
File "/home/appuser/.cache/huggingface/modules/transformers_modules/snowflake/snowflake-arctic-instruct/fd069fdeadfd41eac22caac59261372ee58a63ee/tokenization_arctic.py", line 5, in
from transformers.models.llama import LlamaTokenizer

EDIT 1: Using the AutoTokenizer.from_pretrained('Snowflake/snowflake-arctic-embed-l') worked. Shouldn't the tokenizer be a text-embedding model instead of the "snowflake-arctic-instruct"?

Hey @nishantguvvada ,
Just go to your tokenization_arctic.py config file and change:
transformers.models.llama import LlamaTokenizer
to
from transformers import LlamaTokenizer

Should work for you.

Sign up or log in to comment