Instructions to use royleibov/Jamba-v0.1-ZipNN-Compressed with libraries, inference providers, notebooks, and local apps. Follow these links to get started.

Libraries

How to use royleibov/Jamba-v0.1-ZipNN-Compressed with Transformers:

# Use a pipeline as a high-level helper
from transformers import pipeline

pipe = pipeline("text-generation", model="royleibov/Jamba-v0.1-ZipNN-Compressed", trust_remote_code=True)

# Load model directly
from transformers import AutoTokenizer, AutoModelForCausalLM

tokenizer = AutoTokenizer.from_pretrained("royleibov/Jamba-v0.1-ZipNN-Compressed", trust_remote_code=True)
model = AutoModelForCausalLM.from_pretrained("royleibov/Jamba-v0.1-ZipNN-Compressed", trust_remote_code=True, device_map="auto")

Notebooks
Google Colab
Kaggle
Local Apps Settings

vLLM

How to use royleibov/Jamba-v0.1-ZipNN-Compressed with vLLM:

Install from pip and serve model

# Install vLLM from pip:
pip install vllm
# Start the vLLM server:
vllm serve "royleibov/Jamba-v0.1-ZipNN-Compressed"
# Call the server using curl (OpenAI-compatible API):
curl -X POST "http://localhost:8000/v1/completions" \
	-H "Content-Type: application/json" \
	--data '{
		"model": "royleibov/Jamba-v0.1-ZipNN-Compressed",
		"prompt": "Once upon a time,",
		"max_tokens": 512,
		"temperature": 0.5
	}'

Use Docker

docker model run hf.co/royleibov/Jamba-v0.1-ZipNN-Compressed

SGLang

How to use royleibov/Jamba-v0.1-ZipNN-Compressed with SGLang:

Install from pip and serve model

# Install SGLang from pip:
pip install sglang
# Start the SGLang server:
python3 -m sglang.launch_server \
    --model-path "royleibov/Jamba-v0.1-ZipNN-Compressed" \
    --host 0.0.0.0 \
    --port 30000
# Call the server using curl (OpenAI-compatible API):
curl -X POST "http://localhost:30000/v1/completions" \
	-H "Content-Type: application/json" \
	--data '{
		"model": "royleibov/Jamba-v0.1-ZipNN-Compressed",
		"prompt": "Once upon a time,",
		"max_tokens": 512,
		"temperature": 0.5
	}'

Use Docker images

docker run --gpus all \
    --shm-size 32g \
    -p 30000:30000 \
    -v ~/.cache/huggingface:/root/.cache/huggingface \
    --env "HF_TOKEN=<secret>" \
    --ipc=host \
    lmsysorg/sglang:latest \
    python3 -m sglang.launch_server \
        --model-path "royleibov/Jamba-v0.1-ZipNN-Compressed" \
        --host 0.0.0.0 \
        --port 30000
# Call the server using curl (OpenAI-compatible API):
curl -X POST "http://localhost:30000/v1/completions" \
	-H "Content-Type: application/json" \
	--data '{
		"model": "royleibov/Jamba-v0.1-ZipNN-Compressed",
		"prompt": "Once upon a time,",
		"max_tokens": 512,
		"temperature": 0.5
	}'

Docker Model Runner
How to use royleibov/Jamba-v0.1-ZipNN-Compressed with Docker Model Runner:
```
docker model run hf.co/royleibov/Jamba-v0.1-ZipNN-Compressed
```

tomeras1

michael-go commited on Aug 20, 2024

Commit

c16867f

verified ·

1 Parent(s): ce13f3f

Update LoRA fine-tune example - more target_modules, lower LR, bf16 (#49)

Browse files

- Update LoRA fine-tune example - more target_modules, lower LR, bf16 (a69ca0f303d6079e51f4d323a81e2ec76484fc92)

Co-authored-by: Michael Gokhman <michael-go@users.noreply.huggingface.co>

Files changed (1) hide show

README.md +21 -14

README.md CHANGED Viewed

@@ -96,31 +96,40 @@ model = AutoModelForCausalLM.from_pretrained("ai21labs/Jamba-v0.1",
 </details>
 ### Fine-tuning example
-Jamba is a base model that can be fine-tuned for custom solutions (including for chat/instruct versions). You can fine-tune it using any technique of your choice. Here is an example of fine-tuning with the [PEFT](https://huggingface.co/docs/peft/index) library:
 ```python
 from datasets import load_dataset
-from trl import SFTTrainer
 from peft import LoraConfig
 from transformers import AutoTokenizer, AutoModelForCausalLM, TrainingArguments
 tokenizer = AutoTokenizer.from_pretrained("ai21labs/Jamba-v0.1")
-model = AutoModelForCausalLM.from_pretrained("ai21labs/Jamba-v0.1", device_map='auto')
 dataset = load_dataset("Abirate/english_quotes", split="train")
-training_args = TrainingArguments(
     output_dir="./results",
-    num_train_epochs=3,
     per_device_train_batch_size=4,
     logging_dir='./logs',
     logging_steps=10,
-    learning_rate=2e-3
-)
-lora_config = LoraConfig(
-    r=8,
-    target_modules=["embed_tokens", "x_proj", "in_proj", "out_proj"],
-    task_type="CAUSAL_LM",
-    bias="none"
 )
 trainer = SFTTrainer(
     model=model,
@@ -128,9 +137,7 @@ trainer = SFTTrainer(
     args=training_args,
     peft_config=lora_config,
     train_dataset=dataset,
-    dataset_text_field="quote",
 )
 trainer.train()
 ```

 </details>
 ### Fine-tuning example
+Jamba is a base model that can be fine-tuned for custom solutions (including for chat/instruct versions). You can fine-tune it using any technique of your choice. Here is an example of fine-tuning with the [PEFT](https://huggingface.co/docs/peft/index) library (requires ~120GB GPU RAM, in example 2xA100 80GB):
 ```python
+import torch
 from datasets import load_dataset
+from trl import SFTTrainer, SFTConfig
 from peft import LoraConfig
 from transformers import AutoTokenizer, AutoModelForCausalLM, TrainingArguments
 tokenizer = AutoTokenizer.from_pretrained("ai21labs/Jamba-v0.1")
+model = AutoModelForCausalLM.from_pretrained(
+    "ai21labs/Jamba-v0.1", device_map='auto', torch_dtype=torch.bfloat16)
+lora_config = LoraConfig(
+    r=8,
+    target_modules=[
+        "embed_tokens",
+        "x_proj", "in_proj", "out_proj", # mamba
+        "gate_proj", "up_proj", "down_proj", # mlp
+        "q_proj", "k_proj", "v_proj" # attention
+    ],
+    task_type="CAUSAL_LM",
+    bias="none"
+)
 dataset = load_dataset("Abirate/english_quotes", split="train")
+training_args = SFTConfig(
     output_dir="./results",
+    num_train_epochs=2,
     per_device_train_batch_size=4,
     logging_dir='./logs',
     logging_steps=10,
+    learning_rate=1e-5,
+    dataset_text_field="quote",
 )
 trainer = SFTTrainer(
     model=model,
     args=training_args,
     peft_config=lora_config,
     train_dataset=dataset,
 )
 trainer.train()
 ```