amazingvmsloth 0.1.0__tar.gz

amazingvmsloth-0.1.0/PKG-INFO

@@ -0,0 +1,25 @@
+Metadata-Version: 2.4
+Name: amazingvmsloth
+Version: 0.1.0
+Summary: Blazing-fast LLM fine-tuning with minimal VRAM — multi-GPU, manual LoRA gradients, flash attention, 4-bit quant
+License: MIT
+Requires-Python: >=3.9
+Description-Content-Type: text/markdown
+Requires-Dist: torch>=2.1.0
+Requires-Dist: transformers>=4.36.0
+Requires-Dist: bitsandbytes>=0.41.0
+Requires-Dist: peft>=0.7.0
+Requires-Dist: triton>=2.1.0; sys_platform == "linux"
+Requires-Dist: safetensors>=0.4.0
+Requires-Dist: accelerate>=0.25.0
+Requires-Dist: datasets>=2.14.0
+Requires-Dist: psutil>=5.9.0
+Provides-Extra: flash-attn
+Requires-Dist: flash-attn>=2.3.0; extra == "flash-attn"
+Provides-Extra: multi-gpu
+Requires-Dist: deepspeed>=0.12.0; extra == "multi-gpu"
+Provides-Extra: dev
+Requires-Dist: pytest>=7.0; extra == "dev"
+Requires-Dist: ruff>=0.1.0; extra == "dev"
+Provides-Extra: all
+Requires-Dist: amazingvmsloth[dev,flash-attn,multi-gpu]; extra == "all"

amazingvmsloth-0.1.0/amazingvmsloth/__init__.py

@@ -0,0 +1,25 @@
+import os
+
+# Suppress transformers loading progress bars that can crash on Windows terminal
+os.environ.setdefault("HF_HUB_DISABLE_PROGRESS_BARS", "1")
+
+from amazingvmsloth.lora import apply_lora, LoRAConfig
+from amazingvmsloth.trainer import AmazingTrainer, TrainingConfig
+from amazingvmsloth.quantization import quantize_model_4bit
+from amazingvmsloth.attention import patch_attention
+from amazingvmsloth.models import auto_patch_model
+from amazingvmsloth.multi_gpu import setup_distributed
+from amazingvmsloth.offload import apply_dispatch_offload, estimate_model_size_gb
+from amazingvmsloth.cpu_trainer import CpuTrainer, CpuTrainingConfig
+
+__version__ = "0.1.0"
+__all__ = [
+    "apply_lora",
+    "LoRAConfig",
+    "AmazingTrainer",
+    "TrainingConfig",
+    "quantize_model_4bit",
+    "patch_attention",
+    "auto_patch_model",
+    "setup_distributed",
+]

amazingvmsloth-0.1.0/amazingvmsloth/attention.py

@@ -0,0 +1,236 @@
+import torch
+import torch.nn as nn
+import torch.nn.functional as F
+from typing import Optional, Tuple
+from transformers import PreTrainedModel
+
+_flash_available = False
+try:
+    from flash_attn import flash_attn_func
+    _flash_available = True
+except ImportError:
+    pass
+
+_xformers_available = False
+try:
+    from xformers.ops import memory_efficient_attention
+    _xformers_available = True
+except ImportError:
+    pass
+
+
+def is_flash_available() -> bool:
+    return _flash_available
+
+
+class EfficientAttention(nn.Module):
+    def __init__(self, original_module, config=None):
+        super().__init__()
+        self.original_module = original_module
+        self.config = config
+        self._use_flash = False
+        self.num_heads = getattr(original_module, "num_heads", 32)
+        self.num_kv_heads = getattr(original_module, "num_kv_groups", None) or getattr(
+            original_module, "num_key_value_heads", self.num_heads
+        )
+        self.head_dim = getattr(original_module, "head_dim", 128)
+        self.hidden_size = self.num_heads * self.head_dim
+
+    def _try_flash_attention(self, q, k, v, attention_mask=None):
+        if not _flash_available:
+            return None
+        if q.dtype != torch.float16 and q.dtype != torch.bfloat16:
+            return None
+        if q.size(1) > 8192 and q.device.type != "cuda":
+            return None
+
+        try:
+            seq_len = q.size(1)
+            q_4d = q.view(q.size(0), seq_len, self.num_heads, self.head_dim)
+            k_4d = k.view(k.size(0), k.size(1), self.num_kv_heads, self.head_dim)
+            v_4d = v.view(v.size(0), v.size(1), self.num_kv_heads, self.head_dim)
+
+            if self.num_kv_heads != self.num_heads:
+                n_rep = self.num_heads // self.num_kv_heads
+                k_4d = k_4d.unsqueeze(2).expand(-1, -1, n_rep, -1, -1).reshape(q.size(0), -1, self.num_heads, self.head_dim)
+                v_4d = v_4d.unsqueeze(2).expand(-1, -1, n_rep, -1, -1).reshape(q.size(0), -1, self.num_heads, self.head_dim)
+
+            output = flash_attn_func(q_4d, k_4d, v_4d, causal=True)
+            return output.view(q.size(0), seq_len, -1)
+        except Exception:
+            return None
+
+    def _sdpa_attention(self, q, k, v, attention_mask=None):
+        if hasattr(F, "scaled_dot_product_attention"):
+            is_causal = attention_mask is None
+            output = F.scaled_dot_product_attention(
+                q, k, v,
+                attn_mask=attention_mask,
+                is_causal=is_causal,
+            )
+            return output
+        return None
+
+    def forward(self, hidden_states, attention_mask=None, position_ids=None, **kwargs):
+        return self.original_module(
+            hidden_states,
+            attention_mask=attention_mask,
+            position_ids=position_ids,
+            **kwargs,
+        )
+
+
+class PatchedLlamaAttention(nn.Module):
+    def __init__(self, original_attn, config=None):
+        super().__init__()
+        self.q_proj = original_attn.q_proj
+        self.k_proj = original_attn.k_proj
+        self.v_proj = original_attn.v_proj
+        self.o_proj = original_attn.o_proj
+        self.rotary_emb = original_attn.rotary_emb
+        self.config = config
+
+        model_config = getattr(original_attn, "config", None)
+        self.hidden_size = getattr(model_config, "hidden_size", 4096)
+        self.num_heads = getattr(model_config, "num_attention_heads", 32)
+        self.head_dim = getattr(model_config, "head_dim", self.hidden_size // self.num_heads)
+        self.num_key_value_heads = getattr(model_config, "num_key_value_heads", self.num_heads)
+        self.num_key_value_groups = self.num_heads // self.num_key_value_heads
+        self.max_position_embeddings = getattr(model_config, "max_position_embeddings", 4096)
+
+    def _shape(self, tensor, seq_len, bsz):
+        return tensor.view(bsz, seq_len, self.num_heads, self.head_dim).transpose(1, 2)
+
+    def forward(
+        self,
+        hidden_states: torch.Tensor,
+        attention_mask: Optional[torch.Tensor] = None,
+        position_ids: Optional[torch.LongTensor] = None,
+        past_key_value=None,
+        output_attentions: bool = False,
+        use_cache: bool = False,
+        **kwargs,
+    ) -> Tuple[torch.Tensor, Optional[torch.Tensor], Optional[Tuple[torch.Tensor]]]:
+        bsz, q_len, _ = hidden_states.size()
+
+        query_states = self.q_proj(hidden_states)
+        key_states = self.k_proj(hidden_states)
+        value_states = self.v_proj(hidden_states)
+
+        query_states = query_states.view(bsz, q_len, self.num_heads, self.head_dim).transpose(1, 2)
+        key_states = key_states.view(bsz, q_len, self.num_key_value_heads, self.head_dim).transpose(1, 2)
+        value_states = value_states.view(bsz, q_len, self.num_key_value_heads, self.head_dim).transpose(1, 2)
+
+        cos, sin = self.rotary_emb(value_states, position_ids)
+        query_states, key_states = _apply_rotary_pos_emb(query_states, key_states, cos, sin)
+
+        if past_key_value is not None:
+            key_states = torch.cat([past_key_value[0], key_states], dim=2)
+            value_states = torch.cat([past_key_value[1], value_states], dim=2)
+        past_key_value = (key_states, value_states) if use_cache else None
+
+        if self.num_key_value_groups > 1:
+            key_states = _repeat_kv(key_states, self.num_key_value_groups)
+            value_states = _repeat_kv(value_states, self.num_key_value_groups)
+
+        attn_output = _efficient_attn(query_states, key_states, value_states, attention_mask)
+
+        attn_output = attn_output.transpose(1, 2).contiguous()
+        attn_output = attn_output.reshape(bsz, q_len, self.hidden_size)
+        attn_output = self.o_proj(attn_output)
+
+        if not output_attentions:
+            attn_weights = None
+
+        return attn_output, attn_weights, past_key_value
+
+
+def _apply_rotary_pos_emb(q, k, cos, sin):
+    def _rotate_half(x):
+        x1 = x[..., : x.shape[-1] // 2]
+        x2 = x[..., x.shape[-1] // 2 :]
+        return torch.cat((-x2, x1), dim=-1)
+
+    q_embed = (q * cos) + (_rotate_half(q) * sin)
+    k_embed = (k * cos) + (_rotate_half(k) * sin)
+    return q_embed, k_embed
+
+
+def _repeat_kv(hidden_states, n_rep):
+    if n_rep == 1:
+        return hidden_states
+    bsz, num_heads, seq_len, head_dim = hidden_states.shape
+    hidden_states = hidden_states[:, :, None, :, :].expand(bsz, num_heads, n_rep, seq_len, head_dim)
+    return hidden_states.reshape(bsz, num_heads * n_rep, seq_len, head_dim)
+
+
+def _efficient_attn(q, k, v, mask=None):
+    bsz, n_heads, seq_len, head_dim = q.shape
+
+    if _flash_available and q.is_cuda and q.dtype in (torch.float16, torch.bfloat16):
+        try:
+            q_4d = q.transpose(1, 2)
+            k_4d = k.transpose(1, 2)
+            v_4d = v.transpose(1, 2)
+            out = flash_attn_func(q_4d, k_4d, v_4d, causal=True)
+            return out.transpose(1, 2)
+        except Exception:
+            pass
+
+    if _xformers_available and q.is_cuda and q.dtype in (torch.float16, torch.bfloat16):
+        try:
+            q_3d = q.transpose(1, 2).reshape(bsz * seq_len, n_heads, head_dim)
+            k_3d = k.transpose(1, 2).reshape(bsz * seq_len, n_heads, head_dim)
+            v_3d = v.transpose(1, 2).reshape(bsz * seq_len, n_heads, head_dim)
+            attn_bias = mask if mask is not None else None
+            out = memory_efficient_attention(q_3d, k_3d, v_3d, attn_bias=attn_bias)
+            return out.reshape(bsz, seq_len, n_heads, head_dim).transpose(1, 2)
+        except Exception:
+            pass
+
+    if hasattr(F, "scaled_dot_product_attention"):
+        is_causal = mask is None
+        return F.scaled_dot_product_attention(q, k, v, attn_mask=mask, is_causal=is_causal)
+
+    attn_weights = torch.matmul(q, k.transpose(2, 3)) / (q.shape[-1] ** 0.5)
+    if mask is not None:
+        attn_weights = attn_weights + mask
+    attn_weights = F.softmax(attn_weights, dim=-1, dtype=torch.float32).to(q.dtype)
+    return torch.matmul(attn_weights, v)
+
+
+def patch_attention(model: PreTrainedModel) -> PreTrainedModel:
+    if hasattr(model, "config"):
+        model.config._attn_implementation = "sdpa"
+
+    modules_dict = dict(model.named_modules())
+    patched_count = 0
+
+    for name, module in list(modules_dict.items()):
+        module_type = type(module).__name__
+        if "Attention" in module_type or "SdpaAttention" in module_type:
+            has_rotary = hasattr(module, "rotary_emb")
+            if not has_rotary:
+                continue
+            try:
+                patched = PatchedLlamaAttention(module, getattr(model, "config", None))
+                parts = name.rsplit(".", 1)
+                if len(parts) == 2:
+                    parent = modules_dict.get(parts[0])
+                    if parent is not None:
+                        setattr(parent, parts[1], patched)
+                        patched_count += 1
+            except Exception as e:
+                print(f"[amazingvmsloth] Could not patch attention {name}: {e}")
+
+    skipped = sum(1 for name, m in modules_dict.items()
+                  if ("Attention" in type(m).__name__ or "SdpaAttention" in type(m).__name__)
+                  and not hasattr(m, "rotary_emb"))
+
+    print(f"[amazingvmsloth] Patched {patched_count} attention layers, {skipped} using native SDPA (no rotary_emb on module)")
+    flash_status = "available" if _flash_available else "not available (install flash-attn for speedup)"
+    xformers_status = "available" if _xformers_available else "not available (pip install xformers for speedup)"
+    print(f"[amazingvmsloth] Flash attention: {flash_status}")
+    print(f"[amazingvmsloth] XFormers attention: {xformers_status}")
+    print(f"[amazingvmsloth] Attention implementation: {'xformers' if _xformers_available else 'SDPA'}")
+    return model

amazingvmsloth-0.1.0/amazingvmsloth/bench.py

@@ -0,0 +1,291 @@
+import time
+import torch
+import argparse
+import os
+from typing import Dict, Any
+
+# Suppress transformers loading progress bars (can crash on Windows + cause UI issues)
+os.environ.setdefault("HF_HUB_DISABLE_PROGRESS_BARS", "1")
+
+_unsloth_available = False
+_unsloth_error = None
+
+
+def _try_import_unsloth():
+    """Lazy import unsloth only when unsloth benchmark is requested."""
+    global _unsloth_available, _unsloth_error
+    if _unsloth_available or _unsloth_error is not None:
+        return _unsloth_available
+    import warnings
+    with warnings.catch_warnings():
+        warnings.simplefilter("ignore", UserWarning)
+        try:
+            import unsloth
+            from unsloth import FastLanguageModel
+            _unsloth_available = True
+        except Exception as e:
+            _unsloth_error = str(e)
+    return _unsloth_available
+
+
+def run_amazingvmsloth(model_name: str, dataset_name: str, max_samples: int, max_seq_length: int) -> Dict[str, Any]:
+    from transformers import AutoTokenizer
+    from datasets import load_dataset
+    from amazingvmsloth import quantize_model_4bit, auto_patch_model, LoRAConfig, AmazingTrainer, TrainingConfig
+    from amazingvmsloth.utils.banner import print_banner, get_system_info
+
+    info = get_system_info()
+    print_banner(info, model_name)
+
+    tokenizer = AutoTokenizer.from_pretrained(model_name)
+    if tokenizer.pad_token is None:
+        tokenizer.pad_token = tokenizer.eos_token
+
+    try:
+        model = quantize_model_4bit(model_name)
+    except (torch.OutOfMemoryError, RuntimeError) as e:
+        if "out of memory" in str(e).lower() or "cuda" in str(e).lower():
+            print(f"\n[bench] ERROR: GPU out of memory while loading {model_name}")
+            print(f"[bench] This model may be too large for your GPU.")
+            print(f"[bench] Try: --model Qwen/Qwen2.5-0.5B (smaller model)")
+            print(f"[bench] Or:  --low-vram (auto-tune settings)")
+            raise
+        else:
+            # Try pre-quantized unsloth model
+            unsloth_model = f"unsloth/{model_name.split('/')[-1]}-bnb-4bit"
+            print(f"[bench] Loading failed, trying pre-quantized: {unsloth_model}")
+            model = quantize_model_4bit(unsloth_model)
+    lora_config = LoRAConfig(r=16, lora_alpha=16, use_rslora=True, use_manual_gradients=False)
+    model = auto_patch_model(model, apply_lora_patch=True, lora_config=lora_config)
+
+    dataset = load_dataset(dataset_name, split="train")
+    if max_samples > 0:
+        dataset = dataset.select(range(min(max_samples, len(dataset))))
+
+    def tokenize_fn(examples):
+        texts = [
+            f"### Instruction:\n{inst}\n\n### Response:\n{out}"
+            for inst, out in zip(examples["instruction"], examples["output"])
+        ]
+        return tokenizer(texts, truncation=True, max_length=max_seq_length)
+
+    tokenized = dataset.map(tokenize_fn, batched=True, remove_columns=dataset.column_names)
+    tokenized = tokenized.map(lambda x: {"labels": x["input_ids"]}, batched=True)
+    # Don't set_format("torch") — datasets' torch formatter crashes on Windows
+    # because it tries to import torchvision.io.VideoReader which isn't available.
+    # Our collator handles list->tensor conversion manually.
+
+    # Pre-pack: concatenate all sequences, split into max_seq_length chunks.
+    # This gives 2-5x fewer forward passes vs padding single short sequences.
+    from amazingvmsloth.packing import prepack_dataset, BatchCollator
+    packed = prepack_dataset(tokenized, tokenizer, max_seq_length=max_seq_length)
+    print(f"[bench] Pre-packed {len(tokenized)} samples into {len(packed)} chunks ({max_seq_length} tokens each)")
+
+    config = TrainingConfig(
+        output_dir="./bench_amazingvmsloth",
+        num_train_epochs=1,
+        per_device_train_batch_size=1,
+        gradient_accumulation_steps=1,
+        learning_rate=2e-4,
+        bf16=True,
+        optim="paged_adamw_8bit",
+        max_seq_length=max_seq_length,
+        logging_steps=1,
+        save_steps=999999,
+        packing=False,  # Already pre-packed above
+        compile_model=False,  # Compile overhead dominates on short runs despite pre-packing
+        gradient_checkpointing=True,  # Enable for VRAM optimization (match unsloth's 2.3GB)
+        chunked_loss=False,  # 512 seq fits full logits in 4GB, skip chunked overhead
+        silent=True,  # Skip tqdm/postfix overhead in benchmark
+    )
+
+    trainer = AmazingTrainer(
+        model=model,
+        tokenizer=tokenizer,
+        train_dataset=packed,
+        config=config,
+        data_collator=BatchCollator(),
+    )
+
+    start = time.time()
+    result = trainer.train()
+    elapsed = time.time() - start
+
+    vram_peak = torch.cuda.max_memory_allocated(0) / 1024**3 if torch.cuda.is_available() else 0
+    torch.cuda.reset_peak_memory_stats()
+
+    return {
+        "time_s": round(elapsed, 1),
+        "vram_peak_gb": round(vram_peak, 2),
+        "steps": result["global_step"],
+        "loss": result["log_history"][-1]["loss"] if result["log_history"] else -1,
+    }
+
+
+def run_unsloth(model_name: str, dataset_name: str, max_samples: int, max_seq_length: int) -> Dict[str, Any]:
+    if not _try_import_unsloth():
+        msg = "unsloth not installed" if _unsloth_error is None else f"unsloth import failed: {_unsloth_error}"
+        print(f"[bench] {msg}. Install with: pip install unsloth")
+        return {"time_s": -1, "vram_peak_gb": -1, "steps": -1, "loss": -1, "error": msg}
+
+    from unsloth import FastLanguageModel
+    from trl import SFTTrainer
+    from transformers import TrainingArguments
+    from datasets import load_dataset
+
+    try:
+        model, tokenizer = FastLanguageModel.from_pretrained(
+            model_name=model_name,
+            max_seq_length=max_seq_length,
+            dtype=None,
+            load_in_4bit=True,
+        )
+    except (OSError, MemoryError) as e:
+        err_msg = str(e)
+        if "paging file" in err_msg.lower() or "memory" in err_msg.lower():
+            print(f"[bench] unsloth failed: System out of memory (Windows paging file too small)")
+            print(f"[bench] unsloth requires more RAM/virtual memory than your system has.")
+            print(f"[bench] Tip: Increase Windows virtual memory or close other applications.")
+        else:
+            print(f"[bench] unsloth model loading failed: {err_msg}")
+        return {"time_s": -1, "vram_peak_gb": -1, "steps": -1, "loss": -1, "error": err_msg}
+
+    try:
+        model = FastLanguageModel.get_peft_model(
+            model,
+            r=16,
+            lora_alpha=16,
+            lora_dropout=0,
+            target_modules=["q_proj", "k_proj", "v_proj", "o_proj", "gate_proj", "up_proj", "down_proj"],
+            use_rslora=True,
+        )
+
+        dataset = load_dataset(dataset_name, split="train")
+        if max_samples > 0:
+            dataset = dataset.select(range(min(max_samples, len(dataset))))
+
+        def formatting_func(examples):
+            return [
+                f"### Instruction:\n{inst}\n\n### Response:\n{out}"
+                for inst, out in zip(examples["instruction"], examples["output"])
+            ]
+
+        training_args = TrainingArguments(
+            output_dir="./bench_unsloth",
+            num_train_epochs=1,
+            per_device_train_batch_size=1,
+            gradient_accumulation_steps=8,
+            learning_rate=2e-4,
+            bf16=True,
+            logging_steps=1,
+            save_steps=999999,
+            optim="paged_adamw_8bit",
+            max_grad_norm=1.0,
+            warmup_ratio=0.1,
+            report_to="none",
+        )
+
+        trainer = SFTTrainer(
+            model=model,
+            tokenizer=tokenizer,
+            train_dataset=dataset,
+            formatting_func=formatting_func,
+            max_seq_length=max_seq_length,
+            args=training_args,
+        )
+
+        torch.cuda.reset_peak_memory_stats()
+        start = time.time()
+        trainer.train()
+        elapsed = time.time() - start
+    except Exception as e:
+        elapsed = time.time() - start if 'start' in locals() else 0
+        err_msg = str(e)
+        if "paging file" in err_msg.lower() or "memory" in err_msg.lower() or "os error 1455" in err_msg.lower():
+            print(f"[bench] unsloth failed: System out of memory")
+            print(f"[bench] unsloth requires more RAM/virtual memory than your system has.")
+        else:
+            print(f"[bench] unsloth training failed: {err_msg}")
+        return {
+            "time_s": round(elapsed, 1) if elapsed > 0 else -1,
+            "vram_peak_gb": -1,
+            "steps": getattr(trainer.state, "global_step", 0) if 'trainer' in locals() else 0,
+            "loss": -1,
+            "error": err_msg,
+        }
+
+    vram_peak = torch.cuda.max_memory_allocated(0) / 1024**3 if torch.cuda.is_available() else 0
+    torch.cuda.reset_peak_memory_stats()
+
+    # Find last logged loss
+    loss = -1
+    for entry in reversed(trainer.state.log_history):
+        if "loss" in entry:
+            loss = entry["loss"]
+            break
+
+    return {
+        "time_s": round(elapsed, 1),
+        "vram_peak_gb": round(vram_peak, 2),
+        "steps": trainer.state.global_step,
+        "loss": round(loss, 4) if isinstance(loss, float) else -1,
+    }
+
+
+def main():
+    parser = argparse.ArgumentParser(description="Benchmark: amazingvmsloth vs unsloth")
+    parser.add_argument("--model", default="Qwen/Qwen2.5-1.5B", help="Model name")
+    parser.add_argument("--dataset", default="tatsu-lab/alpaca", help="Dataset")
+    parser.add_argument("--max-samples", type=int, default=500, help="Samples to use")
+    parser.add_argument("--max-seq-length", type=int, default=512, help="Max seq length")
+    parser.add_argument("--skip-unsloth", action="store_true", help="Skip unsloth benchmark")
+    args = parser.parse_args()
+
+    print(f"\n{'='*60}")
+    print(f"  Benchmark: amazingvmsloth vs unsloth")
+    print(f"{'='*60}")
+    print(f"  Model:      {args.model}")
+    print(f"  Dataset:    {args.dataset}")
+    print(f"  Samples:    {args.max_samples}")
+    print(f"  Seq length: {args.max_seq_length}")
+    print(f"{'='*60}\n")
+
+    print("[bench] Running amazingvmsloth...")
+    av_result = run_amazingvmsloth(args.model, args.dataset, args.max_samples, args.max_seq_length)
+    print(f"[bench] amazingvmsloth done: {av_result['time_s']}s, {av_result['vram_peak_gb']}GB peak VRAM\n")
+
+    us_result = None
+    if not args.skip_unsloth:
+        print("[bench] Running unsloth...")
+        us_result = run_unsloth(args.model, args.dataset, args.max_samples, args.max_seq_length)
+        if us_result.get("error"):
+            print(f"[bench] unsloth skipped: {us_result['error']}\n")
+            us_result = None
+        else:
+            print(f"[bench] unsloth done: {us_result['time_s']}s, {us_result['vram_peak_gb']}GB peak VRAM\n")
+
+    print(f"\n{'='*60}")
+    print(f"  Results")
+    print(f"{'='*60}")
+    print(f"  {'Metric':<20} {'amazingvmsloth':>15} {'unsloth':>15}")
+    print(f"  {'-'*20} {'-'*15} {'-'*15}")
+    print(f"  {'Time (s)':<20} {av_result['time_s']:>15} {us_result['time_s'] if us_result else 'N/A':>15}")
+    print(f"  {'Peak VRAM (GB)':<20} {av_result['vram_peak_gb']:>15} {us_result['vram_peak_gb'] if us_result else 'N/A':>15}")
+    print(f"  {'Final Loss':<20} {av_result['loss']:>15} {us_result['loss'] if us_result else 'N/A':>15}")
+    print(f"  {'Steps':<20} {av_result['steps']:>15} {us_result['steps'] if us_result else 'N/A':>15}")
+
+    if us_result and us_result["time_s"] > 0 and av_result["time_s"] > 0:
+        speedup = us_result["time_s"] / av_result["time_s"]
+        vram_saved = us_result["vram_peak_gb"] - av_result["vram_peak_gb"]
+        print(f"\n  Speedup:    {speedup:.2f}x")
+        print(f"  VRAM saved: {vram_saved:.2f} GB")
+        if speedup > 1:
+            print(f"  amazingvmsloth is {speedup:.2f}x faster!")
+        else:
+            print(f"  unsloth is {1/speedup:.2f}x faster")
+
+    print(f"{'='*60}\n")
+
+
+if __name__ == "__main__":
+    main()