初始化项目，由ModelHub XC社区提供模型
Model: celestialcreator/Llama-3.2-1B-MTP-k8 Source: Original Platform
This commit is contained in:
ModelHub XC
36
.gitattributes
vendored
Normal file
36
.gitattributes
vendored
Normal file
@@ -0,0 +1,36 @@
|
||||
*.7z filter=lfs diff=lfs merge=lfs -text
|
||||
*.arrow filter=lfs diff=lfs merge=lfs -text
|
||||
*.bin filter=lfs diff=lfs merge=lfs -text
|
||||
*.bz2 filter=lfs diff=lfs merge=lfs -text
|
||||
*.ckpt filter=lfs diff=lfs merge=lfs -text
|
||||
*.ftz filter=lfs diff=lfs merge=lfs -text
|
||||
*.gz filter=lfs diff=lfs merge=lfs -text
|
||||
*.h5 filter=lfs diff=lfs merge=lfs -text
|
||||
*.joblib filter=lfs diff=lfs merge=lfs -text
|
||||
*.lfs.* filter=lfs diff=lfs merge=lfs -text
|
||||
*.mlmodel filter=lfs diff=lfs merge=lfs -text
|
||||
*.model filter=lfs diff=lfs merge=lfs -text
|
||||
*.msgpack filter=lfs diff=lfs merge=lfs -text
|
||||
*.npy filter=lfs diff=lfs merge=lfs -text
|
||||
*.npz filter=lfs diff=lfs merge=lfs -text
|
||||
*.onnx filter=lfs diff=lfs merge=lfs -text
|
||||
*.ot filter=lfs diff=lfs merge=lfs -text
|
||||
*.parquet filter=lfs diff=lfs merge=lfs -text
|
||||
*.pb filter=lfs diff=lfs merge=lfs -text
|
||||
*.pickle filter=lfs diff=lfs merge=lfs -text
|
||||
*.pkl filter=lfs diff=lfs merge=lfs -text
|
||||
*.pt filter=lfs diff=lfs merge=lfs -text
|
||||
*.pth filter=lfs diff=lfs merge=lfs -text
|
||||
*.rar filter=lfs diff=lfs merge=lfs -text
|
||||
*.safetensors filter=lfs diff=lfs merge=lfs -text
|
||||
saved_model/**/* filter=lfs diff=lfs merge=lfs -text
|
||||
*.tar.* filter=lfs diff=lfs merge=lfs -text
|
||||
*.tar filter=lfs diff=lfs merge=lfs -text
|
||||
*.tflite filter=lfs diff=lfs merge=lfs -text
|
||||
*.tgz filter=lfs diff=lfs merge=lfs -text
|
||||
*.wasm filter=lfs diff=lfs merge=lfs -text
|
||||
*.xz filter=lfs diff=lfs merge=lfs -text
|
||||
*.zip filter=lfs diff=lfs merge=lfs -text
|
||||
*.zst filter=lfs diff=lfs merge=lfs -text
|
||||
*tfevents* filter=lfs diff=lfs merge=lfs -text
|
||||
tokenizer.json filter=lfs diff=lfs merge=lfs -text
|
||||
162
README.md
Normal file
162
README.md
Normal file
@@ -0,0 +1,162 @@
|
||||
---
|
||||
license: llama3.2
|
||||
base_model: meta-llama/Llama-3.2-1B
|
||||
tags:
|
||||
- multi-token-prediction
|
||||
- speculative-decoding
|
||||
- self-distillation
|
||||
- mtp
|
||||
- llama
|
||||
- consumer-gpu
|
||||
- rtx-5090
|
||||
- paper-reproduction
|
||||
datasets:
|
||||
- jwkirchenbauer/metamathqa-grouped-split
|
||||
language:
|
||||
- en
|
||||
pipeline_tag: text-generation
|
||||
library_name: transformers
|
||||
model-index:
|
||||
- name: Llama-3.2-1B-MTP-k8
|
||||
results:
|
||||
- task:
|
||||
type: text-generation
|
||||
name: GSM8K (8-shot CoT)
|
||||
dataset:
|
||||
type: gsm8k
|
||||
name: GSM8K
|
||||
metrics:
|
||||
- name: exact_match (flexible)
|
||||
type: exact_match
|
||||
value: 5.08
|
||||
- name: exact_match (strict)
|
||||
type: exact_match
|
||||
value: 3.03
|
||||
---
|
||||
|
||||
# Llama-3.2-1B-MTP-k8: Multi-Token Prediction on a Single Consumer GPU
|
||||
|
||||
This is a reproduction of **"Multi-Token Prediction via Self-Distillation"** ([arXiv 2602.06019](https://arxiv.org/abs/2602.06019)) adapted for a single NVIDIA RTX 5090 (32GB). The original paper used 4x NVIDIA GH200 (384GB total) with Llama-3.1-8B. We scaled it down to Llama-3.2-1B on consumer hardware.
|
||||
|
||||
## What is Multi-Token Prediction (MTP)?
|
||||
|
||||
Standard language models predict **one token at a time** (autoregressive decoding). MTP trains the model to predict **multiple future tokens simultaneously** using online self-distillation:
|
||||
|
||||
1. A **frozen teacher** (the original model) generates soft probability distributions
|
||||
2. A **trainable student** (same architecture) learns to predict k future tokens at each position
|
||||
3. At inference, **ConfAdapt decoding** emits multiple tokens when the model is confident, falling back to single-token when uncertain
|
||||
|
||||
The result: **faster inference with minimal quality loss**.
|
||||
|
||||
## Results: GSM8K 8-shot Chain-of-Thought
|
||||
|
||||
| Configuration | Exact Match (flexible) | Exact Match (strict) | Throughput |
|
||||
|---|---|---|---|
|
||||
| **Baseline** (Llama-3.2-1B, standard AR) | **7.13%** ± 0.71 | **6.07%** ± 0.66 | ~1.5 s/sample |
|
||||
| **MTP k=1** (single token, quality check) | 5.23% ± 0.61 | 2.96% ± 0.47 | ~2.4 s/sample |
|
||||
| **MTP k=8 + ConfAdapt 90%** | 5.08% ± 0.60 | 3.03% ± 0.47 | **~1.3 s/sample** |
|
||||
|
||||
### Key Findings
|
||||
|
||||
- **ConfAdapt works:** k=8 with ConfAdapt matches k=1 quality while being **1.8x faster** (avg 2.82 tokens emitted per step)
|
||||
- **Quality drop is expected:** The ~2% accuracy drop from baseline is consistent with our smaller setup (1B model, 500M training tokens vs paper's 8B model, 2B tokens)
|
||||
- **The core claim holds:** Multi-token decoding via ConfAdapt preserves generation quality while improving throughput, even on a tiny 1B model
|
||||
|
||||
## Training Details
|
||||
|
||||
### What We Changed from the Paper
|
||||
|
||||
| Parameter | Paper (8B / 4x GH200) | Ours (1B / 1x RTX 5090) |
|
||||
|---|---|---|
|
||||
| Base model | Llama-3.1-8B | Llama-3.2-1B |
|
||||
| GPUs | 4x GH200 (96GB each) | 1x RTX 5090 (32GB) |
|
||||
| FSDP mesh | 1x4 | 1x1 (no FSDP) |
|
||||
| k_toks | Randomized 2-16 across ranks | Fixed 8 |
|
||||
| Training tokens | 2B | 500M |
|
||||
| micro_batch_size | 32 | 8 |
|
||||
| global_batch_size | 128 | 64 (grad accumulation) |
|
||||
| mask_region_ct | 5 | 1 |
|
||||
| rollout_multiplier | 4 | 2 |
|
||||
| Template | Chat (Instruct tokenizer) | Plain text (base tokenizer) |
|
||||
|
||||
### What We Kept the Same
|
||||
|
||||
- **Supervision method:** Soft teacher via KL divergence (paper's recommended self-distillation)
|
||||
- **Dataset:** MetaMathQA (`jwkirchenbauer/metamathqa-grouped-split`)
|
||||
- **Sequence length:** 160 tokens
|
||||
- **Peak learning rate:** 1e-5
|
||||
- **Optimizer:** AdamW with cosine decay
|
||||
|
||||
### Training Metrics
|
||||
|
||||
- **Total steps:** 48,828
|
||||
- **Training time:** ~17 hours on RTX 5090
|
||||
- **Final train loss:** ~0.9
|
||||
- **Final val loss:** 1.895 (perplexity 6.65)
|
||||
|
||||
## Why k=8 Instead of the Paper's Randomized k=2-16?
|
||||
|
||||
The paper's approach randomizes k across GPU ranks each step. With 4 GPUs, the model sees k=2, k=5, k=12, k=16 simultaneously in a single batch, learning to handle any prediction horizon.
|
||||
|
||||
With a single GPU, we can only train one k value per step. We chose k=8 as a middle ground — large enough to demonstrate meaningful multi-token speedup, small enough to fit in 32GB VRAM.
|
||||
|
||||
This is an important tradeoff: our model is specialized for k=8, while the paper's model generalizes across all k values. A production deployment would benefit from the paper's multi-GPU randomized approach.
|
||||
|
||||
## Infrastructure: Running on Consumer Hardware
|
||||
|
||||
This reproduction ran entirely on a home Kubernetes cluster:
|
||||
|
||||
- **GPU:** NVIDIA RTX 5090 (32GB, Blackwell architecture / sm_120)
|
||||
- **System:** 16GB RAM, Debian 13
|
||||
- **Stack:** Kubernetes + containerd + NVIDIA device plugin
|
||||
- **PyTorch:** Nightly build with CUDA 12.8 (required for Blackwell sm_120 support)
|
||||
|
||||
### Challenges We Solved
|
||||
|
||||
1. **Blackwell GPU support:** RTX 5090 (sm_120) requires PyTorch nightly with cu128 — stable releases don't include sm_120 yet
|
||||
2. **Single-GPU checkpoint saving:** The original code uses `torch.distributed.all_reduce()` for checkpoint state sync, which crashes when distributed is not initialized. We added an `is_initialized()` guard
|
||||
3. **W&B configuration:** Default config points to the paper authors' organization. Override with `--wandb.entity=null`
|
||||
4. **HuggingFace checkpoint format:** The litgpt converter outputs `model.pth` but transformers expects `pytorch_model.bin`
|
||||
|
||||
## Usage
|
||||
|
||||
```python
|
||||
from transformers import AutoModelForCausalLM, AutoTokenizer
|
||||
|
||||
model = AutoModelForCausalLM.from_pretrained(
|
||||
"celestialcreator/Llama-3.2-1B-MTP-k8",
|
||||
trust_remote_code=True,
|
||||
torch_dtype="float16",
|
||||
)
|
||||
tokenizer = AutoTokenizer.from_pretrained("celestialcreator/Llama-3.2-1B-MTP-k8")
|
||||
|
||||
# Standard generation (single token, works like any Llama model)
|
||||
inputs = tokenizer("The capital of France is", return_tensors="pt")
|
||||
outputs = model.generate(**inputs, max_new_tokens=50)
|
||||
print(tokenizer.decode(outputs[0], skip_special_tokens=True))
|
||||
```
|
||||
|
||||
For MTP inference with ConfAdapt decoding, use the [mtp-lm evaluation harness fork](https://github.com/jwkirchenbauer/lm-evaluation-harness-mtp-lm).
|
||||
|
||||
## Reproduction Guide
|
||||
|
||||
Full reproduction instructions with Kubernetes manifests and configs: [GitHub Fork](https://github.com/CelestialCreator/mtp-lm)
|
||||
|
||||
## Citation
|
||||
|
||||
If you use this model, please cite the original paper:
|
||||
|
||||
```bibtex
|
||||
@article{kirchenbauer2025multitokenpredictionselfdistillation,
|
||||
title={Multi-Token Prediction via Self-Distillation},
|
||||
author={John Kirchenbauer and Jonas Geiping and Yuxin Wen and Tom Goldstein},
|
||||
journal={arXiv preprint arXiv:2602.06019},
|
||||
year={2025}
|
||||
}
|
||||
```
|
||||
|
||||
## Acknowledgments
|
||||
|
||||
- Original paper and code by [John Kirchenbauer et al.](https://github.com/jwkirchenbauer/mtp-lm)
|
||||
- Built with [LitGPT](https://github.com/Lightning-AI/litgpt), [PyTorch](https://pytorch.org/), and [lm-evaluation-harness](https://github.com/EleutherAI/lm-evaluation-harness)
|
||||
|
||||
41
config.json
Normal file
41
config.json
Normal file
@@ -0,0 +1,41 @@
|
||||
{
|
||||
"architectures": [
|
||||
"LlamaForCausalLM"
|
||||
],
|
||||
"attention_bias": false,
|
||||
"attention_dropout": 0.0,
|
||||
"bos_token_id": 128000,
|
||||
"dtype": "float32",
|
||||
"eos_token_id": 128001,
|
||||
"head_dim": 64,
|
||||
"hidden_act": "silu",
|
||||
"hidden_size": 2048,
|
||||
"initializer_range": 0.02,
|
||||
"intermediate_size": 8192,
|
||||
"max_position_embeddings": 131072,
|
||||
"mlp_bias": false,
|
||||
"model_type": "llama",
|
||||
"num_attention_heads": 32,
|
||||
"num_hidden_layers": 16,
|
||||
"num_key_value_heads": 8,
|
||||
"pad_token_id": 128004,
|
||||
"pretraining_tp": 1,
|
||||
"rms_norm_eps": 1e-05,
|
||||
"rope_scaling": {
|
||||
"factor": 32.0,
|
||||
"high_freq_factor": 4.0,
|
||||
"low_freq_factor": 1.0,
|
||||
"original_max_position_embeddings": 8192,
|
||||
"rope_type": "llama3"
|
||||
},
|
||||
"rope_theta": 500000.0,
|
||||
"tie_word_embeddings": true,
|
||||
"transformers_version": "4.56.2",
|
||||
"unsloth_fixed": true,
|
||||
"use_cache": true,
|
||||
"vocab_size": 128384,
|
||||
"auto_map": {
|
||||
"AutoConfig": "configuration_llama.LlamaConfig",
|
||||
"AutoModelForCausalLM": "modeling_llama.LlamaForCausalLM"
|
||||
}
|
||||
}
|
||||
226
configuration_llama.py
Normal file
226
configuration_llama.py
Normal file
@@ -0,0 +1,226 @@
|
||||
# fmt: off
|
||||
# coding=utf-8
|
||||
# Copyright 2022 EleutherAI and the HuggingFace Inc. team. All rights reserved.
|
||||
#
|
||||
# This code is based on EleutherAI's GPT-NeoX library and the GPT-NeoX
|
||||
# and OPT implementations in this library. It has been modified from its
|
||||
# original forms to accommodate minor architectural differences compared
|
||||
# to GPT-NeoX and OPT used by the Meta AI team that trained the model.
|
||||
#
|
||||
# Licensed under the Apache License, Version 2.0 (the "License");
|
||||
# you may not use this file except in compliance with the License.
|
||||
# You may obtain a copy of the License at
|
||||
#
|
||||
# http://www.apache.org/licenses/LICENSE-2.0
|
||||
#
|
||||
# Unless required by applicable law or agreed to in writing, software
|
||||
# distributed under the License is distributed on an "AS IS" BASIS,
|
||||
# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
|
||||
# See the License for the specific language governing permissions and
|
||||
# limitations under the License.
|
||||
"""LLaMA model configuration"""
|
||||
|
||||
from transformers.configuration_utils import PretrainedConfig
|
||||
from transformers.modeling_rope_utils import rope_config_validation
|
||||
|
||||
|
||||
class LlamaConfig(PretrainedConfig):
|
||||
r"""
|
||||
This is the configuration class to store the configuration of a [`LlamaModel`]. It is used to instantiate an LLaMA
|
||||
model according to the specified arguments, defining the model architecture. Instantiating a configuration with the
|
||||
defaults will yield a similar configuration to that of the LLaMA-7B.
|
||||
e.g. [meta-llama/Llama-2-7b-hf](https://huggingface.co/meta-llama/Llama-2-7b-hf)
|
||||
|
||||
Configuration objects inherit from [`PretrainedConfig`] and can be used to control the model outputs. Read the
|
||||
documentation from [`PretrainedConfig`] for more information.
|
||||
|
||||
|
||||
Args:
|
||||
vocab_size (`int`, *optional*, defaults to 32000):
|
||||
Vocabulary size of the LLaMA model. Defines the number of different tokens that can be represented by the
|
||||
`inputs_ids` passed when calling [`LlamaModel`]
|
||||
hidden_size (`int`, *optional*, defaults to 4096):
|
||||
Dimension of the hidden representations.
|
||||
intermediate_size (`int`, *optional*, defaults to 11008):
|
||||
Dimension of the MLP representations.
|
||||
num_hidden_layers (`int`, *optional*, defaults to 32):
|
||||
Number of hidden layers in the Transformer decoder.
|
||||
num_attention_heads (`int`, *optional*, defaults to 32):
|
||||
Number of attention heads for each attention layer in the Transformer decoder.
|
||||
num_key_value_heads (`int`, *optional*):
|
||||
This is the number of key_value heads that should be used to implement Grouped Query Attention. If
|
||||
`num_key_value_heads=num_attention_heads`, the model will use Multi Head Attention (MHA), if
|
||||
`num_key_value_heads=1` the model will use Multi Query Attention (MQA) otherwise GQA is used. When
|
||||
converting a multi-head checkpoint to a GQA checkpoint, each group key and value head should be constructed
|
||||
by meanpooling all the original heads within that group. For more details, check out [this
|
||||
paper](https://huggingface.co/papers/2305.13245). If it is not specified, will default to
|
||||
`num_attention_heads`.
|
||||
hidden_act (`str` or `function`, *optional*, defaults to `"silu"`):
|
||||
The non-linear activation function (function or string) in the decoder.
|
||||
max_position_embeddings (`int`, *optional*, defaults to 2048):
|
||||
The maximum sequence length that this model might ever be used with. Llama 1 supports up to 2048 tokens,
|
||||
Llama 2 up to 4096, CodeLlama up to 16384.
|
||||
initializer_range (`float`, *optional*, defaults to 0.02):
|
||||
The standard deviation of the truncated_normal_initializer for initializing all weight matrices.
|
||||
rms_norm_eps (`float`, *optional*, defaults to 1e-06):
|
||||
The epsilon used by the rms normalization layers.
|
||||
use_cache (`bool`, *optional*, defaults to `True`):
|
||||
Whether or not the model should return the last key/values attentions (not used by all models). Only
|
||||
relevant if `config.is_decoder=True`.
|
||||
pad_token_id (`int`, *optional*):
|
||||
Padding token id.
|
||||
bos_token_id (`int`, *optional*, defaults to 1):
|
||||
Beginning of stream token id.
|
||||
eos_token_id (`int`, *optional*, defaults to 2):
|
||||
End of stream token id.
|
||||
pretraining_tp (`int`, *optional*, defaults to 1):
|
||||
Experimental feature. Tensor parallelism rank used during pretraining. Please refer to [this
|
||||
document](https://huggingface.co/docs/transformers/main/perf_train_gpu_many#tensor-parallelism) to
|
||||
understand more about it. This value is necessary to ensure exact reproducibility of the pretraining
|
||||
results. Please refer to [this issue](https://github.com/pytorch/pytorch/issues/76232).
|
||||
tie_word_embeddings (`bool`, *optional*, defaults to `False`):
|
||||
Whether to tie weight embeddings
|
||||
rope_theta (`float`, *optional*, defaults to 10000.0):
|
||||
The base period of the RoPE embeddings.
|
||||
rope_scaling (`Dict`, *optional*):
|
||||
Dictionary containing the scaling configuration for the RoPE embeddings. NOTE: if you apply new rope type
|
||||
and you expect the model to work on longer `max_position_embeddings`, we recommend you to update this value
|
||||
accordingly.
|
||||
Expected contents:
|
||||
`rope_type` (`str`):
|
||||
The sub-variant of RoPE to use. Can be one of ['default', 'linear', 'dynamic', 'yarn', 'longrope',
|
||||
'llama3'], with 'default' being the original RoPE implementation.
|
||||
`factor` (`float`, *optional*):
|
||||
Used with all rope types except 'default'. The scaling factor to apply to the RoPE embeddings. In
|
||||
most scaling types, a `factor` of x will enable the model to handle sequences of length x *
|
||||
original maximum pre-trained length.
|
||||
`original_max_position_embeddings` (`int`, *optional*):
|
||||
Used with 'dynamic', 'longrope' and 'llama3'. The original max position embeddings used during
|
||||
pretraining.
|
||||
`attention_factor` (`float`, *optional*):
|
||||
Used with 'yarn' and 'longrope'. The scaling factor to be applied on the attention
|
||||
computation. If unspecified, it defaults to value recommended by the implementation, using the
|
||||
`factor` field to infer the suggested value.
|
||||
`beta_fast` (`float`, *optional*):
|
||||
Only used with 'yarn'. Parameter to set the boundary for extrapolation (only) in the linear
|
||||
ramp function. If unspecified, it defaults to 32.
|
||||
`beta_slow` (`float`, *optional*):
|
||||
Only used with 'yarn'. Parameter to set the boundary for interpolation (only) in the linear
|
||||
ramp function. If unspecified, it defaults to 1.
|
||||
`short_factor` (`list[float]`, *optional*):
|
||||
Only used with 'longrope'. The scaling factor to be applied to short contexts (<
|
||||
`original_max_position_embeddings`). Must be a list of numbers with the same length as the hidden
|
||||
size divided by the number of attention heads divided by 2
|
||||
`long_factor` (`list[float]`, *optional*):
|
||||
Only used with 'longrope'. The scaling factor to be applied to long contexts (<
|
||||
`original_max_position_embeddings`). Must be a list of numbers with the same length as the hidden
|
||||
size divided by the number of attention heads divided by 2
|
||||
`low_freq_factor` (`float`, *optional*):
|
||||
Only used with 'llama3'. Scaling factor applied to low frequency components of the RoPE
|
||||
`high_freq_factor` (`float`, *optional*):
|
||||
Only used with 'llama3'. Scaling factor applied to high frequency components of the RoPE
|
||||
attention_bias (`bool`, *optional*, defaults to `False`):
|
||||
Whether to use a bias in the query, key, value and output projection layers during self-attention.
|
||||
attention_dropout (`float`, *optional*, defaults to 0.0):
|
||||
The dropout ratio for the attention probabilities.
|
||||
mlp_bias (`bool`, *optional*, defaults to `False`):
|
||||
Whether to use a bias in up_proj, down_proj and gate_proj layers in the MLP layers.
|
||||
head_dim (`int`, *optional*):
|
||||
The attention head dimension. If None, it will default to hidden_size // num_attention_heads
|
||||
|
||||
```python
|
||||
>>> from transformers import LlamaModel, LlamaConfig
|
||||
|
||||
>>> # Initializing a LLaMA llama-7b style configuration
|
||||
>>> configuration = LlamaConfig()
|
||||
|
||||
>>> # Initializing a model from the llama-7b style configuration
|
||||
>>> model = LlamaModel(configuration)
|
||||
|
||||
>>> # Accessing the model configuration
|
||||
>>> configuration = model.config
|
||||
```"""
|
||||
|
||||
model_type = "llama"
|
||||
keys_to_ignore_at_inference = ["past_key_values"]
|
||||
# Default tensor parallel plan for base model `LlamaModel`
|
||||
base_model_tp_plan = {
|
||||
"layers.*.self_attn.q_proj": "colwise",
|
||||
"layers.*.self_attn.k_proj": "colwise",
|
||||
"layers.*.self_attn.v_proj": "colwise",
|
||||
"layers.*.self_attn.o_proj": "rowwise",
|
||||
"layers.*.mlp.gate_proj": "colwise",
|
||||
"layers.*.mlp.up_proj": "colwise",
|
||||
"layers.*.mlp.down_proj": "rowwise",
|
||||
}
|
||||
base_model_pp_plan = {
|
||||
"embed_tokens": (["input_ids"], ["inputs_embeds"]),
|
||||
"layers": (["hidden_states", "attention_mask"], ["hidden_states"]),
|
||||
"norm": (["hidden_states"], ["hidden_states"]),
|
||||
}
|
||||
|
||||
def __init__(
|
||||
self,
|
||||
vocab_size=32000,
|
||||
hidden_size=4096,
|
||||
intermediate_size=11008,
|
||||
num_hidden_layers=32,
|
||||
num_attention_heads=32,
|
||||
num_key_value_heads=None,
|
||||
hidden_act="silu",
|
||||
max_position_embeddings=2048,
|
||||
initializer_range=0.02,
|
||||
rms_norm_eps=1e-6,
|
||||
use_cache=True,
|
||||
pad_token_id=None,
|
||||
bos_token_id=1,
|
||||
eos_token_id=2,
|
||||
pretraining_tp=1,
|
||||
tie_word_embeddings=False,
|
||||
rope_theta=10000.0,
|
||||
rope_scaling=None,
|
||||
attention_bias=False,
|
||||
attention_dropout=0.0,
|
||||
mlp_bias=False,
|
||||
head_dim=None,
|
||||
**kwargs,
|
||||
):
|
||||
self.vocab_size = vocab_size
|
||||
self.max_position_embeddings = max_position_embeddings
|
||||
self.hidden_size = hidden_size
|
||||
self.intermediate_size = intermediate_size
|
||||
self.num_hidden_layers = num_hidden_layers
|
||||
self.num_attention_heads = num_attention_heads
|
||||
|
||||
# for backward compatibility
|
||||
if num_key_value_heads is None:
|
||||
num_key_value_heads = num_attention_heads
|
||||
|
||||
self.num_key_value_heads = num_key_value_heads
|
||||
self.hidden_act = hidden_act
|
||||
self.initializer_range = initializer_range
|
||||
self.rms_norm_eps = rms_norm_eps
|
||||
self.pretraining_tp = pretraining_tp
|
||||
self.use_cache = use_cache
|
||||
self.rope_theta = rope_theta
|
||||
self.rope_scaling = rope_scaling
|
||||
self.attention_bias = attention_bias
|
||||
self.attention_dropout = attention_dropout
|
||||
self.mlp_bias = mlp_bias
|
||||
self.head_dim = head_dim if head_dim is not None else self.hidden_size // self.num_attention_heads
|
||||
# Validate the correctness of rotary position embeddings parameters
|
||||
# BC: if there is a 'type' field, copy it it to 'rope_type'.
|
||||
if self.rope_scaling is not None and "type" in self.rope_scaling:
|
||||
self.rope_scaling["rope_type"] = self.rope_scaling["type"]
|
||||
rope_config_validation(self)
|
||||
|
||||
super().__init__(
|
||||
pad_token_id=pad_token_id,
|
||||
bos_token_id=bos_token_id,
|
||||
eos_token_id=eos_token_id,
|
||||
tie_word_embeddings=tie_word_embeddings,
|
||||
**kwargs,
|
||||
)
|
||||
|
||||
|
||||
__all__ = ["LlamaConfig"]
|
||||
11
generation_config.json
Normal file
11
generation_config.json
Normal file
@@ -0,0 +1,11 @@
|
||||
{
|
||||
"_from_model_config": true,
|
||||
"bos_token_id": 128000,
|
||||
"do_sample": true,
|
||||
"eos_token_id": 128001,
|
||||
"max_length": 131072,
|
||||
"pad_token_id": 128004,
|
||||
"temperature": 0.6,
|
||||
"top_p": 0.9,
|
||||
"transformers_version": "4.56.2"
|
||||
}
|
||||
969
modeling_llama.py
Normal file
969
modeling_llama.py
Normal file
@@ -0,0 +1,969 @@
|
||||
# fmt: off
|
||||
|
||||
# Adaptation recipe lifted from Jonas et al. :>
|
||||
# https://github.com/seal-rg/recurrent-pretraining/blob/main/recpre/raven_modeling_minimal.py
|
||||
|
||||
# coding=utf-8
|
||||
# Copyright 2022 EleutherAI and the HuggingFace Inc. team. All rights reserved.
|
||||
#
|
||||
# This code is based on EleutherAI's GPT-NeoX library and the GPT-NeoX
|
||||
# and OPT implementations in this library. It has been modified from its
|
||||
# original forms to accommodate minor architectural differences compared
|
||||
# to GPT-NeoX and OPT used by the Meta AI team that trained the model.
|
||||
#
|
||||
# Licensed under the Apache License, Version 2.0 (the "License");
|
||||
# you may not use this file except in compliance with the License.
|
||||
# You may obtain a copy of the License at
|
||||
#
|
||||
# http://www.apache.org/licenses/LICENSE-2.0
|
||||
#
|
||||
# Unless required by applicable law or agreed to in writing, software
|
||||
# distributed under the License is distributed on an "AS IS" BASIS,
|
||||
# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
|
||||
# See the License for the specific language governing permissions and
|
||||
# limitations under the License.
|
||||
from typing import Callable, Optional, Union
|
||||
import time
|
||||
|
||||
import torch
|
||||
from torch import nn
|
||||
import torch.nn.functional as F
|
||||
|
||||
from transformers.activations import ACT2FN
|
||||
from transformers.cache_utils import Cache, DynamicCache
|
||||
from transformers.generation import GenerationMixin
|
||||
from transformers.integrations import use_kernel_forward_from_hub
|
||||
from transformers.masking_utils import create_causal_mask
|
||||
from transformers.modeling_layers import (
|
||||
GenericForQuestionAnswering,
|
||||
GenericForSequenceClassification,
|
||||
GenericForTokenClassification,
|
||||
GradientCheckpointingLayer,
|
||||
)
|
||||
from transformers.modeling_outputs import (
|
||||
BaseModelOutputWithPast,
|
||||
CausalLMOutputWithPast,
|
||||
)
|
||||
from transformers.modeling_rope_utils import ROPE_INIT_FUNCTIONS, dynamic_rope_update
|
||||
from transformers.modeling_utils import ALL_ATTENTION_FUNCTIONS, PreTrainedModel
|
||||
from transformers.processing_utils import Unpack
|
||||
from transformers.utils import TransformersKwargs, auto_docstring, can_return_tuple, logging
|
||||
from transformers.utils.deprecation import deprecate_kwarg
|
||||
from transformers.utils.generic import check_model_inputs
|
||||
from .configuration_llama import LlamaConfig
|
||||
|
||||
|
||||
# Glue
|
||||
from transformers.generation.utils import GenerateDecoderOnlyOutput
|
||||
|
||||
|
||||
logger = logging.get_logger(__name__)
|
||||
|
||||
|
||||
@use_kernel_forward_from_hub("RMSNorm")
|
||||
class LlamaRMSNorm(nn.Module):
|
||||
def __init__(self, hidden_size, eps=1e-6):
|
||||
"""
|
||||
LlamaRMSNorm is equivalent to T5LayerNorm
|
||||
"""
|
||||
super().__init__()
|
||||
self.weight = nn.Parameter(torch.ones(hidden_size))
|
||||
self.variance_epsilon = eps
|
||||
|
||||
def forward(self, hidden_states):
|
||||
input_dtype = hidden_states.dtype
|
||||
hidden_states = hidden_states.to(torch.float32)
|
||||
variance = hidden_states.pow(2).mean(-1, keepdim=True)
|
||||
hidden_states = hidden_states * torch.rsqrt(variance + self.variance_epsilon)
|
||||
return self.weight * hidden_states.to(input_dtype)
|
||||
|
||||
def extra_repr(self):
|
||||
return f"{tuple(self.weight.shape)}, eps={self.variance_epsilon}"
|
||||
|
||||
|
||||
class LlamaRotaryEmbedding(nn.Module):
|
||||
inv_freq: torch.Tensor # fix linting for `register_buffer`
|
||||
|
||||
def __init__(self, config: LlamaConfig, device=None):
|
||||
super().__init__()
|
||||
# BC: "rope_type" was originally "type"
|
||||
if hasattr(config, "rope_scaling") and isinstance(config.rope_scaling, dict):
|
||||
self.rope_type = config.rope_scaling.get("rope_type", config.rope_scaling.get("type"))
|
||||
else:
|
||||
self.rope_type = "default"
|
||||
self.max_seq_len_cached = config.max_position_embeddings
|
||||
self.original_max_seq_len = config.max_position_embeddings
|
||||
|
||||
self.config = config
|
||||
self.rope_init_fn = ROPE_INIT_FUNCTIONS[self.rope_type]
|
||||
|
||||
inv_freq, self.attention_scaling = self.rope_init_fn(self.config, device)
|
||||
self.register_buffer("inv_freq", inv_freq, persistent=False)
|
||||
self.original_inv_freq = self.inv_freq
|
||||
|
||||
@torch.no_grad()
|
||||
@dynamic_rope_update # power user: used with advanced RoPE types (e.g. dynamic rope)
|
||||
def forward(self, x, position_ids):
|
||||
inv_freq_expanded = self.inv_freq[None, :, None].float().expand(position_ids.shape[0], -1, 1).to(x.device)
|
||||
position_ids_expanded = position_ids[:, None, :].float()
|
||||
|
||||
device_type = x.device.type if isinstance(x.device.type, str) and x.device.type != "mps" else "cpu"
|
||||
with torch.autocast(device_type=device_type, enabled=False): # Force float32
|
||||
freqs = (inv_freq_expanded.float() @ position_ids_expanded.float()).transpose(1, 2)
|
||||
emb = torch.cat((freqs, freqs), dim=-1)
|
||||
cos = emb.cos() * self.attention_scaling
|
||||
sin = emb.sin() * self.attention_scaling
|
||||
|
||||
return cos.to(dtype=x.dtype), sin.to(dtype=x.dtype)
|
||||
|
||||
|
||||
def rotate_half(x):
|
||||
"""Rotates half the hidden dims of the input."""
|
||||
x1 = x[..., : x.shape[-1] // 2]
|
||||
x2 = x[..., x.shape[-1] // 2 :]
|
||||
return torch.cat((-x2, x1), dim=-1)
|
||||
|
||||
|
||||
def apply_rotary_pos_emb(q, k, cos, sin, position_ids=None, unsqueeze_dim=1):
|
||||
"""Applies Rotary Position Embedding to the query and key tensors.
|
||||
|
||||
Args:
|
||||
q (`torch.Tensor`): The query tensor.
|
||||
k (`torch.Tensor`): The key tensor.
|
||||
cos (`torch.Tensor`): The cosine part of the rotary embedding.
|
||||
sin (`torch.Tensor`): The sine part of the rotary embedding.
|
||||
position_ids (`torch.Tensor`, *optional*):
|
||||
Deprecated and unused.
|
||||
unsqueeze_dim (`int`, *optional*, defaults to 1):
|
||||
The 'unsqueeze_dim' argument specifies the dimension along which to unsqueeze cos[position_ids] and
|
||||
sin[position_ids] so that they can be properly broadcasted to the dimensions of q and k. For example, note
|
||||
that cos[position_ids] and sin[position_ids] have the shape [batch_size, seq_len, head_dim]. Then, if q and
|
||||
k have the shape [batch_size, heads, seq_len, head_dim], then setting unsqueeze_dim=1 makes
|
||||
cos[position_ids] and sin[position_ids] broadcastable to the shapes of q and k. Similarly, if q and k have
|
||||
the shape [batch_size, seq_len, heads, head_dim], then set unsqueeze_dim=2.
|
||||
Returns:
|
||||
`tuple(torch.Tensor)` comprising of the query and key tensors rotated using the Rotary Position Embedding.
|
||||
"""
|
||||
cos = cos.unsqueeze(unsqueeze_dim)
|
||||
sin = sin.unsqueeze(unsqueeze_dim)
|
||||
q_embed = (q * cos) + (rotate_half(q) * sin)
|
||||
k_embed = (k * cos) + (rotate_half(k) * sin)
|
||||
return q_embed, k_embed
|
||||
|
||||
|
||||
class LlamaMLP(nn.Module):
|
||||
def __init__(self, config):
|
||||
super().__init__()
|
||||
self.config = config
|
||||
self.hidden_size = config.hidden_size
|
||||
self.intermediate_size = config.intermediate_size
|
||||
self.gate_proj = nn.Linear(self.hidden_size, self.intermediate_size, bias=config.mlp_bias)
|
||||
self.up_proj = nn.Linear(self.hidden_size, self.intermediate_size, bias=config.mlp_bias)
|
||||
self.down_proj = nn.Linear(self.intermediate_size, self.hidden_size, bias=config.mlp_bias)
|
||||
self.act_fn = ACT2FN[config.hidden_act]
|
||||
|
||||
def forward(self, x):
|
||||
down_proj = self.down_proj(self.act_fn(self.gate_proj(x)) * self.up_proj(x))
|
||||
return down_proj
|
||||
|
||||
|
||||
def repeat_kv(hidden_states: torch.Tensor, n_rep: int) -> torch.Tensor:
|
||||
"""
|
||||
This is the equivalent of torch.repeat_interleave(x, dim=1, repeats=n_rep). The hidden states go from (batch,
|
||||
num_key_value_heads, seqlen, head_dim) to (batch, num_attention_heads, seqlen, head_dim)
|
||||
"""
|
||||
batch, num_key_value_heads, slen, head_dim = hidden_states.shape
|
||||
if n_rep == 1:
|
||||
return hidden_states
|
||||
hidden_states = hidden_states[:, :, None, :, :].expand(batch, num_key_value_heads, n_rep, slen, head_dim)
|
||||
return hidden_states.reshape(batch, num_key_value_heads * n_rep, slen, head_dim)
|
||||
|
||||
|
||||
def eager_attention_forward(
|
||||
module: nn.Module,
|
||||
query: torch.Tensor,
|
||||
key: torch.Tensor,
|
||||
value: torch.Tensor,
|
||||
attention_mask: Optional[torch.Tensor],
|
||||
scaling: float,
|
||||
dropout: float = 0.0,
|
||||
**kwargs: Unpack[TransformersKwargs],
|
||||
):
|
||||
key_states = repeat_kv(key, module.num_key_value_groups)
|
||||
value_states = repeat_kv(value, module.num_key_value_groups)
|
||||
|
||||
attn_weights = torch.matmul(query, key_states.transpose(2, 3)) * scaling
|
||||
if attention_mask is not None:
|
||||
causal_mask = attention_mask[:, :, :, : key_states.shape[-2]]
|
||||
attn_weights = attn_weights + causal_mask
|
||||
|
||||
attn_weights = nn.functional.softmax(attn_weights, dim=-1, dtype=torch.float32).to(query.dtype)
|
||||
attn_weights = nn.functional.dropout(attn_weights, p=dropout, training=module.training)
|
||||
attn_output = torch.matmul(attn_weights, value_states)
|
||||
attn_output = attn_output.transpose(1, 2).contiguous()
|
||||
|
||||
return attn_output, attn_weights
|
||||
|
||||
|
||||
class LlamaAttention(nn.Module):
|
||||
"""Multi-headed attention from 'Attention Is All You Need' paper"""
|
||||
|
||||
def __init__(self, config: LlamaConfig, layer_idx: int):
|
||||
super().__init__()
|
||||
self.config = config
|
||||
self.layer_idx = layer_idx
|
||||
self.head_dim = getattr(config, "head_dim", config.hidden_size // config.num_attention_heads)
|
||||
self.num_key_value_groups = config.num_attention_heads // config.num_key_value_heads
|
||||
self.scaling = self.head_dim**-0.5
|
||||
self.attention_dropout = config.attention_dropout
|
||||
self.is_causal = True
|
||||
|
||||
self.q_proj = nn.Linear(
|
||||
config.hidden_size, config.num_attention_heads * self.head_dim, bias=config.attention_bias
|
||||
)
|
||||
self.k_proj = nn.Linear(
|
||||
config.hidden_size, config.num_key_value_heads * self.head_dim, bias=config.attention_bias
|
||||
)
|
||||
self.v_proj = nn.Linear(
|
||||
config.hidden_size, config.num_key_value_heads * self.head_dim, bias=config.attention_bias
|
||||
)
|
||||
self.o_proj = nn.Linear(
|
||||
config.num_attention_heads * self.head_dim, config.hidden_size, bias=config.attention_bias
|
||||
)
|
||||
|
||||
@deprecate_kwarg("past_key_value", new_name="past_key_values", version="4.58")
|
||||
def forward(
|
||||
self,
|
||||
hidden_states: torch.Tensor,
|
||||
position_embeddings: tuple[torch.Tensor, torch.Tensor],
|
||||
attention_mask: Optional[torch.Tensor],
|
||||
past_key_values: Optional[Cache] = None,
|
||||
cache_position: Optional[torch.LongTensor] = None,
|
||||
**kwargs: Unpack[TransformersKwargs],
|
||||
) -> tuple[torch.Tensor, torch.Tensor]:
|
||||
input_shape = hidden_states.shape[:-1]
|
||||
hidden_shape = (*input_shape, -1, self.head_dim)
|
||||
|
||||
query_states = self.q_proj(hidden_states).view(hidden_shape).transpose(1, 2)
|
||||
key_states = self.k_proj(hidden_states).view(hidden_shape).transpose(1, 2)
|
||||
value_states = self.v_proj(hidden_states).view(hidden_shape).transpose(1, 2)
|
||||
|
||||
cos, sin = position_embeddings
|
||||
query_states, key_states = apply_rotary_pos_emb(query_states, key_states, cos, sin)
|
||||
|
||||
if past_key_values is not None:
|
||||
# sin and cos are specific to RoPE models; cache_position needed for the static cache
|
||||
cache_kwargs = {"sin": sin, "cos": cos, "cache_position": cache_position}
|
||||
key_states, value_states = past_key_values.update(key_states, value_states, self.layer_idx, cache_kwargs)
|
||||
|
||||
attention_interface: Callable = eager_attention_forward
|
||||
if self.config._attn_implementation != "eager":
|
||||
attention_interface = ALL_ATTENTION_FUNCTIONS[self.config._attn_implementation]
|
||||
|
||||
attn_output, attn_weights = attention_interface(
|
||||
self,
|
||||
query_states,
|
||||
key_states,
|
||||
value_states,
|
||||
attention_mask,
|
||||
dropout=0.0 if not self.training else self.attention_dropout,
|
||||
scaling=self.scaling,
|
||||
**kwargs,
|
||||
)
|
||||
|
||||
attn_output = attn_output.reshape(*input_shape, -1).contiguous()
|
||||
attn_output = self.o_proj(attn_output)
|
||||
return attn_output, attn_weights
|
||||
|
||||
|
||||
class LlamaDecoderLayer(GradientCheckpointingLayer):
|
||||
def __init__(self, config: LlamaConfig, layer_idx: int):
|
||||
super().__init__()
|
||||
self.hidden_size = config.hidden_size
|
||||
|
||||
self.self_attn = LlamaAttention(config=config, layer_idx=layer_idx)
|
||||
|
||||
self.mlp = LlamaMLP(config)
|
||||
self.input_layernorm = LlamaRMSNorm(config.hidden_size, eps=config.rms_norm_eps)
|
||||
self.post_attention_layernorm = LlamaRMSNorm(config.hidden_size, eps=config.rms_norm_eps)
|
||||
|
||||
@deprecate_kwarg("past_key_value", new_name="past_key_values", version="4.58")
|
||||
def forward(
|
||||
self,
|
||||
hidden_states: torch.Tensor,
|
||||
attention_mask: Optional[torch.Tensor] = None,
|
||||
position_ids: Optional[torch.LongTensor] = None,
|
||||
past_key_values: Optional[Cache] = None,
|
||||
use_cache: Optional[bool] = False,
|
||||
cache_position: Optional[torch.LongTensor] = None,
|
||||
position_embeddings: Optional[tuple[torch.Tensor, torch.Tensor]] = None, # necessary, but kept here for BC
|
||||
**kwargs: Unpack[TransformersKwargs],
|
||||
) -> torch.Tensor:
|
||||
residual = hidden_states
|
||||
hidden_states = self.input_layernorm(hidden_states)
|
||||
# Self Attention
|
||||
hidden_states, _ = self.self_attn(
|
||||
hidden_states=hidden_states,
|
||||
attention_mask=attention_mask,
|
||||
position_ids=position_ids,
|
||||
past_key_values=past_key_values,
|
||||
use_cache=use_cache,
|
||||
cache_position=cache_position,
|
||||
position_embeddings=position_embeddings,
|
||||
**kwargs,
|
||||
)
|
||||
hidden_states = residual + hidden_states
|
||||
|
||||
# Fully Connected
|
||||
residual = hidden_states
|
||||
hidden_states = self.post_attention_layernorm(hidden_states)
|
||||
hidden_states = self.mlp(hidden_states)
|
||||
hidden_states = residual + hidden_states
|
||||
return hidden_states
|
||||
|
||||
|
||||
@auto_docstring
|
||||
class LlamaPreTrainedModel(PreTrainedModel):
|
||||
config: LlamaConfig
|
||||
base_model_prefix = "model"
|
||||
supports_gradient_checkpointing = True
|
||||
_no_split_modules = ["LlamaDecoderLayer"]
|
||||
_skip_keys_device_placement = ["past_key_values"]
|
||||
_supports_flash_attn = True
|
||||
_supports_sdpa = True
|
||||
_supports_flex_attn = True
|
||||
|
||||
_can_compile_fullgraph = True
|
||||
_supports_attention_backend = True
|
||||
_can_record_outputs = {
|
||||
"hidden_states": LlamaDecoderLayer,
|
||||
"attentions": LlamaAttention,
|
||||
}
|
||||
|
||||
|
||||
@auto_docstring
|
||||
class LlamaModel(LlamaPreTrainedModel):
|
||||
def __init__(self, config: LlamaConfig):
|
||||
super().__init__(config)
|
||||
self.padding_idx = config.pad_token_id
|
||||
self.vocab_size = config.vocab_size
|
||||
|
||||
self.embed_tokens = nn.Embedding(config.vocab_size, config.hidden_size, self.padding_idx)
|
||||
self.layers = nn.ModuleList(
|
||||
[LlamaDecoderLayer(config, layer_idx) for layer_idx in range(config.num_hidden_layers)]
|
||||
)
|
||||
self.norm = LlamaRMSNorm(config.hidden_size, eps=config.rms_norm_eps)
|
||||
self.rotary_emb = LlamaRotaryEmbedding(config=config)
|
||||
self.gradient_checkpointing = False
|
||||
|
||||
# Initialize weights and apply final processing
|
||||
self.post_init()
|
||||
|
||||
@check_model_inputs
|
||||
@auto_docstring
|
||||
def forward(
|
||||
self,
|
||||
input_ids: Optional[torch.LongTensor] = None,
|
||||
attention_mask: Optional[torch.Tensor] = None,
|
||||
position_ids: Optional[torch.LongTensor] = None,
|
||||
past_key_values: Optional[Cache] = None,
|
||||
inputs_embeds: Optional[torch.FloatTensor] = None,
|
||||
cache_position: Optional[torch.LongTensor] = None,
|
||||
use_cache: Optional[bool] = None,
|
||||
**kwargs: Unpack[TransformersKwargs],
|
||||
) -> BaseModelOutputWithPast:
|
||||
if (input_ids is None) ^ (inputs_embeds is not None):
|
||||
raise ValueError("You must specify exactly one of input_ids or inputs_embeds")
|
||||
|
||||
if inputs_embeds is None:
|
||||
inputs_embeds: torch.Tensor = self.embed_tokens(input_ids)
|
||||
|
||||
if use_cache and past_key_values is None:
|
||||
past_key_values = DynamicCache(config=self.config)
|
||||
|
||||
if cache_position is None:
|
||||
past_seen_tokens = past_key_values.get_seq_length() if past_key_values is not None else 0
|
||||
cache_position: torch.Tensor = torch.arange(
|
||||
past_seen_tokens, past_seen_tokens + inputs_embeds.shape[1], device=inputs_embeds.device
|
||||
)
|
||||
|
||||
if position_ids is None:
|
||||
position_ids = cache_position.unsqueeze(0)
|
||||
|
||||
causal_mask = create_causal_mask(
|
||||
config=self.config,
|
||||
input_embeds=inputs_embeds,
|
||||
attention_mask=attention_mask,
|
||||
cache_position=cache_position,
|
||||
past_key_values=past_key_values,
|
||||
position_ids=position_ids,
|
||||
)
|
||||
|
||||
hidden_states = inputs_embeds
|
||||
position_embeddings = self.rotary_emb(hidden_states, position_ids)
|
||||
|
||||
for decoder_layer in self.layers[: self.config.num_hidden_layers]:
|
||||
hidden_states = decoder_layer(
|
||||
hidden_states,
|
||||
attention_mask=causal_mask,
|
||||
position_ids=position_ids,
|
||||
past_key_values=past_key_values,
|
||||
cache_position=cache_position,
|
||||
position_embeddings=position_embeddings,
|
||||
**kwargs,
|
||||
)
|
||||
|
||||
hidden_states = self.norm(hidden_states)
|
||||
return BaseModelOutputWithPast(
|
||||
last_hidden_state=hidden_states,
|
||||
past_key_values=past_key_values,
|
||||
)
|
||||
|
||||
|
||||
@auto_docstring
|
||||
class LlamaForCausalLM(LlamaPreTrainedModel, GenerationMixin):
|
||||
_tied_weights_keys = ["lm_head.weight"]
|
||||
_tp_plan = {"lm_head": "colwise_rep"}
|
||||
_pp_plan = {"lm_head": (["hidden_states"], ["logits"])}
|
||||
|
||||
def __init__(self, config):
|
||||
super().__init__(config)
|
||||
self.model = LlamaModel(config)
|
||||
self.vocab_size = config.vocab_size
|
||||
self.lm_head = nn.Linear(config.hidden_size, config.vocab_size, bias=False)
|
||||
|
||||
# Initialize weights and apply final processing
|
||||
self.post_init()
|
||||
|
||||
print(f"Loading local LlamaForCausalLM!")
|
||||
|
||||
@can_return_tuple
|
||||
@auto_docstring
|
||||
def forward(
|
||||
self,
|
||||
input_ids: Optional[torch.LongTensor] = None,
|
||||
attention_mask: Optional[torch.Tensor] = None,
|
||||
position_ids: Optional[torch.LongTensor] = None,
|
||||
past_key_values: Optional[Cache] = None,
|
||||
inputs_embeds: Optional[torch.FloatTensor] = None,
|
||||
labels: Optional[torch.LongTensor] = None,
|
||||
use_cache: Optional[bool] = None,
|
||||
cache_position: Optional[torch.LongTensor] = None,
|
||||
logits_to_keep: Union[int, torch.Tensor] = 0,
|
||||
**kwargs: Unpack[TransformersKwargs],
|
||||
) -> CausalLMOutputWithPast:
|
||||
r"""
|
||||
Example:
|
||||
|
||||
```python
|
||||
>>> from transformers import AutoTokenizer, LlamaForCausalLM
|
||||
|
||||
>>> model = LlamaForCausalLM.from_pretrained("meta-llama/Llama-2-7b-hf")
|
||||
>>> tokenizer = AutoTokenizer.from_pretrained("meta-llama/Llama-2-7b-hf")
|
||||
|
||||
>>> prompt = "Hey, are you conscious? Can you talk to me?"
|
||||
>>> inputs = tokenizer(prompt, return_tensors="pt")
|
||||
|
||||
>>> # Generate
|
||||
>>> generate_ids = model.generate(inputs.input_ids, max_length=30)
|
||||
>>> tokenizer.batch_decode(generate_ids, skip_special_tokens=True, clean_up_tokenization_spaces=False)[0]
|
||||
"Hey, are you conscious? Can you talk to me?\nI'm not conscious, but I can talk to you."
|
||||
```"""
|
||||
outputs: BaseModelOutputWithPast = self.model(
|
||||
input_ids=input_ids,
|
||||
attention_mask=attention_mask,
|
||||
position_ids=position_ids,
|
||||
past_key_values=past_key_values,
|
||||
inputs_embeds=inputs_embeds,
|
||||
use_cache=use_cache,
|
||||
cache_position=cache_position,
|
||||
**kwargs,
|
||||
)
|
||||
|
||||
hidden_states = outputs.last_hidden_state
|
||||
# Only compute necessary logits, and do not upcast them to float if we are not computing the loss
|
||||
slice_indices = slice(-logits_to_keep, None) if isinstance(logits_to_keep, int) else logits_to_keep
|
||||
logits = self.lm_head(hidden_states[:, slice_indices, :])
|
||||
|
||||
loss = None
|
||||
if labels is not None:
|
||||
loss = self.loss_function(logits=logits, labels=labels, vocab_size=self.config.vocab_size, **kwargs)
|
||||
|
||||
return CausalLMOutputWithPast(
|
||||
loss=loss,
|
||||
logits=logits,
|
||||
past_key_values=outputs.past_key_values,
|
||||
hidden_states=outputs.hidden_states,
|
||||
attentions=outputs.attentions,
|
||||
)
|
||||
|
||||
@torch.inference_mode()
|
||||
def generate(self, input_ids: Optional[torch.LongTensor] = None, *args, **kwargs):
|
||||
"""
|
||||
Custom generate entry point.
|
||||
|
||||
If `do_mtp=True` is passed, strictly enforces MTP-only arguments and routes to `_mtp_generate`.
|
||||
Otherwise, routes to standard Hugging Face `generate`.
|
||||
"""
|
||||
|
||||
# 1. Standard Path
|
||||
if not kwargs.get("do_mtp", False):
|
||||
print("Executing standard generate() codepath.")
|
||||
return super().generate(input_ids, *args, **kwargs)
|
||||
|
||||
# 2. MTP Path
|
||||
print("Executing custom MTP generation codepath!")
|
||||
|
||||
# Handle input_ids: HF can pass it as positional (first arg) or keyword
|
||||
# We consolidate it into 'prompt' for the MTP signature
|
||||
prompt = input_ids if input_ids is not None else kwargs.pop("input_ids", None)
|
||||
if prompt is None:
|
||||
# If standard generate was called without input_ids, it might be in *args or handled deeper,
|
||||
# but for MTP we require it explicitly.
|
||||
raise ValueError("MTP generation requires 'input_ids' to be passed.")
|
||||
|
||||
# --- Argument Extraction & Strict Validation ---
|
||||
|
||||
# Keys strictly allowed for MTP (will be passed to _mtp_generate)
|
||||
mtp_allowed_keys = {
|
||||
"do_mtp", "k_toks", "mask_id", "min_mask_id", "max_mask_id", "strategy", "return_mtp_result_dict", "include_prompt", "streamer"
|
||||
}
|
||||
|
||||
# Keys from HF that we know how to map to MTP equivalents
|
||||
# max_length -> max_returned_tokens
|
||||
max_length = kwargs.pop("max_length", None)
|
||||
max_returned_tokens = kwargs.pop("max_returned_tokens", None)
|
||||
if max_returned_tokens is None and max_length is not None:
|
||||
print(f"Renaming max_length={max_length} to max_returned_tokens for MTP generation.")
|
||||
max_returned_tokens = max_length
|
||||
|
||||
# eos_token_id -> eos_id
|
||||
eos_token_id = kwargs.pop("eos_token_id", None)
|
||||
eos_id = kwargs.pop("eos_id", None)
|
||||
if eos_id is None:
|
||||
eos_id = eos_token_id
|
||||
|
||||
# Standard HF args that we SILENTLY IGNORE because they are passed automatically
|
||||
# by the GenerationMixin but are not relevant or supported in this MTP implementation.
|
||||
ignored_hf_keys = {
|
||||
"attention_mask", "use_cache", "do_sample", "stopping_criteria",
|
||||
"pad_token_id", "logits_processor", "max_new_tokens", "generation_config",
|
||||
}
|
||||
|
||||
# Check for explicit incompatibility
|
||||
if kwargs.get("do_sample", False):
|
||||
raise ValueError("MTP generation does not support sampling (do_sample=True).")
|
||||
|
||||
# Extract valid MTP args
|
||||
mtp_kwargs = {}
|
||||
for k in list(kwargs.keys()):
|
||||
if k in mtp_allowed_keys:
|
||||
mtp_kwargs[k] = kwargs.pop(k)
|
||||
|
||||
# Remove ignored HF keys
|
||||
for k in list(kwargs.keys()):
|
||||
if k in ignored_hf_keys:
|
||||
kwargs.pop(k)
|
||||
|
||||
# FAIL LOUDLY if anything is left in kwargs
|
||||
if kwargs:
|
||||
raise ValueError(
|
||||
f"Unsupported argument(s) passed to MTP generate: {list(kwargs.keys())}.\n"
|
||||
f"When do_mtp=True, only these args are supported: {list(mtp_allowed_keys) + ['max_returned_tokens', 'eos_id']}."
|
||||
)
|
||||
|
||||
# Pre-flight checks
|
||||
if len(prompt.shape) > 1 and prompt.shape[0] > 1:
|
||||
raise NotImplementedError("MTP generation currently only supports single-example generation (no batching).")
|
||||
|
||||
# Execute Unified Implementation
|
||||
# Note: We remove 'do_mtp' from kwargs before passing, as the impl doesn't need it
|
||||
mtp_kwargs.pop("do_mtp", None)
|
||||
|
||||
return self._mtp_generate(
|
||||
prompt=prompt,
|
||||
max_returned_tokens=max_returned_tokens,
|
||||
eos_id=eos_id,
|
||||
**mtp_kwargs
|
||||
)
|
||||
|
||||
@torch.inference_mode()
|
||||
def _mtp_generate(
|
||||
self,
|
||||
prompt: torch.Tensor,
|
||||
max_returned_tokens: int = None,
|
||||
k_toks: int = 1,
|
||||
mask_id: int = None,
|
||||
min_mask_id: int = None,
|
||||
max_mask_id: int = None,
|
||||
eos_id: Optional[Union[int, list]] = None,
|
||||
include_prompt: bool = True,
|
||||
streamer = None,
|
||||
strategy: Optional[list] = None,
|
||||
return_mtp_result_dict: bool = False,
|
||||
):
|
||||
"""
|
||||
Implementation of MTP generation logic.
|
||||
"""
|
||||
# --- Setup Stop Tokens ---
|
||||
if isinstance(eos_id, int):
|
||||
stop_tokens = ([eos_id,],)
|
||||
elif isinstance(eos_id, list):
|
||||
assert all(isinstance(eid, int) for eid in eos_id), "If eos_id is a list, all elements must be ints."
|
||||
stop_tokens = tuple([list([eid,]) for eid in eos_id])
|
||||
elif eos_id is None:
|
||||
stop_tokens = ()
|
||||
else:
|
||||
raise ValueError(f"eos_id must be None, int, or list of lists, got {type(eos_id)}")
|
||||
|
||||
# --- Validation ---
|
||||
if k_toks > 1: assert (mask_id is not None) or (min_mask_id is not None), "mask_id must be provided when k_toks > 1"
|
||||
|
||||
input_ids = prompt.clone()
|
||||
prompt_size = prompt.size(1)
|
||||
device = prompt.device
|
||||
|
||||
# Get generation config (defaulting to model's if not present)
|
||||
generation_config = self.generation_config
|
||||
|
||||
# --- Streaming Prompt ---
|
||||
if include_prompt:
|
||||
if streamer is not None:
|
||||
print(f"\n<BEGIN Streaming Prompt>", flush=True)
|
||||
streamer.put(input_ids)
|
||||
print(f"\n<END Streaming Prompt>", flush=True)
|
||||
|
||||
if streamer is not None:
|
||||
print(f"\n<BEGIN Streaming Generation>", flush=True)
|
||||
|
||||
stop_progress = [0] * len(stop_tokens)
|
||||
|
||||
# --- Generation Loop ---
|
||||
t0_prefill = time.perf_counter()
|
||||
t1_prefill = None
|
||||
t0_gen = None
|
||||
t1_gen = None
|
||||
toks_pre_prefill = input_ids.shape[1]
|
||||
toks_post_prefill = None
|
||||
current_idx = input_ids.shape[1]
|
||||
num_fwd_evals = 0
|
||||
effective_k_values = []
|
||||
|
||||
# Prepare kwargs for the inner loop
|
||||
model_kwargs = {}
|
||||
|
||||
while current_idx + k_toks <= max_returned_tokens:
|
||||
|
||||
# 0 is prefill, 1 is first step which can include compile time, then 2 is steady state
|
||||
if (t0_gen is None and num_fwd_evals == 2):
|
||||
t1_prefill = time.perf_counter()
|
||||
t0_gen = time.perf_counter()
|
||||
toks_post_prefill = input_ids.shape[1]
|
||||
|
||||
# Generate the token
|
||||
if k_toks > 1:
|
||||
input_ids = self._extend_w_mask(input_ids=input_ids, k_toks=k_toks, mask_id=mask_id, min_mask_id=min_mask_id, max_mask_id=max_mask_id)
|
||||
|
||||
# first step prep
|
||||
if num_fwd_evals == 0:
|
||||
model_kwargs, generation_config = self._prep_generate_args(
|
||||
self,
|
||||
input_ids,
|
||||
generation_config,
|
||||
)
|
||||
assert "token_type_ids" not in model_kwargs
|
||||
assert "attention_mask" not in model_kwargs
|
||||
assert "decoder_attention_mask" not in model_kwargs
|
||||
|
||||
model_inputs = self.prepare_inputs_for_generation(input_ids, **model_kwargs)
|
||||
|
||||
next_tokens, model_outputs = self._mtp_next_tokens(
|
||||
self,
|
||||
model_inputs,
|
||||
k_toks=k_toks,
|
||||
strategy=strategy,
|
||||
)
|
||||
effective_k_values.append(next_tokens.shape[1])
|
||||
|
||||
if streamer is not None: streamer.put(next_tokens)
|
||||
|
||||
# remove the masks if any
|
||||
if k_toks > 1:
|
||||
input_ids = input_ids[:, :-(k_toks - 1)]
|
||||
input_ids = torch.cat([input_ids, next_tokens], dim=-1)
|
||||
|
||||
# Update cache / model kwargs
|
||||
model_kwargs["past_key_values"] = model_outputs.past_key_values
|
||||
|
||||
if strategy is None:
|
||||
if num_fwd_evals == 0:
|
||||
# this is the end of the prefill step
|
||||
if k_toks > 1:
|
||||
model_kwargs["cache_position"] = torch.arange(
|
||||
prompt_size,
|
||||
prompt_size + k_toks + (k_toks - 1),
|
||||
device=device,
|
||||
dtype=torch.int64,
|
||||
)
|
||||
else:
|
||||
model_kwargs["cache_position"] = torch.tensor(
|
||||
[prompt_size], device=device, dtype=torch.int64
|
||||
)
|
||||
else:
|
||||
model_kwargs["cache_position"].add_(k_toks)
|
||||
else: # we can assume that all strats produce variable number of tokens
|
||||
num_new_tokens = next_tokens.shape[1]
|
||||
if num_fwd_evals == 0:
|
||||
# this is the end of the prefill step
|
||||
if k_toks > 1:
|
||||
model_kwargs["cache_position"] = torch.arange(
|
||||
prompt_size,
|
||||
prompt_size + num_new_tokens + (k_toks - 1),
|
||||
device=device,
|
||||
dtype=torch.int64,
|
||||
)
|
||||
else:
|
||||
model_kwargs["cache_position"] = torch.tensor(
|
||||
[prompt_size], device=device, dtype=torch.int64
|
||||
)
|
||||
else:
|
||||
if k_toks > 1:
|
||||
recomputation_positions = model_kwargs["cache_position"][: -(k_toks - 1)]
|
||||
previous_num_new_tokens = recomputation_positions.size(0)
|
||||
new_start_pos = recomputation_positions[0] + previous_num_new_tokens
|
||||
model_kwargs["cache_position"] = torch.arange(
|
||||
new_start_pos,
|
||||
new_start_pos + num_new_tokens + (k_toks - 1),
|
||||
device=device,
|
||||
dtype=torch.int64,
|
||||
)
|
||||
else:
|
||||
model_kwargs["cache_position"].add_(1)
|
||||
|
||||
current_idx += next_tokens.shape[1]
|
||||
num_fwd_evals += 1
|
||||
|
||||
# Crop cache
|
||||
model_kwargs["past_key_values"].crop(model_kwargs["cache_position"][0])
|
||||
|
||||
# Check for stop sequences
|
||||
hit_stop_seq = False
|
||||
for int_tok in next_tokens.tolist()[0]: # assuming batch size 1
|
||||
for i, seq in enumerate(stop_tokens):
|
||||
if int_tok == seq[stop_progress[i]]:
|
||||
stop_progress[i] += 1
|
||||
if stop_progress[i] == len(seq):
|
||||
hit_stop_seq = True
|
||||
break
|
||||
else:
|
||||
stop_progress[i] = 0
|
||||
if hit_stop_seq:
|
||||
break
|
||||
if hit_stop_seq:
|
||||
break
|
||||
|
||||
# End of generation loop
|
||||
if streamer is not None:
|
||||
streamer.end()
|
||||
print(f"\n<END Streaming Generation>", flush=True)
|
||||
|
||||
t1_gen = time.perf_counter()
|
||||
|
||||
# Calculate stats
|
||||
if t1_prefill is not None and t0_gen is not None and toks_post_prefill is not None:
|
||||
t_prefill = t1_prefill - t0_prefill
|
||||
t_gen = t1_gen - t0_gen
|
||||
tokens_generated = input_ids.shape[1] - toks_post_prefill
|
||||
toks_gend_incl_prefillplus1 = input_ids.shape[1] - toks_pre_prefill
|
||||
else:
|
||||
t_prefill = t1_gen - t0_prefill
|
||||
t_gen = t1_gen - t0_prefill
|
||||
tokens_generated = input_ids.shape[1] - toks_pre_prefill
|
||||
toks_gend_incl_prefillplus1 = input_ids.shape[1] - toks_pre_prefill
|
||||
|
||||
print(
|
||||
f"Using a total of {f'1+1+{num_fwd_evals-2}' if t0_gen is not None else f'{num_fwd_evals}'} forward evals, time for prefill plus first/compilation step {f'(1+1)' if t0_gen is not None else ''} was {t_prefill:.02f} sec, generation time was {t_gen:.02f} sec @ {tokens_generated / t_gen:.02f} tokens/sec {f'steady state' if t0_gen is not None else ''} over {tokens_generated} tokens.",
|
||||
flush=True,
|
||||
)
|
||||
print(f"Strategy used: {strategy}", flush=True)
|
||||
avg_effective_k = sum(effective_k_values) / len(effective_k_values) if effective_k_values else 0
|
||||
print(
|
||||
f"Average effective k_toks over generation: {avg_effective_k:.02f}, full array of effective k_toks: {effective_k_values}",
|
||||
flush=True,
|
||||
)
|
||||
|
||||
if generation_config.return_dict_in_generate:
|
||||
raise NotImplementedError(f"Only basic return type implemented for MTP generate.")
|
||||
|
||||
if return_mtp_result_dict:
|
||||
token_ids = input_ids if include_prompt else input_ids[:,prompt_size:]
|
||||
|
||||
# testing for generation vs. aux data order integrity
|
||||
import hashlib
|
||||
leading_toks = token_ids[0,prompt_size:prompt_size+10].tolist()
|
||||
leading_toks_hash = hashlib.shake_128(str(leading_toks).encode()).hexdigest(4)
|
||||
|
||||
mtp_result_dict = {
|
||||
"token_ids":token_ids,
|
||||
"leading_toks_hash": leading_toks_hash,
|
||||
"num_fwd_evals": num_fwd_evals,
|
||||
"t_prefill": t_prefill,
|
||||
"t_gen": t_gen,
|
||||
"tokens_generated": tokens_generated,
|
||||
"toks_gend_incl_prefillplus1": toks_gend_incl_prefillplus1,
|
||||
"avg_effective_k": avg_effective_k,
|
||||
"effective_k_values": effective_k_values,
|
||||
"tps": tokens_generated / t_gen if t_gen > 0 else 0.0,
|
||||
}
|
||||
return mtp_result_dict
|
||||
|
||||
return input_ids if include_prompt else input_ids[:,prompt_size:]
|
||||
|
||||
@torch.inference_mode()
|
||||
def _prep_generate_args(
|
||||
self,
|
||||
model,
|
||||
input_ids: torch.Tensor,
|
||||
generation_config = None,
|
||||
model_kwargs: dict = None,
|
||||
):
|
||||
# Setup
|
||||
if model_kwargs is None:
|
||||
model_kwargs = {}
|
||||
|
||||
if generation_config is None:
|
||||
generation_config = model.generation_config
|
||||
model_kwargs["use_cache"] = True
|
||||
if "past_key_values" in model_kwargs:
|
||||
print(f"Before _get_initial_cache_position past_key_values cache length: {model_kwargs['past_key_values'].layers[0].get_seq_length()}", flush=True)
|
||||
model_kwargs = model._get_initial_cache_position(input_ids.shape[1], input_ids.device, model_kwargs)
|
||||
if "past_key_values" in model_kwargs:
|
||||
print(f"After _get_initial_cache_position past_key_values cache length: {model_kwargs['past_key_values'].layers[0].get_seq_length()}", flush=True)
|
||||
return model_kwargs, generation_config
|
||||
|
||||
@torch.inference_mode()
|
||||
def _top1_confidence(
|
||||
self,
|
||||
logits: torch.Tensor = None
|
||||
):
|
||||
probs = F.softmax(logits, dim=-1) # LxV
|
||||
top1_idx = torch.argmax(probs, dim=-1) # Lx1
|
||||
top1_confs = probs[torch.arange(probs.shape[0], device=probs.device), top1_idx] # Lx1
|
||||
return top1_confs
|
||||
|
||||
@torch.inference_mode()
|
||||
def _extend_w_mask(
|
||||
self,
|
||||
input_ids=None,
|
||||
k_toks=None,
|
||||
mask_id=None,
|
||||
min_mask_id=None,
|
||||
max_mask_id=None
|
||||
):
|
||||
bsz, _ = input_ids.shape
|
||||
|
||||
if k_toks-1 > 0:
|
||||
if min_mask_id is not None:
|
||||
mask_tensor = torch.arange(min_mask_id, min_mask_id + k_toks - 1, dtype=torch.int64, device=input_ids.device).unsqueeze(0).expand(bsz, -1)
|
||||
# check that we didn't insert something larger than max_mask_id
|
||||
assert mask_tensor.max().item() <= max_mask_id, "Inserted mask ID exceeds specified max_mask_id."
|
||||
else:
|
||||
mask_tensor = torch.ones((bsz,k_toks-1),dtype=torch.int64, device=input_ids.device) * mask_id
|
||||
return torch.cat([input_ids, mask_tensor], dim=-1)
|
||||
|
||||
return input_ids
|
||||
|
||||
@torch.inference_mode()
|
||||
def _mtp_next_tokens(
|
||||
self,
|
||||
model,
|
||||
model_inputs,
|
||||
k_toks: int = 1,
|
||||
strategy: Optional[list] = None,
|
||||
) -> torch.Tensor:
|
||||
outputs = model(**model_inputs)
|
||||
# logits = outputs.logits[:, -k_toks:, :]
|
||||
logits = outputs.logits[0, -k_toks:, :]
|
||||
if strategy is None:
|
||||
_next = torch.argmax(logits, dim=-1, keepdim=False)
|
||||
elif strategy[0] == "conf_adapt" or strategy[0] == "conf_adapt_sample@1":
|
||||
# we compute the position wise confidences using the _top1_confidence function
|
||||
top1_conf = self._top1_confidence(logits)
|
||||
# print(f"top1_conf: {top1_conf}", flush=True)
|
||||
# now we compute the position of the farthest token geq the threshold
|
||||
# but contiguously, so if we have [0.95, 0.92, 0.85, 0.97] and threshold 0.9
|
||||
# we want to get position 1 not 3, since position 2 is below the threshold
|
||||
# also being careful of situation like [0.85, 0.88, 0.95] where nothing meets the threshold
|
||||
# falling back to the first token in that case
|
||||
threshold = strategy[1]
|
||||
lt_thresh_mask = top1_conf < threshold
|
||||
# now we find the first case where the mask is true, and go back one position
|
||||
if torch.all(~lt_thresh_mask):
|
||||
# then all positions are above the threshold, we take the last position
|
||||
last_pos = k_toks - 1
|
||||
else:
|
||||
last_pos = torch.argmax(lt_thresh_mask.int()).item() - 1
|
||||
if last_pos < 0:
|
||||
last_pos = 0
|
||||
|
||||
# print(f"last_pos: {last_pos}", flush=True)
|
||||
|
||||
# then we slice the logits to only keep up to that position
|
||||
logits = logits[: last_pos + 1]
|
||||
|
||||
if last_pos == 0 and strategy[0] == "conf_adapt_sample@1":
|
||||
# if k is 1 this step, then draw from the distribution
|
||||
probs = torch.softmax(logits, dim=-1)
|
||||
# print(f"Sampling from probs at k={logits.size(0)}: {probs.shape}", flush=True)
|
||||
temperature = strategy[2]
|
||||
if 0.0 < temperature < 1.0:
|
||||
probs = probs.pow(1.0 / temperature)
|
||||
probs = probs / probs.sum(dim=-1, keepdim=True)
|
||||
else:
|
||||
print(f"Using temperature={temperature} has no effect.", flush=True)
|
||||
_next = torch.multinomial(probs, num_samples=1).squeeze(0)
|
||||
|
||||
else:
|
||||
_next = torch.argmax(logits, dim=-1, keepdim=False)
|
||||
elif strategy[0] == "random":
|
||||
sampling_weights = strategy[1]
|
||||
# we sample k for this step according to the provided weights
|
||||
k_toks = int(
|
||||
torch.multinomial(torch.tensor(sampling_weights), num_samples=1).item()
|
||||
)
|
||||
logits = logits[: k_toks + 1]
|
||||
_next = torch.argmax(logits, dim=-1, keepdim=False)
|
||||
else:
|
||||
raise ValueError(f"Unknown strategy: {strategy}")
|
||||
|
||||
_next = _next.unsqueeze(0) # add batch dim back
|
||||
return _next, outputs
|
||||
|
||||
|
||||
class LlamaForSequenceClassification(GenericForSequenceClassification, LlamaPreTrainedModel): ...
|
||||
|
||||
|
||||
class LlamaForQuestionAnswering(GenericForQuestionAnswering, LlamaPreTrainedModel):
|
||||
base_model_prefix = "transformer" # For BC, where `transformer` was used instead of `model`
|
||||
|
||||
|
||||
class LlamaForTokenClassification(GenericForTokenClassification, LlamaPreTrainedModel): ...
|
||||
|
||||
|
||||
#################################### HF registration ############################################################
|
||||
|
||||
from transformers import AutoConfig, AutoModel, AutoModelForCausalLM
|
||||
|
||||
LlamaConfig.register_for_auto_class()
|
||||
|
||||
LlamaForCausalLM.register_for_auto_class("AutoModel")
|
||||
LlamaForCausalLM.register_for_auto_class("AutoModelForCausalLM")
|
||||
|
||||
__all__ = [
|
||||
"LlamaForCausalLM",
|
||||
"LlamaModel",
|
||||
"LlamaPreTrainedModel",
|
||||
"LlamaForSequenceClassification",
|
||||
"LlamaForQuestionAnswering",
|
||||
"LlamaForTokenClassification",
|
||||
]
|
||||
3
pytorch_model.bin
Normal file
3
pytorch_model.bin
Normal file
@@ -0,0 +1,3 @@
|
||||
version https://git-lfs.github.com/spec/v1
|
||||
oid sha256:885a89f2e07794a47c503a490966a9466a0cc8a153f0a13ed3c6757b83db4972
|
||||
size 3400702686
|
||||
249
special_tokens_map.json
Normal file
249
special_tokens_map.json
Normal file
@@ -0,0 +1,249 @@
|
||||
{
|
||||
"additional_special_tokens": [
|
||||
{
|
||||
"content": "<|mtp_special_token_0|>",
|
||||
"lstrip": false,
|
||||
"normalized": false,
|
||||
"rstrip": false,
|
||||
"single_word": false
|
||||
},
|
||||
{
|
||||
"content": "<|mtp_special_token_1|>",
|
||||
"lstrip": false,
|
||||
"normalized": false,
|
||||
"rstrip": false,
|
||||
"single_word": false
|
||||
},
|
||||
{
|
||||
"content": "<|mtp_special_token_2|>",
|
||||
"lstrip": false,
|
||||
"normalized": false,
|
||||
"rstrip": false,
|
||||
"single_word": false
|
||||
},
|
||||
{
|
||||
"content": "<|mtp_special_token_3|>",
|
||||
"lstrip": false,
|
||||
"normalized": false,
|
||||
"rstrip": false,
|
||||
"single_word": false
|
||||
},
|
||||
{
|
||||
"content": "<|mtp_special_token_4|>",
|
||||
"lstrip": false,
|
||||
"normalized": false,
|
||||
"rstrip": false,
|
||||
"single_word": false
|
||||
},
|
||||
{
|
||||
"content": "<|mtp_special_token_5|>",
|
||||
"lstrip": false,
|
||||
"normalized": false,
|
||||
"rstrip": false,
|
||||
"single_word": false
|
||||
},
|
||||
{
|
||||
"content": "<|mtp_special_token_6|>",
|
||||
"lstrip": false,
|
||||
"normalized": false,
|
||||
"rstrip": false,
|
||||
"single_word": false
|
||||
},
|
||||
{
|
||||
"content": "<|mtp_special_token_7|>",
|
||||
"lstrip": false,
|
||||
"normalized": false,
|
||||
"rstrip": false,
|
||||
"single_word": false
|
||||
},
|
||||
{
|
||||
"content": "<|mtp_special_token_8|>",
|
||||
"lstrip": false,
|
||||
"normalized": false,
|
||||
"rstrip": false,
|
||||
"single_word": false
|
||||
},
|
||||
{
|
||||
"content": "<|mtp_special_token_9|>",
|
||||
"lstrip": false,
|
||||
"normalized": false,
|
||||
"rstrip": false,
|
||||
"single_word": false
|
||||
},
|
||||
{
|
||||
"content": "<|mtp_special_token_10|>",
|
||||
"lstrip": false,
|
||||
"normalized": false,
|
||||
"rstrip": false,
|
||||
"single_word": false
|
||||
},
|
||||
{
|
||||
"content": "<|mtp_special_token_11|>",
|
||||
"lstrip": false,
|
||||
"normalized": false,
|
||||
"rstrip": false,
|
||||
"single_word": false
|
||||
},
|
||||
{
|
||||
"content": "<|mtp_special_token_12|>",
|
||||
"lstrip": false,
|
||||
"normalized": false,
|
||||
"rstrip": false,
|
||||
"single_word": false
|
||||
},
|
||||
{
|
||||
"content": "<|mtp_special_token_13|>",
|
||||
"lstrip": false,
|
||||
"normalized": false,
|
||||
"rstrip": false,
|
||||
"single_word": false
|
||||
},
|
||||
{
|
||||
"content": "<|mtp_special_token_14|>",
|
||||
"lstrip": false,
|
||||
"normalized": false,
|
||||
"rstrip": false,
|
||||
"single_word": false
|
||||
},
|
||||
{
|
||||
"content": "<|mtp_special_token_15|>",
|
||||
"lstrip": false,
|
||||
"normalized": false,
|
||||
"rstrip": false,
|
||||
"single_word": false
|
||||
},
|
||||
{
|
||||
"content": "<|mtp_special_token_16|>",
|
||||
"lstrip": false,
|
||||
"normalized": false,
|
||||
"rstrip": false,
|
||||
"single_word": false
|
||||
},
|
||||
{
|
||||
"content": "<|mtp_special_token_17|>",
|
||||
"lstrip": false,
|
||||
"normalized": false,
|
||||
"rstrip": false,
|
||||
"single_word": false
|
||||
},
|
||||
{
|
||||
"content": "<|mtp_special_token_18|>",
|
||||
"lstrip": false,
|
||||
"normalized": false,
|
||||
"rstrip": false,
|
||||
"single_word": false
|
||||
},
|
||||
{
|
||||
"content": "<|mtp_special_token_19|>",
|
||||
"lstrip": false,
|
||||
"normalized": false,
|
||||
"rstrip": false,
|
||||
"single_word": false
|
||||
},
|
||||
{
|
||||
"content": "<|mtp_special_token_20|>",
|
||||
"lstrip": false,
|
||||
"normalized": false,
|
||||
"rstrip": false,
|
||||
"single_word": false
|
||||
},
|
||||
{
|
||||
"content": "<|mtp_special_token_21|>",
|
||||
"lstrip": false,
|
||||
"normalized": false,
|
||||
"rstrip": false,
|
||||
"single_word": false
|
||||
},
|
||||
{
|
||||
"content": "<|mtp_special_token_22|>",
|
||||
"lstrip": false,
|
||||
"normalized": false,
|
||||
"rstrip": false,
|
||||
"single_word": false
|
||||
},
|
||||
{
|
||||
"content": "<|mtp_special_token_23|>",
|
||||
"lstrip": false,
|
||||
"normalized": false,
|
||||
"rstrip": false,
|
||||
"single_word": false
|
||||
},
|
||||
{
|
||||
"content": "<|mtp_special_token_24|>",
|
||||
"lstrip": false,
|
||||
"normalized": false,
|
||||
"rstrip": false,
|
||||
"single_word": false
|
||||
},
|
||||
{
|
||||
"content": "<|mtp_special_token_25|>",
|
||||
"lstrip": false,
|
||||
"normalized": false,
|
||||
"rstrip": false,
|
||||
"single_word": false
|
||||
},
|
||||
{
|
||||
"content": "<|mtp_special_token_26|>",
|
||||
"lstrip": false,
|
||||
"normalized": false,
|
||||
"rstrip": false,
|
||||
"single_word": false
|
||||
},
|
||||
{
|
||||
"content": "<|mtp_special_token_27|>",
|
||||
"lstrip": false,
|
||||
"normalized": false,
|
||||
"rstrip": false,
|
||||
"single_word": false
|
||||
},
|
||||
{
|
||||
"content": "<|mtp_special_token_28|>",
|
||||
"lstrip": false,
|
||||
"normalized": false,
|
||||
"rstrip": false,
|
||||
"single_word": false
|
||||
},
|
||||
{
|
||||
"content": "<|mtp_special_token_29|>",
|
||||
"lstrip": false,
|
||||
"normalized": false,
|
||||
"rstrip": false,
|
||||
"single_word": false
|
||||
},
|
||||
{
|
||||
"content": "<|mtp_special_token_30|>",
|
||||
"lstrip": false,
|
||||
"normalized": false,
|
||||
"rstrip": false,
|
||||
"single_word": false
|
||||
},
|
||||
{
|
||||
"content": "<|mtp_special_token_31|>",
|
||||
"lstrip": false,
|
||||
"normalized": false,
|
||||
"rstrip": false,
|
||||
"single_word": false
|
||||
}
|
||||
],
|
||||
"bos_token": {
|
||||
"content": "<|begin_of_text|>",
|
||||
"lstrip": false,
|
||||
"normalized": false,
|
||||
"rstrip": false,
|
||||
"single_word": false
|
||||
},
|
||||
"eos_token": {
|
||||
"content": "<|end_of_text|>",
|
||||
"lstrip": false,
|
||||
"normalized": false,
|
||||
"rstrip": false,
|
||||
"single_word": false
|
||||
},
|
||||
"pad_token": {
|
||||
"content": "<|finetune_right_pad_id|>",
|
||||
"lstrip": false,
|
||||
"normalized": false,
|
||||
"rstrip": false,
|
||||
"single_word": false
|
||||
}
|
||||
}
|
||||
3
tokenizer.json
Normal file
3
tokenizer.json
Normal file
@@ -0,0 +1,3 @@
|
||||
version https://git-lfs.github.com/spec/v1
|
||||
oid sha256:195adb4002a7756f15ec76beeaf85e29c969d45cea2670574d8e705369f9574a
|
||||
size 17216342
|
||||
2356
tokenizer_config.json
Normal file
2356
tokenizer_config.json
Normal file
File diff suppressed because it is too large
Load Diff
Reference in New Issue
Block a user