crumbs-playground/27b-3bit-testing

from transformers import AutoConfig, AutoModelForCausalLM, AutoTokenizer
import torch

model_id = "crumbs-playground/27b-3bit-testing"
tokenizer = AutoTokenizer.from_pretrained(model_id)
model     = AutoModelForCausalLM.from_pretrained(
    model_id, dtype=torch.bfloat16, 
    device_map="cuda",
    trust_remote_code=True,
)

model.add_adapter("my_lora_adapter", r=32, alpha=32**0.5, trainable=True)
model.generation_config.cache_implementation = "static"
torch.compile(model._c23_base_model, mode="reduce-overhead", fullgraph=False)

with model.disable_adapters(["my_lora_adapter"]): # <- disable loras if you want!
    inputs = tokenizer("Once upon a time,", return_tensors="pt")
    inputs = {k:v.to("cuda:0") for k,v in inputs.items()}
    outputs = model.generate(
        **inputs, do_sample=True,
        max_new_tokens=32,
        temperature=1.0,
        top_k=0,
    )
    print(tokenizer.decode(outputs[0]))


# reinforcement learning!
advantage = 10.0 # <- what it generated was awesome

optimizer = torch.optim.AdamW(
    [p for n,p in model.named_parameters() if "my_lora_adapter" in n],
    3e-5, weight_decay=0.1, betas=(0.9, 0.95)
)
logits = model(outputs.clone().detach()).logits
logits = logits[:, inputs['input_ids'].shape[-1]:, :]
labels = outputs[:, inputs['input_ids'].shape[-1]:].clone()
loss = torch.nn.functional.cross_entropy(
    logits.reshape(-1, logits.size(-1)),
    labels.reshape(-1),
    reduction="none",
).view(labels.shape)
policy_loss = (-loss * -advantage).mean()
policy_loss.backward()

grad_norm = torch.nn.utils.clip_grad_norm_(model.parameters(), 0.5)
optimizer.step()

# save your adapters to the huggingface hub!
model.push_to_hub("storybook", "me/CHQQ-adapter-proof-of-concept")

# load other people's uploaded adapters from the hub!
model.load_adapter(
    adapter_name="custom_adapter", 
    load_directory="someone_else/CHQQ-proof-of-concept"
)