框架模型层

Triton 可以与 PyTorch 框架无缝集成，虽然 PyTorch 模型不会直接转换为 Triton，但可以利用 Triton 编写自定义的 CUDA 核心，从而优化特定的操作。这种方式让开发者可以在 PyTorch 中使用 Triton 优化的操作，提升性能。

例如，在 PyTorch 模型中包装 Triton 核心的代码：

class MyModel(torch.nn.Module):
    def forward(self, x, y):
        z = triton_add_wrapper(x, y)
        return z

Unsloth 是一个高效的库，使用 Triton 编写的算子，能够实现高性能的模型训练和推理，且没有准确性损失。下面是使用 Unsloth 的 FastLanguageModel 来加载一个预训练的 LLaMA 3 模型并进行推理的示例代码：

import time 
import torch
from unsloth import FastLanguageModel

max_seq_length = 2048 # Choose any! We auto support RoPE Scaling internally!
dtype = None # None for auto detection. Float16 for Tesla T4, V100, Bfloat16 for Ampere+
load_in_4bit = True # Use 4bit quantization to reduce memory usage. Can be False.

model, tokenizer = FastLanguageModel.from_pretrained(
    model_name = "/home/aii-works/llama3/Meta-Llama-3-8B-Instruct",
    max_seq_length = max_seq_length,
    dtype = dtype,
    load_in_4bit = load_in_4bit
)

model = FastLanguageModel.get_peft_model(
    model,
    r = 16, # Choose any number > 0 ! Suggested 8, 16, 32, 64, 128
    target_modules = ["q_proj", "k_proj", "v_proj", "o_proj",
                      "gate_proj", "up_proj", "down_proj",],
    lora_alpha = 16,
    lora_dropout = 0, # Supports any, but = 0 is optimized
    bias = "none",    # Supports any, but = "none" is optimized
    # [NEW] "unsloth" uses 30% less VRAM, fits 2x larger batch sizes!
    use_gradient_checkpointing = "unsloth", # True or "unsloth" for very long context
    random_state = 3407,
    use_rslora = False,  # We support rank stabilized LoRA
    loftq_config = None, # And LoftQ
)

alpaca_prompt = """Below is an instruction that describes a task, paired with an input that provides further context. Write a response that appropriately completes the request.

### Instruction:
{}

### Input:
{}

### Response:
{}"""

FastLanguageModel.for_inference(model) 
       
inputs = tokenizer(
[
    alpaca_prompt.format(
        # "Continue the fibonnaci sequence.", # instruction
        "Q:",
        "Name the planets in the solar system?",
        "", # output - leave this blank for generation!
    )
], return_tensors = "pt").to("cuda")

iterations = 10
with torch.no_grad():
    for _ in range(5):
        outputs = model.generate(**inputs, max_new_tokens = 64, use_cache = True) 
t_start = time.time()
for _ in range(iterations):
    with torch.no_grad():
        outputs = model.generate(**inputs, max_new_tokens = 64, use_cache = True) 
elapsed_time = time.time() - t_start
latency = elapsed_time / iterations * 1000
FPS = 1000 / latency

print(tokenizer.decode(outputs[0], skip_special_tokens=True))
print(f"FPS: {FPS:.2f}")

结果：

🦥 Unsloth: Will patch your computer to enable 2x faster free finetuning.
==((====))==  Unsloth 2024.8: Fast Llama patching. Transformers = 4.44.2.
   \\   /|    GPU: NVIDIA GeForce RTX 4080 SUPER. Max memory: 15.695 GB. Platform = Linux.
O^O/ \_/ \    Pytorch: 2.4.0. CUDA = 8.9. CUDA Toolkit = 12.1.
\        /    Bfloat16 = TRUE. FA [Xformers = 0.0.27.post2. FA2 = False]
 "-____-"     Free Apache license: http://github.com/unslothai/unsloth
Loading checkpoint shards: 100%|███████████████████████████████████████████████████████████████████████████████████████████████████████████████████████████████████████████████████████████████████████| 4/4 [00:02<00:00,  1.76it/s]
/home/aii-works/llama3/Meta-Llama-3-8B-Instruct does not have a padding token! Will use pad_token = <|reserved_special_token_250|>.
Unsloth 2024.8 patched 32 layers with 32 QKV layers, 32 O layers and 32 MLP layers.
huggingface/tokenizers: The current process just got forked, after parallelism has already been used. Disabling parallelism to avoid deadlocks...
To disable this warning, you can either:
        - Avoid using `tokenizers` before the fork if possible
        - Explicitly set the environment variable TOKENIZERS_PARALLELISM=(true | false)
huggingface/tokenizers: The current process just got forked, after parallelism has already been used. Disabling parallelism to avoid deadlocks...
To disable this warning, you can either:
        - Avoid using `tokenizers` before the fork if possible
        - Explicitly set the environment variable TOKENIZERS_PARALLELISM=(true | false)
Below is an instruction that describes a task, paired with an input that provides further context. Write a response that appropriately completes the request.

### Instruction:
Q:

### Input:
Name the planets in the solar system?

### Response:
The eight planets in our solar system, in order from the Sun, are:

1. Mercury
2. Venus
3. Earth
4. Mars
5. Jupiter
6. Saturn
7. Uranus
8. Neptune

Note: Pluto was previously considered a planet, but in 2006,
FPS: 0.89