claude-flow
Version:
Ruflo - Enterprise AI agent orchestration for Claude Code. Deploy 60+ specialized agents in coordinated swarms with self-learning, fault-tolerant consensus, vector memory, and MCP integration
455 lines • 14.9 kB
JavaScript
/**
* LoRA (Low-Rank Adaptation) Implementation
*
* Enables efficient fine-tuning by decomposing weight updates into low-rank matrices.
* Dramatically reduces memory requirements while maintaining adaptation quality.
*
* Features:
* - Rank decomposition (r << d) for memory efficiency
* - Additive weight updates: W' = W + BA (where B ∈ R^{d×r}, A ∈ R^{r×k})
* - Support for multiple adaptation heads
* - Persistence to .swarm/lora-weights.json
*
* Memory savings:
* - Original: d × k parameters
* - LoRA: r × (d + k) parameters
* - For d=384, k=384, r=8: 786,432 → 6,144 (128x reduction)
*
* @module lora-adapter
*/
import { existsSync, mkdirSync, readFileSync, writeFileSync } from 'fs';
import { dirname, join } from 'path';
// ============================================================================
// Types & Constants
// ============================================================================
/**
* Default LoRA rank (determines memory/quality tradeoff)
*/
export const DEFAULT_RANK = 8;
/**
* Input dimension (384 from ONNX MiniLM-L6-v2)
*/
export const INPUT_DIM = 384;
/**
* Default output dimension (same as input for adapter)
*/
export const OUTPUT_DIM = 384;
/**
* Default alpha scaling factor
*/
export const DEFAULT_ALPHA = 16;
// ============================================================================
// Default Configuration
// ============================================================================
const DEFAULT_CONFIG = {
rank: DEFAULT_RANK,
alpha: DEFAULT_ALPHA,
inputDim: INPUT_DIM,
outputDim: OUTPUT_DIM,
learningRate: 0.001,
weightsPath: join(process.cwd(), '.swarm', 'lora-weights.json'),
enableDropout: true,
dropoutProb: 0.1,
autoSaveInterval: 50,
};
// ============================================================================
// LoRA Adapter Class
// ============================================================================
/**
* Low-Rank Adaptation module for efficient embedding fine-tuning
*/
export class LoRAAdapter {
config;
weights;
totalAdaptations = 0;
totalUpdates = 0;
adaptationNormSum = 0;
lastUpdate = null;
updatesSinceLastSave = 0;
constructor(config) {
this.config = { ...DEFAULT_CONFIG, ...config };
this.weights = this.initializeWeights();
}
/**
* Initialize weights with Kaiming/He initialization
*/
initializeWeights() {
const { rank, inputDim, outputDim, alpha } = this.config;
// A: rank × inputDim, initialized with Kaiming normal
const A = new Float32Array(rank * inputDim);
const stdA = Math.sqrt(2.0 / inputDim);
for (let i = 0; i < A.length; i++) {
A[i] = this.gaussianRandom() * stdA;
}
// B: outputDim × rank, initialized to zero (standard LoRA init)
const B = new Float32Array(outputDim * rank);
// B starts at zero so initial adaptation is zero
return {
A,
B,
scaling: alpha / rank,
};
}
/**
* Box-Muller transform for Gaussian random numbers
*/
gaussianRandom() {
let u = 0, v = 0;
while (u === 0)
u = Math.random();
while (v === 0)
v = Math.random();
return Math.sqrt(-2.0 * Math.log(u)) * Math.cos(2.0 * Math.PI * v);
}
/**
* Initialize adapter and load persisted weights
*/
async initialize() {
const loaded = this.loadWeights();
return { success: true, weightsLoaded: loaded };
}
/**
* Apply LoRA adaptation to an embedding
* output = input + scaling * (B @ A @ input)
*/
adapt(input) {
const startTime = performance.now();
const { rank, inputDim, outputDim } = this.config;
const { A, B, scaling } = this.weights;
// Step 1: Compute A @ input (rank-dimensional)
const hidden = new Float32Array(rank);
for (let r = 0; r < rank; r++) {
let sum = 0;
const rowOffset = r * inputDim;
// Unroll by 4 for SIMD-friendly access
let i = 0;
for (; i + 3 < inputDim; i += 4) {
sum += A[rowOffset + i] * input[i];
sum += A[rowOffset + i + 1] * input[i + 1];
sum += A[rowOffset + i + 2] * input[i + 2];
sum += A[rowOffset + i + 3] * input[i + 3];
}
for (; i < inputDim; i++) {
sum += A[rowOffset + i] * input[i];
}
hidden[r] = sum;
}
// Optional dropout (during training inference, skip)
// In LoRA inference, we don't apply dropout
// Step 2: Compute B @ hidden (outputDim-dimensional)
const delta = new Float32Array(outputDim);
for (let o = 0; o < outputDim; o++) {
let sum = 0;
const rowOffset = o * rank;
for (let r = 0; r < rank; r++) {
sum += B[rowOffset + r] * hidden[r];
}
delta[o] = sum * scaling;
}
// Step 3: Add adaptation to input
const adapted = new Float32Array(outputDim);
let adaptationNorm = 0;
for (let i = 0; i < outputDim; i++) {
adapted[i] = input[i] + delta[i];
adaptationNorm += delta[i] * delta[i];
}
adaptationNorm = Math.sqrt(adaptationNorm);
// Update stats
this.totalAdaptations++;
this.adaptationNormSum += adaptationNorm;
const timeMs = performance.now() - startTime;
return { adapted, adaptationNorm, timeMs };
}
/**
* Train the adapter with a gradient signal
* Uses simplified update: A += lr * hidden^T @ grad, B += lr * grad @ hidden^T
*/
train(input, gradOutput, reward = 1.0) {
const { rank, inputDim, outputDim, learningRate } = this.config;
const { A, B, scaling } = this.weights;
// Forward pass to get hidden states
const hidden = new Float32Array(rank);
for (let r = 0; r < rank; r++) {
let sum = 0;
const rowOffset = r * inputDim;
for (let i = 0; i < inputDim; i++) {
sum += A[rowOffset + i] * input[i];
}
hidden[r] = sum;
}
// Compute gradient for B: grad_B = gradOutput @ hidden^T
const scaledLr = learningRate * reward * scaling;
for (let o = 0; o < outputDim; o++) {
const rowOffset = o * rank;
for (let r = 0; r < rank; r++) {
B[rowOffset + r] += scaledLr * gradOutput[o] * hidden[r];
}
}
// Compute gradient for hidden: grad_hidden = B^T @ gradOutput
const gradHidden = new Float32Array(rank);
for (let r = 0; r < rank; r++) {
let sum = 0;
for (let o = 0; o < outputDim; o++) {
sum += B[o * rank + r] * gradOutput[o];
}
gradHidden[r] = sum;
}
// Compute gradient for A: grad_A = gradHidden @ input^T
for (let r = 0; r < rank; r++) {
const rowOffset = r * inputDim;
for (let i = 0; i < inputDim; i++) {
A[rowOffset + i] += scaledLr * gradHidden[r] * input[i];
}
}
// Compute loss (L2 norm of gradient)
let loss = 0;
for (let i = 0; i < gradOutput.length; i++) {
loss += gradOutput[i] * gradOutput[i];
}
loss = Math.sqrt(loss);
// Update counters
this.totalUpdates++;
this.lastUpdate = Date.now();
this.updatesSinceLastSave++;
// Auto-save if needed
if (this.updatesSinceLastSave >= this.config.autoSaveInterval) {
this.saveWeights();
this.updatesSinceLastSave = 0;
}
return { updated: true, loss };
}
/**
* Merge LoRA weights into base weights (for deployment)
* Returns: W' = W + scaling * B @ A
*/
merge(baseWeights) {
const { rank, inputDim, outputDim } = this.config;
const { A, B, scaling } = this.weights;
// Compute BA product
const merged = new Float32Array(baseWeights);
for (let o = 0; o < outputDim; o++) {
for (let i = 0; i < inputDim; i++) {
let sum = 0;
for (let r = 0; r < rank; r++) {
sum += B[o * rank + r] * A[r * inputDim + i];
}
merged[o * inputDim + i] += scaling * sum;
}
}
return merged;
}
/**
* Get current statistics
*/
getStats() {
const { rank, inputDim, outputDim } = this.config;
const originalParams = inputDim * outputDim;
const loraParams = rank * (inputDim + outputDim);
return {
totalAdaptations: this.totalAdaptations,
totalUpdates: this.totalUpdates,
rank: this.config.rank,
compressionRatio: originalParams / loraParams,
avgAdaptationNorm: this.totalAdaptations > 0
? this.adaptationNormSum / this.totalAdaptations
: 0,
lastUpdate: this.lastUpdate,
};
}
/**
* Reset adapter to initial state
*/
reset() {
this.weights = this.initializeWeights();
this.totalAdaptations = 0;
this.totalUpdates = 0;
this.adaptationNormSum = 0;
this.lastUpdate = null;
this.updatesSinceLastSave = 0;
}
/**
* Save weights to disk
*/
saveWeights() {
try {
const dir = dirname(this.config.weightsPath);
if (!existsSync(dir)) {
mkdirSync(dir, { recursive: true });
}
const data = {
version: 1,
config: {
rank: this.config.rank,
alpha: this.config.alpha,
inputDim: this.config.inputDim,
outputDim: this.config.outputDim,
},
weights: {
A: Array.from(this.weights.A),
B: Array.from(this.weights.B),
scaling: this.weights.scaling,
},
stats: {
totalAdaptations: this.totalAdaptations,
totalUpdates: this.totalUpdates,
adaptationNormSum: this.adaptationNormSum,
lastUpdate: this.lastUpdate,
},
savedAt: new Date().toISOString(),
};
writeFileSync(this.config.weightsPath, JSON.stringify(data, null, 2));
return true;
}
catch {
return false;
}
}
/**
* Load weights from disk
*/
loadWeights() {
try {
if (!existsSync(this.config.weightsPath)) {
return false;
}
const content = readFileSync(this.config.weightsPath, 'utf-8');
const data = JSON.parse(content);
if (data.version !== 1) {
return false;
}
// Verify dimensions match
const { rank, inputDim, outputDim } = data.config;
if (rank !== this.config.rank ||
inputDim !== this.config.inputDim ||
outputDim !== this.config.outputDim) {
return false;
}
// Load weights
this.weights = {
A: new Float32Array(data.weights.A),
B: new Float32Array(data.weights.B),
scaling: data.weights.scaling,
};
// Load stats
this.totalAdaptations = data.stats.totalAdaptations || 0;
this.totalUpdates = data.stats.totalUpdates || 0;
this.adaptationNormSum = data.stats.adaptationNormSum || 0;
this.lastUpdate = data.stats.lastUpdate || null;
return true;
}
catch {
return false;
}
}
/**
* Export weights as JSON
*/
exportWeights() {
return {
A: Array.from(this.weights.A),
B: Array.from(this.weights.B),
scaling: this.weights.scaling,
config: {
rank: this.config.rank,
alpha: this.config.alpha,
inputDim: this.config.inputDim,
outputDim: this.config.outputDim,
},
};
}
/**
* Import weights from JSON
*/
importWeights(data) {
try {
const { rank, inputDim, outputDim } = this.config;
if (data.A.length !== rank * inputDim ||
data.B.length !== outputDim * rank) {
return false;
}
this.weights = {
A: new Float32Array(data.A),
B: new Float32Array(data.B),
scaling: data.scaling,
};
return true;
}
catch {
return false;
}
}
}
// ============================================================================
// Singleton & Factory Functions
// ============================================================================
let loraInstance = null;
let initPromise = null;
/**
* Get or create singleton LoRA adapter instance
*/
export async function getLoRAAdapter() {
if (loraInstance) {
return loraInstance;
}
if (initPromise) {
return initPromise;
}
initPromise = (async () => {
const adapter = new LoRAAdapter();
await adapter.initialize();
loraInstance = adapter;
return adapter;
})();
return initPromise;
}
/**
* Reset singleton instance (for testing)
*/
export function resetLoRAAdapter() {
if (loraInstance) {
loraInstance.reset();
}
loraInstance = null;
initPromise = null;
}
/**
* Create new LoRA adapter instance (factory)
*/
export function createLoRAAdapter(config) {
return new LoRAAdapter(config);
}
/**
* Quick adaptation (convenience function)
*/
export async function adaptEmbedding(input) {
const adapter = await getLoRAAdapter();
return adapter.adapt(input);
}
/**
* Quick training (convenience function)
*/
export async function trainLoRA(input, gradOutput, reward) {
const adapter = await getLoRAAdapter();
return adapter.train(input, gradOutput, reward);
}
/**
* Get LoRA statistics (convenience function)
*/
export async function getLoRAStats() {
const adapter = await getLoRAAdapter();
return adapter.getStats();
}
export default {
LoRAAdapter,
getLoRAAdapter,
resetLoRAAdapter,
createLoRAAdapter,
adaptEmbedding,
trainLoRA,
getLoRAStats,
DEFAULT_RANK,
DEFAULT_ALPHA,
INPUT_DIM,
OUTPUT_DIM,
};
//# sourceMappingURL=lora-adapter.js.map