lfm25-350m-linux-grpo

LFM2.5-350M Linux Command Generator

A fine-tuned version of LiquidAI/LFM2.5-350M specialized for converting natural language instructions into Linux shell commands.

What This Model Actually Is

This is a task-specific fine-tune of LiquidAI's 350M parameter language model, trained to generate bash commands wrapped in special tokens. It was created as a research/demonstration project to explore:

1. LoRA fine-tuning for command generation tasks
2. GRPO (Group Relative Policy Optimization) for reinforcement learning from rewards
3. Custom format training using special tokens
4. Production pipeline with Azure OpenAI dataset generation

Architecture & Training

Base Model

• Model: LiquidAI/LFM2.5-350M (350M parameters)
• Architecture: Transformer decoder-only
• Context Length: 4096 tokens

Training Pipeline

OpenAI GPT-4 Dataset (3,000 examples)
    ↓
SFT Training (4 epochs, LoRA r=16, BF16)
    - assistant_only_loss=True
    - Custom data collator with assistant_masks
    ↓
GRPO Training (2 epochs, 7 reward functions)
    - beta=0.04 (KL constraint)
    - num_generations=3 per prompt
    - Temperature annealing 0.7 → 0.45
    ↓
Final Merged Model

Dataset (15 Categories, ~3,000 examples)

Output Format (v30)

The model outputs raw bash commands between special tokens:

<|tool_call_start|>find . -name "*.py" -mtime -7<|tool_call_end|>

Key characteristics:

• No function wrappers (linux_command(...)) - just raw bash
• Uses LFM2.5's native special tokens: <|tool_call_start|>, <|tool_call_end|>
• Designed for direct extraction and execution

Performance Metrics (Actual)

What These Numbers Mean

• 100% Format Accuracy: The model consistently outputs commands in the correct format with proper special tokens
• 98% Tool Name: Almost never uses old function-wrapper format
• 24% Exact Match: Matches the reference command exactly 1 in 4 times (this is actually competitive for 350M parameters)
• 0.58 F1: Moderate token overlap with reference commands

Comparison Context

Takeaway: For a 350M parameter model, 24% EM is reasonable. It's not SOTA, but it's competitive for the size class and runs on minimal hardware.

Technical Highlights

1. LoRA Fine-Tuning

LoraConfig(
    r=16,                    # Rank
    lora_alpha=32,           # Scaling
    lora_dropout=0.05,
    target_modules="all-linear",
    bias="none",
)

2. GRPO Reward Functions (7 total)

• reward_format: Correct special token usage (+2/-1)
• reward_tool_name: Raw command format validation (+2/-2)
• reward_exact_cmd: Exact string match (+2, partial credit)
• reward_similarity: Token F1 similarity (0-1)
• reward_safety: Dangerous command penalty (-3)
• reward_penalties: Termination and structure quality
• reward_structure: Content quality and format

3. Critical Bug Fixes Applied

Tokenizer Patch for GRPO

# TRL's GRPOTrainer calls batch_decode with skip_special_tokens=True
# which strips our format tokens. We monkey-patch to force=False.
def _forced_decode(sequences, skip_special_tokens=True, **kwargs):
    return original(sequences, skip_special_tokens=False, **kwargs)

Pickle Fix for odict_keys

# TRL/Transformers has issues with odict_keys in save checkpoints
# We monkey-patch Trainer._save to convert to list before saving
def patched_save(self, output_dir, state_dict):
    if hasattr(self, 'model_kwarg_keys'):
        if isinstance(keys, (KeysView, ValuesView, ItemsView)):
            self.model_kwarg_keys = list(keys)
    # ... sanitize and save

4. Training Optimizations

• Right padding for training (assistant_only_loss requirement)
• BF16 mixed precision for speed
• Gradient checkpointing for memory
• Temperature annealing (0.7 → 0.45) for exploration → exploitation
• Milestone checkpoints at 10%, 50%, 100%

Usage

Basic Inference

from transformers import AutoModelForCausalLM, AutoTokenizer
import torch
import re

model_id = "2796gauravc/lfm25-350m-linux-grpo"
tokenizer = AutoTokenizer.from_pretrained(model_id)
model = AutoModelForCausalLM.from_pretrained(
    model_id,
    torch_dtype=torch.bfloat16,
    device_map="auto"
)

# Prepare prompt
messages = [
    {"role": "system", "content": "You are a Linux command assistant."},
    {"role": "user", "content": "Find all PDF files modified in the last 7 days"}
]

# Tokenize
enc = tokenizer.apply_chat_template(
    messages, 
    add_generation_prompt=True, 
    return_tensors="pt"
)
input_ids = enc.input_ids.to(model.device)
attention_mask = enc.attention_mask.to(model.device)

# Generate
with torch.no_grad():
    outputs = model.generate(
        input_ids=input_ids,
        attention_mask=attention_mask,
        max_new_tokens=100,
        do_sample=False,
        pad_token_id=tokenizer.pad_token_id
    )

# Decode with special tokens preserved
response = tokenizer.decode(
    outputs[0][input_ids.size(-1):], 
    skip_special_tokens=False
)

# Extract command
match = re.search(
    r"<\|tool_call_start\|>(.*?)<\|tool_call_end\|>", 
    response, 
    re.DOTALL
)
if match:
    command = match.group(1).strip()
    print(f"Generated: {command}")
    # Output: find . -name "*.py" -mtime -7

Hardware Requirements

Limitations & Honest Assessment

What It Does Well

1. ✅ Format compliance: Always uses correct special tokens
2. ✅ Simple commands: Good at basic file operations, text processing
3. ✅ Edge deployment: Small enough to run on consumer hardware
4. ✅ No function wrappers: Clean raw command output

What It Struggles With

1. ❌ Complex pipelines: Multi-stage commands with pipes
2. ❌ Exact match: Only 24% match reference exactly (but many alternatives are valid)
3. ❌ Edge cases: Unusual flags or rare utilities
4. ❌ Context awareness: No memory of previous commands

Known Issues

1. Semantic equivalence not string equivalence: Many valid bash commands exist for the same task. The model may generate a correct alternative that doesn't match the reference string.
2. Safety: While we filter dangerous patterns in training, the model could still suggest risky commands. Always review before execution.
3. Overfitting to training patterns: May repeat common patterns from the training data.

Citation

@misc{lfm25-350m-linux-grpo,
  title={LFM2.5-350M Linux Command Generator},
  author={Gaurav Chauhan},
  year={2026},
  howpublished={\url{https://huggingface.co/2796gauravc/lfm25-350m-linux-grpo}},
  note={350M parameter NL2Bash model with LoRA + GRPO training}
}

License

MIT License - See LICENSE file for details.

Base model: LiquidAI/LFM2.5-350M (Apache 2.0)

Acknowledgments

• LiquidAI for the LFM2.5 base model
• HuggingFace for Transformers and TRL libraries
• Azure OpenAI for dataset generation API