ML Research Plugin for Claude Code - Project Proposal

[nishide-dev]

Feature Type

Feature Type: Plugin Project - Complete ML Research Workflow Support

Problem Statement

Machine learning researchers face significant friction when working with modern ML frameworks:

1. Boilerplate Complexity: Setting up PyTorch Lightning + Hydra + W&B requires extensive boilerplate code and configuration
2. Framework Integration: Integrating Hugging Face Transformers, PyTorch Geometric, and Lightning with proper distributed training strategies is error-prone
3. Debugging Difficulty: Diagnosing training issues (NaN loss, memory leaks, convergence problems) requires deep framework knowledge
4. Configuration Management: Managing experiment configs, hyperparameter sweeps, and reproducibility is time-consuming
5. Best Practices Gap: Knowing industry-standard patterns (FSDP, LoRA, gradient checkpointing) requires extensive documentation review
6. Tool Fragmentation: Package managers (uv, pixi), formatters (ruff), type checkers (ty), loggers (W&B) need to work together seamlessly

Current Workflow Pain Points:

• Researchers spend 40-60% of time on infrastructure setup rather than experimentation
• Copy-pasting code between projects leads to inconsistencies and bugs
• Debugging distributed training failures is extremely time-consuming
• Keeping up with best practices (e.g., LoRA/QLoRA for LLMs) requires constant learning

Proposed Solution

Create a comprehensive Claude Code plugin that transforms ML research workflows by providing:

1. Intelligent Commands (10 slash commands)

Slash commands for common ML tasks:

• /project-init: Scaffold new PyTorch Lightning + Hydra projects with uv or pixi
• /ml-config: Generate and validate Hydra configurations
• /train: Execute training runs with Lightning Trainer
• /experiment: Manage experiments (create, compare, track)
• /debug: Analyze training logs and diagnose issues (NaN loss, OOM, convergence)
• /profile: Profile model performance (memory, compute, throughput)
• /data: Generate Lightning DataModules for common datasets
• /validate: Validate project structure, configs, and dependencies
• /model-export: Export models to ONNX, TorchScript, HF Hub
• /setup: Configure W&B, environment variables, GPU settings

2. Specialized Agents (6 sub-agents)

Domain-expert agents for complex tasks:

• ml-architect 🔵 (Opus): Design ML system architectures, distributed training strategies, scaling plans
• training-debugger 🔴 (Sonnet): Diagnose training issues, analyze logs, suggest fixes
• config-generator 💜 (Sonnet): Generate Hydra configs, validate YAML, manage experiments
• pytorch-expert 🟠 (Sonnet): Implement custom modules, optimization patterns, performance tuning
• geometric-specialist 🔷 (Sonnet): PyTorch Geometric expert for GNNs, heterogeneous graphs, distributed graph training
• transformers-specialist 💜 (Sonnet): Hugging Face Transformers expert for LLM fine-tuning, PEFT (LoRA/QLoRA), tokenization

3. Comprehensive Skills (7 knowledge bases)

Expert guides that provide context:

• ml-lightning-basics (1500+ lines): PyTorch Lightning fundamentals, LightningModule patterns, callbacks, distributed training
• ml-hydra-config (1000+ lines): Hydra configuration management, composition, instantiation, sweeps
• ml-pytorch-geometric (1200+ lines): GNN implementations, heterogeneous graphs, distributed graph training
• ml-wandb-tracking (800+ lines): W&B integration, experiment tracking, artifact management
• tool-pixi (1000+ lines): Pixi package manager for cross-platform ML environments
• ml-cli-tools (1000+ lines): Building CLIs with Typer + Rich, Hydra Compose API, logging patterns
• ml-transformers (1500+ lines): HF Transformers + Lightning integration, PEFT, distributed training (FSDP/DeepSpeed)

4. Automated Rules (3 rule sets)

Enforce best practices automatically:

• coding-standards.md: Type hints, deterministic operations, tensor shape docs, Lightning patterns, CLI best practices, Transformers integration
• security-practices.md: Credential management (Pydantic Settings), API key protection, safe logging
• workflow-constraints.md: Config validation, checkpointing, experiment tagging, reproducibility

5. Smart Hooks (automation triggers)

Auto-formatting and validation:

• SessionStart: Show available commands on plugin load
• BeforeSave (Python): Auto-format with ruff, validate with ty
• AfterEdit (YAML): Validate Hydra configs
• AfterEdit (Python): Run quick validation checks

Why This Plugin?

For ML Researchers

1. Accelerated Setup: Go from idea to training in minutes, not hours
2. Fewer Bugs: Industry-standard patterns prevent common mistakes
3. Better Debugging: Intelligent analysis of training issues
4. Modern Stack: uv/pixi + ruff + ty + Lightning + Hydra + W&B integration
5. LLM-Ready: Full support for Transformers fine-tuning with LoRA/QLoRA

For ML Teams

1. Consistency: Shared patterns across projects
2. Knowledge Transfer: Junior researchers learn from expert agents
3. Reproducibility: Proper config management and experiment tracking
4. Scalability: FSDP, DeepSpeed, distributed training built-in

Unique Value Proposition

Unlike generic ML tools, this plugin:

• Context-Aware: Understands your entire project structure (configs, checkpoints, logs)
• Framework-Native: Deep integration with Lightning, Hydra, HF Transformers, PyG
• Production-Ready: Industry best practices (not toy examples)
• Extensible: Easy to add new commands, agents, skills as frameworks evolve

Technical Architecture

Plugin Structure

claude-code-ml-research/
├── .claude-plugin/
│   └── plugin.json           # Plugin manifest
├── commands/                  # 10 slash commands
│   ├── project-init.md
│   ├── ml-config.md
│   ├── train.md
│   ├── experiment.md
│   ├── debug.md
│   ├── profile.md
│   ├── data.md
│   ├── validate.md
│   ├── model-export.md
│   └── setup.md
├── agents/                    # 6 specialized agents
│   ├── ml-architect.md
│   ├── training-debugger.md
│   ├── config-generator.md
│   ├── pytorch-expert.md
│   ├── geometric-specialist.md
│   └── transformers-specialist.md
├── skills/                    # 7 knowledge bases
│   ├── ml-lightning-basics/SKILL.md
│   ├── ml-hydra-config/SKILL.md
│   ├── ml-pytorch-geometric/SKILL.md
│   ├── ml-wandb-tracking/SKILL.md
│   ├── tool-pixi/SKILL.md
│   ├── ml-cli-tools/SKILL.md
│   └── ml-transformers/SKILL.md
├── rules/ml/                  # 3 rule sets
│   ├── coding-standards.md
│   ├── security-practices.md
│   └── workflow-constraints.md
└── hooks/
    └── hooks.json             # 5 automation hooks

Technology Stack

Core Frameworks:

• PyTorch Lightning 2.x (training framework)
• Hydra 1.3+ (configuration management)
• Hugging Face Transformers 4.x (LLMs/NLP)
• PyTorch Geometric 2.x (GNNs)
• Weights & Biases (experiment tracking)

Package Management:

• uv (Python package manager - fast, modern)
• pixi (cross-platform conda alternative)

Code Quality:

• ruff (linter + formatter - 10-100x faster than black/flake8/isort)
• ty (type checker - fast, user-friendly)
• pre-commit (automation)

Testing:

• pytest (unit tests)
• markdownlint (documentation quality)

Usage Examples

Example 1: Start New Project

# User in Claude Code:
/project-init

# Plugin prompts for:
# - Project name: "llm-sentiment-analysis"
# - Package manager: uv
# - Frameworks: Lightning + Transformers + W&B
# - Distributed strategy: FSDP

# Creates:
# - pyproject.toml with dependencies
# - src/models/lightning_module.py
# - configs/ with Hydra configs
# - scripts/train.py
# - .pre-commit-config.yaml
# - README.md with setup instructions

Example 2: Debug Training Issue

# User in Claude Code:
/debug logs/train_20250207.log

# training-debugger agent analyzes:
# - Detects NaN loss at step 1250
# - Identifies gradient explosion in layer 42
# - Suggests fixes:
#   1. Enable gradient clipping (max_norm=1.0)
#   2. Reduce learning rate by 10x
#   3. Check for numerical instability in custom loss

Example 3: Generate Hydra Config

# User in Claude Code:
/ml-config create experiment=bert_finetune

# config-generator creates:
# configs/experiment/bert_finetune.yaml:
#   - Model: bert-base-uncased
#   - Data: GLUE SST-2
#   - Trainer: DDP with 4 GPUs
#   - Callbacks: ModelCheckpoint, EarlyStopping
#   - Logger: W&B

Example 4: Fine-tune LLM with LoRA

# User in Claude Code:
I need to fine-tune Llama 2-7B on my custom dataset with LoRA

# transformers-specialist agent:
# 1. Generates LightningModule with LoRA config
# 2. Sets up QLoRA (4-bit quantization) for 24GB GPU
# 3. Creates DataModule for custom dataset
# 4. Configures FSDP for multi-GPU
# 5. Adds W&B tracking
# 6. Provides training script with best practices

Success Metrics

Quantitative

• Time to First Training Run: < 15 minutes (vs. hours manually)
• Bug Reduction: 60% fewer common mistakes (missing save_hyperparameters, wrong padding side, etc.)
• Code Reuse: 80% less boilerplate code
• Debugging Speed: 3x faster issue resolution

Qualitative

• Researchers spend more time on science, less on infrastructure
• Junior researchers onboard faster with expert agents
• Teams adopt consistent patterns and best practices
• Projects are more reproducible and maintainable

Target Audience

Primary

• ML Researchers: Academia and industry research labs
• ML Engineers: Building production ML systems
• PhD Students: Learning modern ML frameworks
• Kaggle Competitors: Rapid experimentation

Secondary

• ML Teams: Standardizing workflows across projects
• ML Educators: Teaching modern ML development
• Open Source Contributors: Contributing to ML frameworks

Alternatives Considered

1. Generic Templates: Too rigid, don't adapt to context
2. Cookiecutter Projects: Static, no intelligent assistance
3. Framework CLIs: Limited to single framework, no cross-framework patterns
4. LLM Prompting: No persistent knowledge, inconsistent advice

Why Claude Code Plugin Wins:

• Context-aware (sees entire project)
• Intelligent agents (specialized expertise)
• Persistent knowledge (skills library)
• Automated enforcement (rules + hooks)
• Extensible (easy to add new frameworks)

Risks and Mitigations

Risk 1: Framework Version Compatibility

Mitigation: Pin tested versions in templates, provide upgrade guides

Risk 2: Plugin Maintenance Overhead

Mitigation: Modular architecture, community contributions, automated tests

Risk 3: Learning Curve for Plugin Usage

Mitigation: Interactive tutorials, video guides, example projects

Risk 4: Divergence from Framework Best Practices

Mitigation: Regular reviews of official docs, community feedback, expert validation

Related ML Framework

• ✅ PyTorch Lightning
• ✅ Hugging Face Transformers
• ✅ PyTorch Geometric
• ✅ Hydra
• ✅ Weights & Biases

Priority

High - This plugin significantly improves ML research workflows and reduces common friction points.

Additional Context

Inspiration

This plugin addresses real pain points experienced in ML research:

• Hours spent setting up new projects
• Debugging obscure training failures
• Keeping configs in sync across experiments
• Learning new frameworks (Transformers PEFT, PyG MessagePassing)
• Integrating tools (Lightning + Hydra + W&B)

Design Philosophy

1. Expert Knowledge Encoded: Skills contain industry best practices, not toy examples
2. Context-Aware Assistance: Agents understand your project structure
3. Progressive Disclosure: Simple commands for common tasks, agents for complex ones
4. Framework Native: Deep integration, not shallow wrappers
5. Production Ready: Patterns work at scale (distributed, quantization, optimization)

Community Engagement

• Open Source: MIT License
• Contributions Welcome: Issues, PRs, skill additions
• Documentation First: Every feature well-documented
• User Feedback Loop: Regular surveys, feature requests

Contribution

• I am willing to implement this feature and submit a PR (Initial implementation complete)
• I can help with documentation for this feature
• I can help test this feature

Checklist

• I have searched existing issues and PRs to avoid duplicates
• This feature aligns with the plugin's ML research focus
• Initial implementation (Phase 1) is complete
• Ready for community testing and feedback (Phase 2)

---

Project Status

Current Status: Phase 1 Complete ✅

The core infrastructure is fully implemented and validated:

• 10 commands available
• 6 specialized agents ready
• 7 comprehensive skills (9000+ lines)
• 3 rule sets enforcing best practices
• 5 automation hooks configured
• CI/CD pipeline operational
• Documentation complete

Next Steps: Community testing, feedback gathering, iterative improvements based on real-world usage.

Timeline:

• Phase 1: ✅ Complete
• Phase 2: February-March 2026 (Testing & Validation)
• Phase 3: April 2026 (Polish & Documentation)
• Phase 4: Q3 2026+ (Advanced Features)

---

Issue Labels: enhancement, project-proposal, documentation, high-priority, phase-1-complete

Assignees: @nishide-dev

Projects: ML Research Plugin Development

Milestone: v0.1.0 - Core Infrastructure