Views
No views yet
KernelBookSDFTTrainer on top of transformers.Trainer with DeepSpeed ZeRO-3.1from transformers import AutoModelForCausalLM, AutoTokenizer
2
3model_id = "aadityabuilds/qwen2-5-coder-7b-kernelbook-sdft"
4tokenizer = AutoTokenizer.from_pretrained(model_id, trust_remote_code=True)
5model = AutoModelForCausalLM.from_pretrained(
6 model_id, torch_dtype="auto", device_map="auto", trust_remote_code=True
7)
8
9messages = [
10 {
11 "role": "user",
12 "content": "Convert the following PyTorch code to an equivalent Triton kernel...",
13 }
14]
15prompt = tokenizer.apply_chat_template(messages, tokenize=False, add_generation_prompt=True)
16inputs = tokenizer(prompt, return_tensors="pt").to(model.device)
17outputs = model.generate(**inputs, max_new_tokens=1200, do_sample=False)
18print(tokenizer.decode(outputs[0][inputs["input_ids"].shape[1] :], skip_special_tokens=True))| Setting | Value |
|---|---|
| Base model | Qwen2.5-Coder-7B-Instruct |
| Method | SDFT (forced-completion distillation) |
| Epochs | 1 |
| Hardware | 4× H100 (Modal) |
| Parallelism | DeepSpeed ZeRO-3, bf16 |