Views
No views yet
1{
2 "_name_or_path": "output/hermes-llama2-4k/checkpoint-2259",
3 "architectures": [
4 "LlamaForCausalLM"
5 ],
6 "bos_token_id": 1,
7 "eos_token_id": 2,
8 "hidden_act": "silu",
9 "hidden_size": 4096,
10 "initializer_range": 0.02,
11 "intermediate_size": 11008,
12 "max_position_embeddings": 4096,
13 "model_type": "llama",
14 "num_attention_heads": 32,
15 "num_hidden_layers": 32,
16 "num_key_value_heads": 32,
17 "pad_token_id": 0,
18 "pretraining_tp": 1,
19 "rms_norm_eps": 1e-05,
20 "rope_scaling": null,
21 "tie_word_embeddings": false,
22 "torch_dtype": "bfloat16",
23 "transformers_version": "4.32.0.dev0",
24 "use_cache": false,
25 "vocab_size": 32000
26}DATASET = "abideen/Cosmopedia-100k-pretrain" # @param
from datasets import load_dataset
# converted to BitLinear
class BitLinear(nn.Linear):
def forward(self, x):
w = self.weight # a weight tensor with shape [d, k]
x = x.to(w.device)
RMSNorm = LlamaRMSNorm(x.shape[-1]).to(w.device)
x_norm = RMSNorm(x)
# A trick for implementing Straight−Through−Estimator (STE) using detach()
x_quant = x_norm + (activation_quant(x_norm) - x_norm).detach()
w_quant = w + (weight_quant(w) - w).detach()
y = F.linear(x_quant, w_quant)
return y
### Create the llama model with our custom config. Convert it to bitnet.
model = LlamaForCausalLM(config)
convert_to_bitnet(model, copy_weights=False)1args = TrainingArguments(
2 output_dir=output_path,
3 per_device_train_batch_size=BATCH_SIZE,
4 logging_steps=100,
5 gradient_accumulation_steps=2,
6 num_train_epochs=EPOCHS,
7 weight_decay=0.01,
8 warmup_steps=0.1,
9 lr_scheduler_type="cosine",
10 learning_rate=LEARNING_RATE,
11 # max_steps=5000,
12 save_steps=0.25,
13 fp16=True,
14 report_to="wandb"
15)
16
17trainer = Trainer(
18 model=model,
19 tokenizer=tokenizer,
20 args=args,
21 data_collator=data_collator,
22 train_dataset=tokenized_data["train"],
23)
24
25trainer.train()1from transformers import AutoModelForCausalLM, AutoTokenizer
2from transformers.models.llama.modeling_llama import *
3# Load a pretrained BitNet model
4model = "saadnaeem/Llama2-70M-Cosmopedia-100k-Pretrain"
5tokenizer = AutoTokenizer.from_pretrained(model)
6model = AutoModelForCausalLM.from_pretrained(model)
7
8
9def activation_quant(x):
10 scale = 127.0 / x.abs().max(dim=-1, keepdim=True).values.clamp_(min=1e-5)
11 y = (x * scale).round().clamp_(-128, 127)
12 y = y / scale
13 return y
14def weight_quant(w):
15 scale = 1.0 / w.abs().mean().clamp_(min=1e-5)
16 u = (w * scale).round().clamp_(-1, 1)
17 u = u / scale
18 return u
19
20class BitLinear(nn.Linear):
21 def forward(self, x):
22 w = self.weight # a weight tensor with shape [d, k]
23 x = x.to(w.device)
24 RMSNorm = LlamaRMSNorm(x.shape[-1]).to(w.device)
25 x_norm = RMSNorm(x)
26 # A trick for implementing Straight−Through−Estimator (STE) using detach()
27 x_quant = x_norm + (activation_quant(x_norm) - x_norm).detach()
28 w_quant = w + (weight_quant(w) - w).detach()
29 y = F.linear(x_quant, w_quant)
30 return y
31
32def convert_to_bitnet(model, copy_weights):
33 for name, module in model.named_modules():
34 # Replace linear layers with BitNet
35 if isinstance(module, LlamaSdpaAttention) or isinstance(module, LlamaMLP):
36 for child_name, child_module in module.named_children():
37 if isinstance(child_module, nn.Linear):
38 bitlinear = BitLinear(child_module.in_features, child_module.out_features, child_module.bias is not None).to(device="cuda:0")
39 if copy_weights:
40 bitlinear.weight = child_module.weight
41 if child_module.bias is not None:
42 bitlinear.bias = child_module.bias
43 setattr(module, child_name, bitlinear)
44 # Remove redundant input_layernorms
45 elif isinstance(module, LlamaDecoderLayer):
46 for child_name, child_module in module.named_children():
47 if isinstance(child_module, LlamaRMSNorm) and child_name == "input_layernorm":
48 setattr(module, child_name, nn.Identity().to(device="cuda:0"))
49
50
51convert_to_bitnet(model, copy_weights=True)
52model.to(device="cuda:0")
53
54prompt = "What is Machine Learning?"
55inputs = tokenizer(prompt, return_tensors="pt").to(model.device)
56generate_ids = model.generate(inputs.input_ids, max_length=50)
57tokenizer.batch_decode(generate_ids, skip_special_tokens=True, clean_up_tokenization_spaces=False)[0]