initial example script exploration (only pirate script tested)

Author: Titus-von-Koeller

examples/compile_inference.py

@@ -0,0 +1,59 @@
+import logging
+
+import torch
+import torch._dynamo
+
+# Enable verbose logging using PyTorch's new logging system
+from torch._logging import set_logs
+from transformers import AutoModelForCausalLM, AutoTokenizer, BitsAndBytesConfig
+
+set_logs(dynamo=logging.DEBUG, graph_breaks=True, recompiles=True)  # Enable specific artifacts
+
+torch.set_float32_matmul_precision("high")
+
+quantization_config = BitsAndBytesConfig(load_in_8bit=True)
+
+model_id = "google/gemma-2-2b-it"
+tokenizer = AutoTokenizer.from_pretrained(model_id)
+model = AutoModelForCausalLM.from_pretrained(
+    model_id,
+    quantization_config=quantization_config,
+    device_map="auto",
+    torch_dtype=torch.bfloat16,
+)
+
+input_text = "Write me a poem about Machine Learning."
+input_ids = tokenizer(input_text, return_tensors="pt").to(model.device)
+
+# Improved explanation with artifact capture
+from torch._dynamo import explain
+
+explanation, graphs, guards, log_location = explain(
+    model,
+    **input_ids,
+    traceback=True,  # Show origin of graph breaks
+)
+print(f"Graph breaks: {len(explanation)}\nDetails:")
+for break_info in explanation:
+    print(f"- {break_info}")
+
+
+# Enhanced AOTAutograd tracing
+def trace_handler(gm: torch.fx.GraphModule, example_inputs):
+    print(f"AOT traced graph with {len(gm.graph.nodes)} nodes")
+    gm.graph.print_tabular()  # Show full graph structure
+    return gm.forward
+
+
+model = torch.compile(
+    model,
+    backend="aot_eager",
+    options={
+        "fw_compiler": trace_handler,
+        "track_graph_metrics": True,  # Additional compilation metrics
+    },
+)
+
+# Generate with compilation diagnostics
+outputs = model.generate(**input_ids, max_new_tokens=32)
+print(tokenizer.decode(outputs[0]))

examples/compile_inference.py.tmp

@@ -0,0 +1,108 @@
+# for torch.compile trace run:
+#
+# TORCH_TRACE="./tracedir" TORCH_LOGS="graph_breaks" CUDA_VISIBLE_DEVICES=0 python examples/compile_pirate_qlora.py
+
+# 🏴☠️⛵ Pirate Coder's Delight: Fine-Tune Mistral-7B to Speak Like a Buccaneer Hacker
+# Using bitsandbytes 4-bit + torch.compile - 100% Original Dataset
+
+from datasets import Dataset
+from peft import LoraConfig
+import torch
+from transformers import AutoModelForCausalLM, AutoTokenizer
+from trl import SFTConfig, SFTTrainer
+
+# 1. Load Model with Pirate-Optimized Quantization 🏴☠️
+model_id = "Qwen/Qwen1.5-1.8B"
+tokenizer = AutoTokenizer.from_pretrained(model_id)
+tokenizer.pad_token = tokenizer.eos_token
+
+model = AutoModelForCausalLM.from_pretrained(
+    model_id,
+    quantization_config={
+        "load_in_4bit": True,
+        "bnb_4bit_quant_type": "nf4",
+        "bnb_4bit_compute_dtype": torch.bfloat16,
+        "bnb_4bit_use_double_quant": True,
+    },
+    device_map="auto",
+    attn_implementation="flash_attention_2",
+    torch_dtype=torch.bfloat16,
+)
+
+
+# 2. Original Pirate Programmer Dataset 🦜
+def pirate_formatting_func(example):
+    return {"text": f"### Instruction:\n{example['instruction']}\n\n### Response:\n{example['response']}"}
+
+
+train_dataset = Dataset.from_list(
+    [
+        {
+            "instruction": "Explain quantum computing using pirate slang",
+            "response": "Arrr, matey! 'Tis like sailin' parallel seas...",
+        },
+        {
+            "instruction": "Write Python code to find buried treasure",
+            "response": "def find_booty():\n    return (sum(coordinates) / len(coordinates))",
+        },
+        {
+            "instruction": "Why do pirates hate distributed systems?",
+            "response": "Too many captains sink the ship, ye scallywag!",
+        },
+    ]
+).map(pirate_formatting_func)
+
+
+# 2. Prepare Pirate Dataset
+def tokenize_pirate_data(examples):
+    return tokenizer(examples["text"], padding="max_length", max_length=256, truncation=True, return_tensors="pt")
+
+
+train_dataset = train_dataset.map(
+    tokenize_pirate_data,
+    batched=True,
+    remove_columns=["instruction", "response"],  # Keep only tokenized fields
+)
+
+# 3. Configure QLoRA ⚙️
+peft_config = LoraConfig(
+    r=32,
+    lora_alpha=64,
+    target_modules=["q_proj", "v_proj"],
+    lora_dropout=0.05,
+    bias="none",
+    task_type="CAUSAL_LM",
+)
+
+# 5. Training Configuration
+training_args = SFTConfig(
+    per_device_train_batch_size=2,
+    gradient_accumulation_steps=1,
+    max_steps=5,
+    learning_rate=2e-5,
+    max_seq_length=256,
+    remove_unused_columns=False,
+    output_dir="./pirate_coder",
+    optim="paged_adamw_8bit",
+    dataset_text_field="text",
+    packing=True,
+    torch_compile={
+        "mode": "reduce-overhead",
+        "fullgraph": False,
+        "dynamic": False,
+    },
+    report_to="none",
+    logging_steps=1,
+)
+
+# 6. Launch Training with Pirate Flair! 🚀
+trainer = SFTTrainer(
+    model=model,
+    train_dataset=train_dataset,
+    args=training_args,
+    peft_config=peft_config,
+    formatting_func=pirate_formatting_func,
+)
+
+print("⚡ Batten down the hatches - training with torch.compile!")
+trainer.train()

examples/compile_pirate_qlora.py

@@ -0,0 +1,131 @@
+# 🦄 Fine-Tune Llama-3.8B to Write Pirate Jokes & Shakespearean Sonnets
+# Using bitsandbytes 4-bit + torch.compile 🏴☠️
+
+# !pip install -qU "unsloth[colab] @ git+https://github.com/unslothai/unsloth.git" \
+#     transformers==4.40.0 accelerate==0.30.0 bitsandbytes==0.43.0
+
+from datasets import load_dataset
+import torch
+from unsloth import FastLanguageModel
+
+# 1. Load Pre-Quantized Model with bitsandbytes 🎯
+model, tokenizer = FastLanguageModel.from_pretrained(
+    model_name="unsloth/llama-3-8B-bnb-4bit",
+    max_seq_length=2048,
+    dtype=torch.bfloat16 if torch.cuda.is_bf16_supported() else torch.float16,
+    load_in_4bit=True,
+    quantization_config={
+        "bnb_4bit_quant_type": "nf4",
+        "bnb_4bit_compute_dtype": torch.bfloat16,
+        "bnb_4bit_use_double_quant": True,
+    },
+)
+
+# 2. Prepare Creative Dataset 🎭
+pirate_dataset = load_dataset(
+    "json",
+    data_files={"train": "https://huggingface.co/datasets/jondurbin/pirate-jokes/resolve/main/pirate_jokes.json"},
+)
+
+
+def format_creative_prompt(sample):
+    return f"""<|begin_of_text|><|start_header_id|>system<|end_header_id|>
+    You are a swashbuckling pirate poet. Respond ONLY in pirate speak or Shakespearean verse.<|eot_id|>
+    <|start_header_id|>user<|end_header_id|>
+    {sample['prompt']}<|eot_id|>
+    <|start_header_id|>assistant<|end_header_id|>
+    {sample['response']}<|eot_id|>"""
+
+
+dataset = pirate_dataset.map(format_creative_prompt)
+
+# 3. Configure QLoRA with torch.compile Diagnostics 🔍
+model = FastLanguageModel.get_peft_model(
+    model,
+    r=32,
+    target_modules=["q_proj", "k_proj", "v_proj", "o_proj"],
+    lora_alpha=64,
+    lora_dropout=0.1,
+    bias="none",
+    use_gradient_checkpointing=True,
+    torch_compile={
+        "mode": "reduce-overhead",
+        "fullgraph": False,  # Allow partial compilation initially
+        "dynamic": True,
+    },
+    random_state=3407,  # For reproducibility of creative outputs
+)
+
+
+# 4. Custom Training Loop with Graph Break Analysis 🕵️♂️
+def detect_graph_breaks():
+    import torch._dynamo
+
+    original_verbose = torch._dynamo.config.verbose
+    torch._dynamo.config.verbose = True
+
+    # Trigger compilation with sample input
+    sample_input = tokenizer("Arrr! Tell me about yer treasure...", return_tensors="pt").to("cuda")
+    compiled_model = torch.compile(model)
+    _ = compiled_model(**sample_input)
+
+    torch._dynamo.config.verbose = original_verbose
+
+
+detect_graph_breaks()  # Initial graph break detection
+
+
+# 5. Creative Generation During Training 🎩
+class PirateStreamer:
+    def __init__(self, tokenizer):
+        self.tokenizer = tokenizer
+        self.prompt = "Yarrr! Why did the pirate's chicken cross the road?"
+
+    def __call__(self, input_ids, *args, **kwargs):
+        if random.random() < 0.1:  # 10% chance to generate during training
+            with torch.no_grad():
+                outputs = model.generate(
+                    input_ids=self.tokenizer(self.prompt, return_tensors="pt").to("cuda").input_ids,
+                    max_new_tokens=50,
+                    temperature=0.7,
+                    repetition_penalty=1.1,
+                )
+            print("\n🏴☠️ Crew's Update:", self.tokenizer.decode(outputs[0]))
+        return input_ids
+
+
+# 6. Launch Training with Progressive Compilation 🚀
+trainer = SFTTrainer(
+    model=model,
+    train_dataset=dataset,
+    dataset_text_field="text",
+    max_seq_length=1024,
+    packing=True,
+    callbacks=[PirateStreamer(tokenizer)],
+    args=TrainingArguments(
+        per_device_train_batch_size=2,
+        gradient_accumulation_steps=4,
+        warmup_steps=10,
+        max_steps=100,
+        learning_rate=3e-5,
+        fp16=not torch.cuda.is_bf16_supported(),
+        bf16=torch.cuda.is_bf16_supported(),
+        logging_steps=1,
+        optim="paged_adamw_8bit",
+        weight_decay=0.01,
+        lr_scheduler_type="cosine",
+        output_dir="pirate-poet",
+        report_to="none",
+    ),
+)
+
+# Progressive compilation strategy
+trainer.train_step = torch.compile(
+    trainer.train_step,
+    mode="reduce-overhead",
+    fullgraph=False,  # Start with partial graphs
+    dynamic=True,
+)
+
+print("🏁 Starting training - watch for graph breaks and pirate wisdom!")
+trainer.train()