docs/API.md

Fusion-LLM API 文档

模型 API

FusionMini

FusionMini 是 Fusion-LLM 的迷你模型实现。

类定义

class FusionMini(nn.Module):
    """
    Fusion-LLM 迷你模型
    
    Args:
        config (FusionMiniConfig): 模型配置
    """

方法

__init__(config)

初始化 FusionMini 模型。

参数：

• config (FusionMiniConfig): 模型配置对象

示例：

from models.fusion_mini import FusionMini, FusionMiniConfig

config = FusionMiniConfig(
    vocab_size=1000,
    hidden_size=128,
    num_hidden_layers=2,
)
model = FusionMini(config)

forward(input_ids, attention_mask=None, labels=None, past_key_values=None, use_cache=False, return_dict=True)

模型前向传播。

参数：

• input_ids (torch.Tensor): 输入 token IDs，形状为 (batch_size, sequence_length)
• attention_mask (torch.Tensor, optional): 注意力掩码，形状为 (batch_size, sequence_length)
• labels (torch.Tensor, optional): 标签，形状为 (batch_size, sequence_length)
• past_key_values (tuple, optional): 过去的键值缓存
• use_cache (bool): 是否使用 KV 缓存
• return_dict (bool): 是否返回字典格式

返回：

• 如果 return_dict=True：返回字典，包含：
  • loss (torch.Tensor): 损失值（如果提供了 labels）
  • logits (torch.Tensor): 逻辑值，形状为 (batch_size, sequence_length, vocab_size)
  • past_key_values (tuple): 更新的键值缓存（如果 use_cache=True）
• 如果 return_dict=False：返回元组

示例：

# 训练模式
outputs = model(
    input_ids=input_ids,
    labels=labels,
    return_dict=True,
)
loss = outputs["loss"]
logits = outputs["logits"]

# 推理模式（使用 KV 缓存）
outputs = model(
    input_ids=input_ids,
    use_cache=True,
    return_dict=True,
)
logits = outputs["logits"]
past_key_values = outputs["past_key_values"]

generate(input_ids, max_length=50, temperature=1.0, top_k=50, top_p=0.95)

生成文本（简易接口）。

参数：

• input_ids (torch.Tensor): 输入 token IDs
• max_length (int): 最大生成长度
• temperature (float): 温度参数
• top_k (int): Top-K 采样参数
• top_p (float): Top-P 采样参数

返回：

• torch.Tensor: 生成的 token IDs

示例：

generated = model.generate(
    input_ids=input_ids,
    max_length=50,
    temperature=0.8,
)

save_pretrained(save_directory)

保存模型和配置。

参数：

• save_directory (str): 保存目录

示例：

model.save_pretrained("output/my_model")

from_pretrained(pretrained_model_name_or_path)

从预训练路径加载模型。

参数：

• pretrained_model_name_or_path (str): 预训练模型路径

返回：

• FusionMini: 加载的模型

示例：

model = FusionMini.from_pretrained("output/my_model")

---

FusionMiniConfig

FusionMiniConfig 是 FusionMini 模型的配置类。

类定义

class FusionMiniConfig(PretrainedConfig):
    """
    FusionMini 模型配置
    
    Args:
        vocab_size (int): 词汇表大小
        hidden_size (int): 隐藏层大小
        num_hidden_layers (int): 隐藏层数量
        num_attention_heads (int): 注意力头数量
        intermediate_size (int): 中间层大小
        max_position_embeddings (int): 最大位置编码
        num_key_value_heads (int, optional): KV 头数量（用于 GQA）
        window_size (int, optional): SBLA 窗口大小
        ...
    """

属性

属性	类型	默认值	描述
vocab_size	int	50257	词汇表大小
hidden_size	int	768	隐藏层大小
num_hidden_layers	int	12	隐藏层数量
num_attention_heads	int	12	注意力头数量
intermediate_size	int	3072	中间层大小
max_position_embeddings	int	1024	最大位置编码
num_key_value_heads	int	None	KV 头数量（GQA）
window_size	int	16	SBLA 窗口大小

方法

__init__(**kwargs)

初始化配置。

示例：

config = FusionMiniConfig(
    vocab_size=1000,
    hidden_size=128,
    num_hidden_layers=2,
    num_attention_heads=2,
)

---

注意力 API

SBLAttention

SBLAttention 是 SBLA（Sliding Block Latent Attention）注意力实现。

类定义

class SBLAttention(nn.Module):
    """
    SBLA 注意力层
    
    Args:
        hidden_size (int): 隐藏层大小
        num_heads (int): 注意力头数量
        window_size (int): 窗口大小
        num_key_value_heads (int, optional): KV 头数量（用于 GQA）
    """

方法

__init__(hidden_size, num_heads, window_size, num_key_value_heads=None)

初始化 SBLA 注意力层。

forward(hidden_states, attention_mask=None, past_key_value=None, use_cache=False, output_attentions=False)

前向传播。

参数：

• hidden_states (torch.Tensor): 隐藏状态，形状为 (batch_size, sequence_length, hidden_size)
• attention_mask (torch.Tensor, optional): 注意力掩码
• past_key_value (tuple, optional): 过去的键值缓存
• use_cache (bool): 是否使用 KV 缓存
• output_attentions (bool): 是否输出注意力权重

返回：

• tuple: (output, past_key_value, attentions)

示例：

from models.sbla_attention import SBLAttention

attention = SBLAttention(
    hidden_size=128,
    num_heads=2,
    window_size=16,
)

hidden_states = torch.randn(1, 32, 128)
output, past_key_value, _ = attention(hidden_states)
print(output.shape)  # torch.Size([1, 32, 128])

---

Thinking Dial API

ThinkingDialProcessor

ThinkingDialProcessor 是 Thinking Dial（动态推理强度控制）处理器。

类定义

class ThinkingDialProcessor:
    """
    Thinking Dial 处理器
    
    用于处理 think token，动态控制推理强度。
    """

方法

process(text)

处理文本，注入 think token。

参数：

• text (str): 输入文本（可能包含 <|think_depth_N|>）

返回：

• str: 处理后的文本

示例：

from models.thinking_dial import ThinkingDialProcessor

processor = ThinkingDialProcessor()

text = "<|think_depth_2|> 这是一个需要深入思考的问题。"
processed_text = processor.process(text)
print(processed_text)  # 处理后的文本

get_think_depth(text)

获取 think token 的深度。

参数：

• text (str): 输入文本

返回：

• int: think 深度（0-3）

示例：

depth = processor.get_think_depth(text)
print(depth)  # 2

---

量化 API

DyQuant

DyQuant 是动态混合精度量化器（4/8/16-bit）。

类定义

class DyQuant:
    """
    动态混合精度量化器
    
    支持 4-bit、8-bit、16-bit 量化。
    """

方法

quantize(model, bits=8)

量化模型。

参数：

• model (nn.Module): 要量化的模型
• bits (int): 量化位数（4/8/16）

返回：

• nn.Module: 量化后的模型

示例：

from inference.dyquant import DyQuant

quantizer = DyQuant()
quantized_model = quantizer.quantize(model, bits=8)

save(model, path)

保存量化模型。

参数：

• model (nn.Module): 量化模型
• path (str): 保存路径

示例：

quantizer.save(quantized_model, "output/quantized_model")

load(path)

加载量化模型。

参数：

• path (str): 模型路径

返回：

• nn.Module: 加载的量化模型

示例：

loaded_model = quantizer.load("output/quantized_model")

---

训练 API

FullFinetuner

FullFinetuner 是全量微调器。

类定义

class FullFinetuner:
    """
    全量微调器
    
    用于全量微调 Fusion-LLM 模型。
    """

方法

train(model, train_dataset, eval_dataset=None, num_epochs=3, batch_size=4)

训练模型。

参数：

• model (nn.Module): 要训练的模型
• train_dataset (Dataset): 训练数据集
• eval_dataset (Dataset, optional): 评估数据集
• num_epochs (int): 训练轮数
• batch_size (int): 批次大小

示例：

from train.full_finetune import FullFinetuner

finetuner = FullFinetuner()

finetuner.train(
    model=model,
    train_dataset=train_dataset,
    eval_dataset=eval_dataset,
    num_epochs=3,
    batch_size=4,
)

---

评估 API

ModelEvaluator

ModelEvaluator 是模型评估器。

类定义

class ModelEvaluator:
    """
    模型评估器
    
    用于评估模型性能（Perplexity、Loss、Accuracy 等）。
    """

方法

evaluate(model, eval_data, metrics=["perplexity", "loss", "accuracy"])

评估模型。

参数：

• model (nn.Module): 要评估的模型
• eval_data (Dataset): 评估数据集
• metrics (List[str]): 评估指标列表

返回：

• EvaluationMetrics: 评估结果

示例：

from evaluation.metrics import ModelEvaluator

evaluator = ModelEvaluator()

metrics = evaluator.evaluate(
    model=model,
    eval_data=eval_dataset,
    metrics=["perplexity", "loss", "accuracy"],
)

print(f"Perplexity: {metrics.perplexity:.2f}")
print(f"Loss: {metrics.loss:.4f}")
print(f"Accuracy: {metrics.accuracy:.4f}")

---

部署 API

OllamaDeployer

OllamaDeployer 是 Ollama 部署器。

类定义

class OllamaDeployer:
    """
    Ollama 部署器
    
    用于将 Fusion-LLM 模型部署到 Ollama。
    """

方法

deploy(model_path, output_path)

部署模型到 Ollama。

参数：

• model_path (str): 模型路径
• output_path (str): 输出路径

示例：

from inference.ollama_deploy_v2 import OllamaDeployer

deployer = OllamaDeployer()

deployer.deploy(
    model_path="output/real_model",
    output_path="output/ollama_model",
)

---

数据集 API

TextDataset

TextDataset 是文本数据集。

类定义

class TextDataset(Dataset):
    """
    文本数据集
    
    用于加载文本数据并进行编码。
    """

方法

__init__(tokenizer, file_path, block_size=128)

初始化数据集。

参数：

• tokenizer: Tokenizer
• file_path (str): 文件路径
• block_size (int): 块大小

示例：

from torch.utils.data import Dataset

class TextDataset(Dataset):
    def __init__(self, tokenizer, file_path, block_size=128):
        # ...
        pass

---

完整示例

训练完整流程

import torch
from models.fusion_mini import FusionMini, FusionMiniConfig
from train.full_finetune import FullFinetuner
from evaluation.metrics import ModelEvaluator

# 1. 创建模型
config = FusionMiniConfig(
    vocab_size=1000,
    hidden_size=128,
    num_hidden_layers=2,
)
model = FusionMini(config)

# 2. 创建训练器
finetuner = FullFinetuner()

# 3. 训练
finetuner.train(
    model=model,
    train_dataset=train_dataset,
    eval_dataset=eval_dataset,
    num_epochs=3,
    batch_size=4,
)

# 4. 评估
evaluator = ModelEvaluator()
metrics = evaluator.evaluate(
    model=model,
    eval_data=eval_dataset,
)

print(f"Perplexity: {metrics.perplexity:.2f}")
print(f"Loss: {metrics.loss:.4f}")

推理完整流程

import torch
from models.fusion_mini import FusionMini

# 1. 加载模型
model = FusionMini.from_pretrained("output/real_model")
model.eval()

# 2. 创建输入
input_ids = torch.tensor([[1, 2, 3, 4, 5]])

# 3. 推理
with torch.no_grad():
    outputs = model(
        input_ids=input_ids,
        return_dict=True,
    )
    
    logits = outputs["logits"]
    print(f"Logits shape: {logits.shape}")

# 4. 生成
generated = model.generate(
    input_ids=input_ids,
    max_length=50,
)
print(f"Generated: {generated}")

---

常见问题

1. 如何自定义模型配置？

使用 FusionMiniConfig：

config = FusionMiniConfig(
    vocab_size=2000,       # 更大的词汇表
    hidden_size=256,        # 更大的隐藏层
    num_hidden_layers=4,    # 更深的模型
    num_attention_heads=4,  # 更多的注意力头
)

2. 如何使用 GQA？

在配置中设置 num_key_value_heads：

config = FusionMiniConfig(
    num_attention_heads=8,   # 8 个查询头
    num_key_value_heads=2,    # 2 个 KV 头（GQA）
)

3. 如何启用 KV 缓存？

在推理时使用 use_cache=True：

outputs = model(
    input_ids=input_ids,
    use_cache=True,
    return_dict=True,
)
past_key_values = outputs["past_key_values"]

4. 如何量化模型？

使用 DyQuant：

from inference.dyquant import DyQuant

quantizer = DyQuant()
quantized_model = quantizer.quantize(model, bits=8)

---

许可证

Fusion-LLM API 文档采用 Apache 2.0 许可证。