license: apache-2.0
language:

• en
  dataset:
• chargoddard/Open-Platypus-Chat
  tags:
• axolotl
  base_model: ai21labs/Jamba-v0.1

Jambatypus-v0.1

该模型是基于ai21labs/Jamba-v0.1在chargoddard/Open-Platypus-Chat数据集上通过QLoRA微调的版本。

训练使用2块A100 80GB显卡，基于我的LazyAxolotl - Jamba笔记本完成。

本仓库包含FP16精度的适配器和合并后模型。

推荐使用ChatML模板来调用该模型。

base_model: ai21labs/Jamba-v0.1  
trust_remote_code: true  

load_in_8bit: false  
load_in_4bit: true  
strict: false  

datasets:  
  - path: chargoddard/Open-Platypus-Chat  
    type: sharegpt  
chat_template: chatml  
dataset_prepared_path:  
val_set_size: 0.01  
output_dir: ./out  

sequence_len: 4096  
sample_packing: true  
pad_to_sequence_len: true  
eval_sample_packing: false  

use_wandb: true  
wandb_project: axolotl  
wandb_entity:  
wandb_watch:  
wandb_name: Jambatypus-v0.1  
wandb_log_model:  

adapter: qlora  
lora_r: 16  
lora_alpha: 32  
lora_dropout: 0.05  
lora_target_linear: true  

low_cpu_mem_usage: true  
gradient_accumulation_steps: 8  
micro_batch_size: 1  
num_epochs: 1  
optimizer: adamw_bnb_8bit  
adam_beta2: 0.95  
adam_epsilon: 0.00001  
max_grad_norm: 1.0  
lr_scheduler: cosine  
learning_rate: 0.0002  

train_on_inputs: false  
group_by_length: false  
bf16: auto  
fp16:  
tf32: false  

gradient_checkpointing: true  
gradient_checkpointing_kwargs:  
  use_reentrant: false  
early_stopping_patience:  
resume_from_checkpoint:  
local_rank:  
logging_steps: 1  
xformers_attention:  
flash_attention: true  

warmup_steps: 10  
evals_per_epoch: 4  
saves_per_epoch: 4  
save_total_limit: 2  
debug:  
deepspeed:  
weight_decay: 0.0  
special_tokens:  

训练超参数

训练过程中使用的超参数如下：

• 学习率: 0.0002
• 训练批次大小: 1
• 评估批次大小: 1
• 随机种子: 42
• 分布式类型: 多GPU
• 设备数量: 2
• 梯度累积步数: 8
• 总训练批次大小: 16
• 总评估批次大小: 2
• 优化器: Adam，参数为betas=(0.9,0.95)和epsilon=1e-05
• 学习率调度器类型: cosine
• 学习率预热步数: 10
• 训练轮数: 1

训练结果

框架版本

• PEFT 0.10.0
• Transformers 4.40.0.dev0
• Pytorch 2.1.2+cu118
• Datasets 2.18.0
• Tokenizers 0.15.0

💻 使用方式

以下代码创建了一个与Jambatypus交互的Gradio聊天界面。

!pip install -qqq -U git+https://github.com/huggingface/transformers  
!pip install -qqq mamba-ssm causal-conv1d>=1.2.0  
!pip install -qqq accelerate bitsandbytes torch datasets peft gradio  
!pip install -qqq flash-attn --no-build-isolation  


import torch  
import gradio as gr  
from threading import Thread  
from peft import PeftModel, PeftConfig  
from transformers import AutoModelForCausalLM, AutoTokenizer, BitsAndBytesConfig, TextIteratorStreamer  

STOP_TOKEN = "<|im_end|>"  

def predict(message, history, system_prompt, temperature, max_new_tokens, top_k, repetition_penalty, top_p):  
    # 使用给定的聊天模板格式化历史记录  
    stop_token = "<|im_end|>"  
    instruction = '<|im_start|>system\n' + system_prompt + '\n<|im_end|>\n'  
    for human, assistant in history:  
        instruction += '<|im_start|>user\n' + human + '\n<|im_end|>\n<|im_start|>assistant\n' + assistant  
    instruction += '\n<|im_start|>user\n' + message + '\n<|im_end|>\n<|im_start|>assistant\n'  
    
    streamer = TextIteratorStreamer(tokenizer, skip_prompt=True, skip_special_tokens=True)  
    enc = tokenizer([instruction], return_tensors="pt", padding=True, truncation=True)  
    input_ids, attention_mask = enc.input_ids, enc.attention_mask  

    generate_kwargs = dict(  
        {"input_ids": input_ids.to(device), "attention_mask": attention_mask.to(device)},  
        streamer=streamer,  
        do_sample=True,  
        temperature=temperature,  
        max_new_tokens=max_new_tokens,  
        top_k=top_k,  
        repetition_penalty=repetition_penalty,  
        top_p=top_p  
    )  
    t = Thread(target=model.generate, kwargs=generate_kwargs)  
    t.start()  
    outputs = []  
    for new_token in streamer:  
        if STOP_TOKEN in new_token:  
            outputs.append(new_token[:-len(stop_token)-1])  
            yield "".join(outputs)  
            break  
        outputs.append(new_token)  
        yield "".join(outputs)  


# 加载模型  
device = torch.device('cuda' if torch.cuda.is_available() else 'cpu')  
tokenizer = AutoTokenizer.from_pretrained("ai21labs/Jamba-v0.1")  

# 4位精度量化配置  
quantization_config = BitsAndBytesConfig(  
    load_in_8bit=True,  
    llm_int8_skip_modules=["mamba"]  
)  
# 加载模型和ChatML格式的tokenizer  
model = AutoModelForCausalLM.from_pretrained(  
    "ai21labs/Jamba-v0.1",  
    trust_remote_code=True,  
    torch_dtype=torch.cuda.is_bf16_supported() and torch.bfloat16 or torch.float16,  
    attn_implementation="flash_attention_2",  
    low_cpu_mem_usage=True,  
    quantization_config=quantization_config  
)  
config = PeftConfig.from_pretrained("mlabonne/Jambatypus-v0.1")  
model = PeftModel.from_pretrained(model, "mlabonne/Jambatypus-v0.1")  

# 创建Gradio界面  
gr.ChatInterface(  
    predict,  
    title="Jambatypus",  
    description="与Jambatypus聊天！",  
    examples=[  
        ["你能解方程2x + 3 = 11求x吗？"],  
        ["写一首关于古罗马的史诗。"],  
        ["第一个登上月球的人是谁？"],  
        ["使用列表推导式生成1到10的平方列表。"],  
        ["推荐一些受欢迎的科幻书籍。"],  
        ["你能写一个关于时间旅行侦探的短篇故事吗？"]  
    ],  
    additional_inputs_accordion=gr.Accordion(label="⚙️ 参数", open=False),  
    additional_inputs=[  
        gr.Textbox("尽你所能完成任务。", label="系统提示"),  
        gr.Slider(0, 1, 0.8, label="温度"),  
        gr.Slider(128, 4096, 1024, label="最大新令牌数"),  
        gr.Slider(1, 80, 40, label="Top K采样"),  
        gr.Slider(0, 2, 1.1, label="重复惩罚"),  
        gr.Slider(0, 1, 0.95, label="Top P采样"),  
    ],  
    theme=gr.themes.Soft(primary_hue="green"),  
).queue().launch(share=True)