zenlm/zen-identity
1
0
Fork 0
mirror of https://github.com/zenlm/zen-identity.git synced 2026-07-26 22:30:18 +00:00
Code Issues Packages Projects Releases Wiki Activity

Fix model configs: zen-coder=Devstral2, zen-coder-flash=GLM-4.7-Flash, zen-coder-max=Kimi-K2

Browse Source
This commit is contained in:
Zach Kelling
2026-02-04 18:14:31 -08:00
parent e46368712c
commit 5e5583b8ac
2 changed files with 214 additions and 589 deletions
Download Patch File Download Diff File Expand all files Collapse all files
trainer/zen_trainer/models.py
+73 -234
View File
@@ -1,8 +1,10 @@
"""
Zen Coder Model Configurations
5 supported architectures with VRAM requirements and training cost estimates.
Costs based on 8xH200 HGX ($35/hr) with ~1.5TB total VRAM.
Correct model mappings:
- zen-coder: Devstral 2 (24B dense)
- zen-coder-flash: GLM-4.7-Flash (31B MoE, 3B active) - FLAGSHIP
- zen-coder-max: Kimi K2 (671B MoE, 14B active)
"""
from dataclasses import dataclass
@@ -11,284 +13,121 @@ from typing import Optional, Dict, Any
@dataclass
class ModelConfig:
"""Configuration for a model architecture."""
name: str
hf_id: str
size_b: float # Billions of parameters
architecture: str # dense, moe
vram_full: int # GB for full precision
vram_qlora: int # GB for QLoRA (4-bit)
max_seq_length: int
base_model: str
size_total: str
size_active: str
architecture: str
vram_qlora: int
context_window: int
license: str
# Training settings
lora_r: int
lora_alpha: int
batch_size: int
grad_accum: int
learning_rate: float
# Cost estimates (8xH200 @ $35/hr, ~10K samples)
hours_estimate: tuple # (min, max)
cost_estimate: tuple # (min, max) USD
# Backend support
supports_mlx: bool
supports_unsloth: bool
supports_deepspeed: bool
# Quantization
gguf_url: Optional[str] = None
fits_128gb: bool = False
# All 5 Zen Coder model configurations
ZEN_MODELS: Dict[str, ModelConfig] = {
# =========================================
# Tier 1: Lightweight (Local Development)
# =========================================
"qwen3-4b": ModelConfig(
name="Zen Coder 4B (Qwen3)",
hf_id="Qwen/Qwen3-4B-Instruct-2507",
size_b=4.0,
ZEN_CODER_MODELS: Dict[str, ModelConfig] = {
# Edge model
"zen-coder-4b": ModelConfig(
name="Zen Coder 4B",
hf_id="zenlm/zen-coder",
base_model="Qwen/Qwen3-Coder-4B",
size_total="4B",
size_active="4B",
architecture="dense",
vram_full=16,
vram_qlora=8,
max_seq_length=4096,
context_window=32768,
license="Apache 2.0",
lora_r=64,
lora_alpha=128,
batch_size=4,
grad_accum=4,
learning_rate=2e-4,
hours_estimate=(0.5, 1.5),
cost_estimate=(17, 52), # $35/hr × 0.5-1.5h
supports_mlx=True,
supports_unsloth=True,
supports_deepspeed=True,
fits_128gb=True,
),
# =========================================
# Tier 2: Mid-Range (Production)
# =========================================
"devstral-24b": ModelConfig(
name="Zen Coder 24B (Devstral Small 2)",
hf_id="mistralai/Devstral-Small-2-24B-Instruct-2512",
size_b=24.0,
# Main production model
"zen-coder": ModelConfig(
name="Zen Coder",
hf_id="zenlm/zen-coder-24b",
base_model="mistralai/Devstral-Small-2-24B-Instruct-2512",
size_total="24B",
size_active="24B",
architecture="dense",
vram_full=96,
vram_qlora=24,
max_seq_length=8192,
context_window=262144, # 256K context!
context_window=262144,
license="Apache 2.0",
lora_r=32,
lora_alpha=64,
batch_size=2,
grad_accum=8,
learning_rate=1e-4,
hours_estimate=(2, 4),
cost_estimate=(70, 140), # $35/hr × 2-4h
supports_mlx=True, # With quantization
supports_mlx=True,
supports_unsloth=True,
supports_deepspeed=True,
fits_128gb=True,
),
# =========================================
# Tier 3: Large (High-Performance)
# =========================================
"devstral-123b": ModelConfig(
name="Zen Coder 123B (Devstral 2)",
hf_id="mistralai/Devstral-2-123B-Instruct-2512",
size_b=123.0,
architecture="dense",
vram_full=492,
vram_qlora=128, # Fits in 128GB with 4-bit!
max_seq_length=8192,
context_window=262144, # 256K context
license="Mistral Research",
lora_r=16,
lora_alpha=32,
batch_size=1,
grad_accum=16,
learning_rate=5e-5,
hours_estimate=(4, 8),
cost_estimate=(140, 280), # $35/hr × 4-8h
supports_mlx=False, # Too large
supports_unsloth=True,
supports_deepspeed=True,
fits_128gb=True, # QLoRA 4-bit fits!
gguf_url="https://huggingface.co/unsloth/Devstral-2-123B-GGUF",
),
# =========================================
# Tier 4: MoE Large (Maximum Performance)
# =========================================
"glm47-358b": ModelConfig(
name="Zen Coder MAX (GLM-4.7)",
hf_id="zai-org/GLM-4.7",
size_b=358.0, # 358B total (MoE)
architecture="moe", # Mixture of Experts
vram_full=716, # 358B × 2 bytes (BF16)
vram_qlora=180, # Q4 ~90GB + training overhead
max_seq_length=8192,
context_window=200000, # 200K context, 128K output
license="GLM-4 License",
lora_r=16,
lora_alpha=32,
batch_size=1,
grad_accum=16,
learning_rate=5e-6,
hours_estimate=(6, 12),
cost_estimate=(210, 420), # $35/hr × 6-12h
supports_mlx=True, # Can train on Mac Studio 512GB!
supports_unsloth=True,
supports_deepspeed=True,
fits_128gb=False, # Needs ~180GB for QLoRA training
gguf_url="https://huggingface.co/AaryanK/GLM-4.7-GGUF",
),
# =========================================
# Tier 5: MoE Ultra (Research/Frontier)
# =========================================
"kimi-k2-1t": ModelConfig(
name="Zen Coder ULTRA (Kimi K2 Thinking)",
hf_id="moonshotai/Kimi-K2-Instruct",
size_b=1000.0, # 1 Trillion params
architecture="moe", # MoE with 32 experts
vram_full=2000,
vram_qlora=400, # Needs 8xH200 minimum
max_seq_length=8192,
context_window=131072, # 128K context
# FLAGSHIP - Best balance of performance and efficiency
"zen-coder-flash": ModelConfig(
name="Zen Coder Flash",
hf_id="zenlm/zen-coder-flash",
base_model="zai-org/GLM-4.7-Flash",
size_total="31B MoE",
size_active="3B",
architecture="moe",
vram_qlora=24,
context_window=131072,
license="MIT",
lora_r=64,
lora_alpha=128,
batch_size=2,
learning_rate=1e-4,
supports_mlx=True,
supports_unsloth=True,
),
lora_r=8,
lora_alpha=16,
# Frontier model
"zen-coder-max": ModelConfig(
name="Zen Coder Max",
hf_id="zenlm/zen-max",
base_model="moonshotai/Kimi-K2-Instruct",
size_total="671B MoE",
size_active="14B",
architecture="moe",
vram_qlora=160,
context_window=262144,
license="Apache 2.0",
lora_r=16,
lora_alpha=32,
batch_size=1,
grad_accum=32,
learning_rate=1e-6,
hours_estimate=(12, 24),
cost_estimate=(420, 840), # $35/hr × 12-24h
learning_rate=5e-5,
supports_mlx=False,
supports_unsloth=False, # Too large
supports_deepspeed=True, # ZeRO-3 required
fits_128gb=False, # Needs multi-node
supports_unsloth=False,
),
}
def get_model_config(model_key: str) -> ModelConfig:
"""Get configuration for a model by key."""
if model_key not in ZEN_MODELS:
available = ", ".join(ZEN_MODELS.keys())
if model_key not in ZEN_CODER_MODELS:
available = ", ".join(ZEN_CODER_MODELS.keys())
raise ValueError(f"Unknown model: {model_key}. Available: {available}")
return ZEN_MODELS[model_key]
return ZEN_CODER_MODELS[model_key]
def list_models_by_vram(max_vram_gb: int) -> list:
"""List models that fit within a VRAM budget (QLoRA)."""
return [
(key, cfg) for key, cfg in ZEN_MODELS.items()
if cfg.vram_qlora <= max_vram_gb
]
def list_models():
"""Print model summary."""
print("Zen Coder Model Family:")
print("-" * 70)
for key, cfg in ZEN_CODER_MODELS.items():
flag = " ⭐" if key == "zen-coder-flash" else ""
print(f"{key}{flag}")
print(f" Base: {cfg.base_model}")
print(f" Size: {cfg.size_total} ({cfg.size_active} active)")
print(f" Context: {cfg.context_window:,} tokens")
print()
def estimate_training_cost(
model_key: str,
num_samples: int = 10000,
hourly_rate: float = 35.0,
) -> Dict[str, Any]:
"""
Estimate training cost for a model.
Args:
model_key: Model identifier
num_samples: Number of training samples
hourly_rate: Cost per hour (default: $35 for 8xH200)
Returns:
Dict with cost estimates and recommendations
"""
cfg = get_model_config(model_key)
# Scale estimate by sample count (base is 10K)
scale = num_samples / 10000
hours_min = cfg.hours_estimate[0] * scale
hours_max = cfg.hours_estimate[1] * scale
cost_min = hours_min * hourly_rate
cost_max = hours_max * hourly_rate
return {
"model": cfg.name,
"size_b": cfg.size_b,
"architecture": cfg.architecture,
"samples": num_samples,
"hours_estimate": (round(hours_min, 1), round(hours_max, 1)),
"cost_estimate_usd": (round(cost_min), round(cost_max)),
"hourly_rate": hourly_rate,
"vram_required_gb": cfg.vram_qlora,
"fits_128gb": cfg.fits_128gb,
"recommended_platform": (
"MLX (local)" if cfg.supports_mlx and cfg.vram_qlora <= 64
else "8xH200 (cloud)" if cfg.vram_qlora <= 400
else "Multi-node cluster"
),
}
# Quick reference
COST_SUMMARY = """
╔═══════════════════════════════════════════════════════════════════════════════╗
║ ZEN CODER TRAINING COSTS (8xH200 @ $35/hr) ║
╠═══════════════════════════════════════════════════════════════════════════════╣
║ Model │ Size │ Weights │ VRAM (QLoRA) │ Hours │ Cost (10K) ║
╠═════════════════════╪════════╪═════════╪══════════════╪════════╪═════════════╣
║ Qwen3 4B │ 4B │ 8 GB │ 8 GB ✓ │ 0.5-1.5│ $17-52 ║
║ Devstral 24B │ 24B │ 48 GB │ 24 GB ✓ │ 2-4 │ $70-140 ║
║ Devstral 123B │ 123B │ 246 GB │ 128 GB ✓ │ 4-8 │ $140-280 ║
║ GLM-4.7 358B (MoE) │ 358B │ 716 GB │ 180 GB ◆ │ 6-12 │ $210-420 ║
║ Kimi K2 1T (MoE) │ 1T │ ~2 TB │ 400 GB │ 12-24 │ $420-840 ║
╚═══════════════════════════════════════════════════════════════════════════════╝
✓ = Fits in 128GB (M3 Ultra / single GPU node)
◆ = Fits Mac Studio 512GB or 8xH200 (no multi-GPU required)
Weights = BF16 full precision. Q4 inference = Weights ÷ 4
Model Keys:
qwen3-4b - Tier 1: Local dev, M3 Max 64GB / RTX 4090
devstral-24b - Tier 2: Production, M3 Ultra 128GB / A100
devstral-123b - Tier 3: High-perf, M3 Ultra 192GB / 8xH200
glm47-358b - Tier 4: Maximum (MoE), Mac Studio 512GB / 8xH200
kimi-k2-1t - Tier 5: Ultra (MoE), Multi-node cluster only
Benchmarks (12 ARC - Agentic, Reasoning, Coding):
Eval toolkit: https://github.com/zai-org/glm-simple-evals
"""
if __name__ == "__main__":
print(COST_SUMMARY)
print("\nDetailed estimates for 3.35M samples:")
for key in ZEN_MODELS:
est = estimate_training_cost(key, num_samples=3_350_000)
print(f"\n{est['model']}:")
print(f" Hours: {est['hours_estimate'][0]}-{est['hours_estimate'][1]}")
print(f" Cost: ${est['cost_estimate_usd'][0]}-${est['cost_estimate_usd'][1]}")
print(f" Platform: {est['recommended_platform']}")
list_models()
training/app.py
+141 -355
View File
@@ -1,373 +1,159 @@
"""
Zen Training Space - Unified Training for All Zen Models
Train any Zen model with any dataset combination from HuggingFace
Zen Identity Training - HuggingFace Spaces
Fine-tune Zen models with identity data.
"""
import os
import gradio as gr
import torch
from transformers import AutoModel, AutoTokenizer, AutoProcessor, TrainingArguments, Trainer
from datasets import load_dataset, concatenate_datasets
import json
from typing import List, Dict
from transformers import AutoModelForCausalLM, AutoTokenizer, TrainingArguments, Trainer, BitsAndBytesConfig
from peft import LoraConfig, get_peft_model, prepare_model_for_kbit_training
from datasets import load_dataset
# Model configurations
# Zen Coder Model Family - CORRECT MAPPINGS
MODELS = {
"Language Models": {
"zen-nano-0.6b": {
"hf_id": "zenlm/zen-nano-0.6b",
"type": "language",
"size": "0.6B",
"context": "32K"
},
"zen-eco-4b-instruct": {
"hf_id": "zenlm/zen-eco-4b-instruct",
"type": "language",
"size": "4B",
"context": "32K"
},
"zen-eco-4b-agent": {
"hf_id": "zenlm/zen-eco-4b-agent",
"type": "language",
"size": "4B",
"context": "32K"
},
"zen-omni-7b": {
"hf_id": "zenlm/zen-omni-7b",
"type": "language",
"size": "7B",
"context": "32K"
},
"zen-coder-14b": {
"hf_id": "zenlm/zen-coder-14b",
"type": "language",
"size": "14B",
"context": "128K"
},
"zen-next-32b": {
"hf_id": "zenlm/zen-next-32b",
"type": "language",
"size": "32B",
"context": "32K"
},
"zen-coder-4b": {
"base": "Qwen/Qwen3-Coder-4B",
"size": "4B",
"arch": "dense",
"context": "32K",
},
"zen-coder": {
"base": "mistralai/Devstral-Small-2-24B-Instruct-2512",
"size": "24B",
"arch": "dense",
"context": "256K",
},
"zen-coder-flash": {
"base": "zai-org/GLM-4.7-Flash",
"size": "31B MoE (3B active)",
"arch": "moe",
"context": "131K",
"flagship": True,
},
"zen-coder-max": {
"base": "moonshotai/Kimi-K2-Instruct",
"size": "671B MoE (14B active)",
"arch": "moe",
"context": "256K",
},
"Vision-Language Models": {
"zen-vl-4b-instruct": {
"hf_id": "zenlm/zen-vl-4b-instruct",
"type": "vision-language",
"size": "4B",
"context": "32K"
},
"zen-vl-8b-instruct": {
"hf_id": "zenlm/zen-vl-8b-instruct",
"type": "vision-language",
"size": "8B",
"context": "32K"
},
"zen-vl-30b-instruct": {
"hf_id": "zenlm/zen-vl-30b-instruct",
"type": "vision-language",
"size": "30B",
"context": "32K"
},
}
}
# Dataset configurations
DATASETS = {
"Agent Training": {
"ADP - AgentTuning OS": {
"hf_id": "neulab/agent-data-collection",
"config": "agenttuning_os",
"size": "~5k samples"
},
"ADP - AgentTuning KG": {
"hf_id": "neulab/agent-data-collection",
"config": "agenttuning_kg",
"size": "~5k samples"
},
"ADP - AgentTuning DB": {
"hf_id": "neulab/agent-data-collection",
"config": "agenttuning_db",
"size": "~5k samples"
},
"ADP - Synatra": {
"hf_id": "neulab/agent-data-collection",
"config": "synatra",
"size": "99k samples"
},
"ADP - Code Feedback": {
"hf_id": "neulab/agent-data-collection",
"config": "code_feedback",
"size": "66k samples"
},
"ADP - Go Browse": {
"hf_id": "neulab/agent-data-collection",
"config": "go-browse-wa",
"size": "27k samples"
},
},
"Function Calling": {
"xLAM Function Calling 60k": {
"hf_id": "Salesforce/xlam-function-calling-60k",
"config": None,
"size": "60k samples"
},
},
"Instruction Tuning": {
"Alpaca": {
"hf_id": "tatsu-lab/alpaca",
"config": None,
"size": "52k samples"
},
}
}
OUTPUT_DIR = "./zen-identity-lora"
def train_model(
model_name: str,
selected_datasets: List[str],
max_samples: int,
epochs: int,
batch_size: int,
learning_rate: float,
output_repo: str
):
"""Main training function"""
try:
logs = []
def log(msg):
print(msg)
logs.append(msg)
yield "\n".join(logs)
yield from log("=" * 80)
yield from log("🧘 ZEN TRAINING SPACE")
yield from log("=" * 80)
yield from log("")
# GPU info
yield from log(f"🎮 GPU Available: {torch.cuda.is_available()}")
if torch.cuda.is_available():
yield from log(f" Device: {torch.cuda.get_device_name(0)}")
yield from log(f" Memory: {torch.cuda.get_device_properties(0).total_memory / 1e9:.1f}GB")
yield from log("")
# Find model config
model_config = None
for category in MODELS.values():
if model_name in category:
model_config = category[model_name]
break
if not model_config:
yield from log(f"❌ Model {model_name} not found")
return
yield from log(f"📦 Loading model: {model_name}")
yield from log(f" HF ID: {model_config['hf_id']}")
yield from log(f" Size: {model_config['size']}")
yield from log(f" Type: {model_config['type']}")
# Load model
model = AutoModel.from_pretrained(
model_config['hf_id'],
torch_dtype=torch.bfloat16,
device_map="auto",
trust_remote_code=True
)
if model_config['type'] == "vision-language":
processor = AutoProcessor.from_pretrained(model_config['hf_id'])
else:
processor = AutoTokenizer.from_pretrained(model_config['hf_id'])
yield from log("✅ Model loaded")
yield from log("")
# Load datasets
yield from log("📚 Loading datasets...")
all_datasets = []
for dataset_name in selected_datasets:
# Find dataset config
dataset_config = None
for category in DATASETS.values():
if dataset_name in category:
dataset_config = category[dataset_name]
break
if not dataset_config:
yield from log(f"⚠️ Dataset {dataset_name} not found, skipping")
continue
yield from log(f" Loading: {dataset_name}")
yield from log(f" HF ID: {dataset_config['hf_id']}")
try:
if dataset_config['config']:
ds = load_dataset(
dataset_config['hf_id'],
dataset_config['config'],
split="train",
streaming=True
)
else:
ds = load_dataset(
dataset_config['hf_id'],
split="train",
streaming=True
)
# Take limited samples
samples = []
for i, example in enumerate(ds):
if i >= max_samples // len(selected_datasets):
break
samples.append(example)
all_datasets.extend(samples)
yield from log(f" ✅ Loaded {len(samples)} samples")
except Exception as e:
yield from log(f" ❌ Error: {e}")
yield from log(f"\n✅ Total samples loaded: {len(all_datasets)}")
yield from log("")
# Training setup
yield from log("⚙️ Training Configuration:")
yield from log(f" Epochs: {epochs}")
yield from log(f" Batch Size: {batch_size}")
yield from log(f" Learning Rate: {learning_rate}")
yield from log(f" Samples: {len(all_datasets)}")
yield from log(f" Output: {output_repo}")
yield from log("")
training_args = TrainingArguments(
output_dir="./training-output",
num_train_epochs=epochs,
per_device_train_batch_size=batch_size,
learning_rate=learning_rate,
logging_steps=10,
save_steps=100,
bf16=True,
push_to_hub=True,
hub_model_id=output_repo,
report_to="tensorboard",
)
# Create trainer
trainer = Trainer(
model=model,
args=training_args,
train_dataset=all_datasets if len(all_datasets) > 0 else None,
)
# Train!
yield from log("🔥 TRAINING STARTED")
yield from log("=" * 80)
result = trainer.train()
yield from log("")
yield from log("=" * 80)
yield from log("✅ TRAINING COMPLETED!")
yield from log("=" * 80)
yield from log(f"📊 Final Loss: {result.training_loss:.4f}")
yield from log(f"☁️ Model uploaded to: {output_repo}")
yield from log("")
yield from log("🎉 SUCCESS!")
except Exception as e:
yield from log(f"\n❌ ERROR: {str(e)}")
import traceback
yield from log(f"\n{traceback.format_exc()}")
# Build Gradio Interface
with gr.Blocks(title="Zen Training Space", theme=gr.themes.Soft()) as demo:
gr.Markdown("""
# 🧘 Zen Training Space
### Unified Training Platform for All Zen Models
Train any Zen model with any dataset combination from HuggingFace.
All datasets are loaded directly from HF - no local storage needed!
""")
with gr.Row():
with gr.Column(scale=1):
gr.Markdown("### 1. Select Model")
model_choice = gr.Dropdown(
choices=[
*[f"{cat} / {model}" for cat in MODELS for model in MODELS[cat]]
],
label="Model",
value="Vision-Language Models / zen-vl-4b-instruct"
)
gr.Markdown("### 2. Select Datasets")
dataset_choices = gr.CheckboxGroup(
choices=[
*[f"{cat} / {ds}" for cat in DATASETS for ds in DATASETS[cat]]
],
label="Datasets",
value=[
"Agent Training / ADP - Synatra",
"Function Calling / xLAM Function Calling 60k"
]
)
gr.Markdown("### 3. Training Config")
max_samples = gr.Slider(100, 100000, value=10000, step=100, label="Max Samples")
epochs = gr.Slider(1, 10, value=3, step=1, label="Epochs")
batch_size = gr.Slider(1, 8, value=1, step=1, label="Batch Size")
learning_rate = gr.Number(value=2e-5, label="Learning Rate")
output_repo = gr.Textbox(
value="zenlm/zen-vl-4b-agent-custom",
label="Output Repository (HuggingFace)"
)
train_btn = gr.Button("🚀 Start Training", variant="primary", size="lg")
with gr.Column(scale=2):
gr.Markdown("### Training Logs")
output = gr.Textbox(label="", lines=35, max_lines=50, show_label=False)
train_btn.click(
train_model,
inputs=[
model_choice,
dataset_choices,
max_samples,
epochs,
batch_size,
learning_rate,
output_repo
],
outputs=output
def train_model(model_key: str, lr: float, epochs: int, batch_size: int, lora_rank: int, progress=gr.Progress()):
if model_key not in MODELS:
return f"Unknown model: {model_key}"
cfg = MODELS[model_key]
base_model = cfg["base"]
progress(0.1, desc=f"Loading {model_key}...")
device = "cuda" if torch.cuda.is_available() else "cpu"
if device == "cpu":
return "No GPU available. Use a GPU Space."
bnb_config = BitsAndBytesConfig(
load_in_4bit=True,
bnb_4bit_quant_type="nf4",
bnb_4bit_compute_dtype=torch.bfloat16,
)
tokenizer = AutoTokenizer.from_pretrained(base_model, trust_remote_code=True)
model = AutoModelForCausalLM.from_pretrained(
base_model,
quantization_config=bnb_config,
device_map="auto",
trust_remote_code=True,
)
progress(0.3, desc="Setting up LoRA...")
model = prepare_model_for_kbit_training(model)
lora_config = LoraConfig(
r=lora_rank,
lora_alpha=lora_rank * 2,
lora_dropout=0.05,
target_modules=["q_proj", "k_proj", "v_proj", "o_proj"],
bias="none",
task_type="CAUSAL_LM",
)
model = get_peft_model(model, lora_config)
progress(0.4, desc="Loading identity data...")
# Load identity dataset from this repo
dataset = load_dataset("zenlm/zen-identity", split="train")
def tokenize(examples):
return tokenizer(examples["text"], truncation=True, max_length=512, padding="max_length")
tokenized = dataset.map(tokenize, batched=True)
progress(0.5, desc="Training...")
training_args = TrainingArguments(
output_dir=OUTPUT_DIR,
num_train_epochs=epochs,
per_device_train_batch_size=batch_size,
learning_rate=lr,
logging_steps=10,
save_steps=100,
bf16=True,
report_to="none",
)
trainer = Trainer(
model=model,
args=training_args,
train_dataset=tokenized,
)
trainer.train()
progress(0.9, desc="Saving...")
model.save_pretrained(OUTPUT_DIR)
tokenizer.save_pretrained(OUTPUT_DIR)
return f"Done. Saved to {OUTPUT_DIR}"
with gr.Blocks(title="Zen Identity Training") as demo:
gr.Markdown("""
---
### 📊 Available Models
- **Language**: nano (0.6B), eco (4B), omni (7B), coder (14B), next (32B)
- **Vision-Language**: zen-vl (4B, 8B, 30B)
### 📚 Available Datasets
- **Agent Training**: ADP (220k+ trajectories across 15+ configs)
- **Function Calling**: xLAM (60k high-quality examples)
- **Instruction**: Alpaca (52k samples)
### 💰 Cost Estimates (HF Pro GPU)
- 4B model: $3-5 for 10k samples
- 8B model: $8-12 for 10k samples
- 32B model: $30-50 for 10k samples
# Zen Identity Training
Fine-tune Zen models with identity data.
| Model | Base | Size |
|-------|------|------|
| zen-coder-4b | Qwen3-Coder-4B | 4B |
| zen-coder | Devstral-Small-2-24B | 24B |
| **zen-coder-flash** | **GLM-4.7-Flash** | **31B MoE** ⭐ |
| zen-coder-max | Kimi-K2 | 671B MoE |
""")
with gr.Row():
model_select = gr.Dropdown(
choices=list(MODELS.keys()),
value="zen-coder-flash",
label="Model"
)
with gr.Row():
lr = gr.Slider(1e-5, 1e-3, value=1e-4, label="Learning Rate")
epochs = gr.Slider(1, 10, value=3, step=1, label="Epochs")
with gr.Row():
batch = gr.Slider(1, 4, value=1, step=1, label="Batch Size")
rank = gr.Slider(8, 64, value=16, step=8, label="LoRA Rank")
train_btn = gr.Button("Train", variant="primary")
output = gr.Textbox(label="Status", lines=3)
train_btn.click(train_model, [model_select, lr, epochs, batch, rank], output)
if __name__ == "__main__":
demo.launch(server_name="0.0.0.0", server_port=7860)
demo.launch()