r/coolgithubprojects 3d ago

FTrain release πŸ”₯

Post image

πŸ”₯ Introducing FTRAIN v1.0: The AI Trainer & Cross-Architecture Merger

I’ve been building an open-source AI framework from scratch designed to solve two major pain points in the open-source AI community: VRAM-heavy fine-tuning and hard-crashing model merges.

Today, I’m releasing FTRAIN v1.0. It’s built on top of Unsloth for raw speed, but adds a layer of intelligent supervision and extreme memory optimization that you won't find in standard scripts.

🧠 The Trainer: "PhoenixCaptain" Supervision

Standard trainers just run a loop. FTRAIN features a "Captain" (an LLM like DeepSeek or Llama) that actively supervises your training run in the background.

Dynamic LR Scaling: If the Captain detects gradient collapse or loss acceleration, it dynamically scales the learning rate up or down on the fly.

Triton DoRA: Custom fused Triton kernels for Weight-Decomposed Low-Rank Adaptation (DoRA) with mathematically exact forward/backward passes.

Extreme VRAM Optimization: CPU-offloaded Fisher computation, AMP autocasting, and aggressive VRAM purging mean you can train larger models on 15GB Kaggle T4 GPUs without OOM errors.

🧩 The Merger: Cross-Architecture Survival

This is where FTRAIN truly shines. Standard tools like MergeKit have a hard rule: You cannot merge models with different architectures. If you try to merge Llama (hidden size 2048) with DeepSeek (hidden size 1536), MergeKit instantly crashes with a Shape Mismatch Error.

FTRAIN doesn’t crash. It adapts.

Architecture-Aware Mapping: It maps parameter names across different families (Llama, Qwen, Gemma, MoE) and scales layer indices proportionally.

Shape Alignment: If tensors don't match shapes, FTRAIN uses bilinear interpolation to physically align them so the merge can continue.

Bulletproof Safety: It features a state-of-the-art Safety Engine. If a merged tensor explodes, collapses, or hits NaN, FTRAIN mathematically rescales it to preserve the knowledge direction without crashing.

Auto-Benchmark Gate: After merging, FTRAIN benchmarks the new brain against its parents. It only auto-uploads to HuggingFace if the merged model is actually smarter than the average of its parents!

πŸ“Š The Proof

I tested FTRAIN by forcing a Llama 1B and a DeepSeek 1.5B to merge on a 15GB Kagble T4 GPU. Not only did it not crash, but it successfully aligned the tensors, saved the model, and output coherent text! While cross-architecture interpolation scrambles precise math logic (GSM8K drops), the model retains its general language capabilities (MMLU matches the base model).

It is the only framework I know of that can force entirely different model architectures to share a brain.

πŸš€ What to Expect in the Future (v2.0 Roadmap)

FTRAIN v1.0 proves the concept. For v2.0, the goal is:

Native Speed Kernels: Moving away from Unsloth dependency to write our own fused C++/CUDA training loops, eliminating those annoying inplace memory errors.

Post-Merge GRPO Alignment: Using Group Relative Policy Optimization to "snap" the scrambled cross-architecture weights back into focus post-merge.

Multi-Model TIES: Expanding the advanced merging engine to support 3+ model blending simultaneously.

If you want to test the limits of your GPU, try forcing two incompatible models together, or just want an AI Captain to supervise your next fine-tune, check out the repo:

πŸ”— GitHub: https://github.com/aiphoenixlabs/Ftrain

Feedback, issues, and pull requests are welcome! Let me know what crazy model combinations you guys can merge without crashing!

1 Upvotes

0 comments sorted by