Resolve merge conflict

frontend/package.json

@@ -0,0 +1,28 @@
+{
+    "name": "gpu-bartender",
+    "version": "1.0.0",
+    "description": "",
+    "main": "server.js",
+    "scripts": {
+        "build": "webpack --mode production",
+        "start": "ts-node server.ts"
+    },
+    "keywords": [],
+    "author": "",
+    "license": "ISC",
+    "dependencies": {
+        "@types/chart.js": "^2.9.41",
+        "body-parser": "^1.19.0",
+        "chart.js": "^4.4.3",
+        "express": "^4.17.1"
+    },
+    "devDependencies": {
+        "@types/body-parser": "^1.19.2",
+        "@types/express": "^4.17.13",
+        "ts-loader": "^9.5.1",
+        "ts-node": "^10.4.0",
+        "typescript": "^4.9.5",
+        "webpack": "^5.93.0",
+        "webpack-cli": "^5.1.4"
+    }
+}

frontend/server.ts

@@ -0,0 +1,64 @@
+import express, { Request, Response } from 'express';
+import bodyParser from 'body-parser';
+import path from 'path';
+import { VRAMCalculator } from './src/backend/calculator';
+import { DataArgs } from './src/backend/dataArgs';
+import { FinetuningArgs } from './src/backend/finetuningArgs';
+import { ModelArgs } from './src/backend/modelArgs';
+import { OptimizerArgs } from './src/backend/optimizerArgs';
+
+const app = express();
+const port = 3000;
+
+app.use(bodyParser.json());
+app.use(express.static(path.join(__dirname, 'src', 'frontend')));
+app.use('/dist', express.static(path.join(__dirname, 'dist')));
+app.post('/calculate', (req: Request, res: Response) => {
+    const data = req.body;
+
+    const modelArgs = new ModelArgs(
+        data.modelSize,
+        data.vocabSize,
+        data.hiddenSize,
+        data.numAttentionHeads,
+        data.numKeyValueHeads,
+        data.intermediateSize,
+        data.numLayers
+    );
+    const finetuningArgs = new FinetuningArgs(
+        data.trainingPrecision,
+        data.isFsdp,
+        data.loraAlpha,
+        data.loraDropout,
+        data.loraRank,
+        data.loraTarget,
+        data.qloraAlpha,
+        data.qloraDropout
+    );
+    const optimizerArgs = new OptimizerArgs(
+        data.optimizer,
+        data.optimizerSgdMomentum
+    );
+    const dataArgs = new DataArgs(
+        data.batchSize,
+        data.sequenceLength
+    );
+    const calculator = new VRAMCalculator(
+        modelArgs,
+        finetuningArgs,
+        optimizerArgs,
+        dataArgs,
+        data.numGpus,
+        data.unit
+    );
+
+    const result = calculator.estimate_result();
+    res.json(result);
+});
+app.get('*', (req: Request, res: Response) => {
+    res.sendFile(path.join(__dirname, 'src', 'frontend', 'index.html'));
+});
+
+app.listen(port, () => {
+    console.log(`Server running at http://localhost:${port}`);
+});

frontend/src/backend/calculator.ts

@@ -0,0 +1,84 @@
+import { DataArgs } from './dataArgs';
+import { FinetuningArgs } from './finetuningArgs';
+import { ModelArgs } from './modelArgs';
+import { OptimizerArgs } from './optimizerArgs';
+
+export class VRAMCalculator {
+    private bytesPerParam: number;
+    private gpuDivisor: number;
+
+    constructor(
+        private modelArgs: ModelArgs,
+        private finetuningArgs: FinetuningArgs,
+        private optimizerArgs: OptimizerArgs,
+        private dataArgs: DataArgs,
+        private numGpus: number = 1,
+        private unit: string = "MiB"
+    ) {
+        this.bytesPerParam = this.computeBytesPerParam();
+        this.gpuDivisor = this.computeGpuDivisor();
+    }
+
+    private computeBytesPerParam(): number {
+        return this.finetuningArgs.trainingPrecision === 'mixed' ? 6 : 4;
+    }
+
+    private computeGpuDivisor(): number {
+        return this.finetuningArgs.isFsdp && this.numGpus > 1 ? this.numGpus : 1;
+    }
+
+    private computeParameters(): number {
+        return (this.bytesPerParam * this.modelArgs.numParams * 1e9) / this.gpuDivisor;
+    }
+
+    private computeActivations(): number {
+        const { hiddenSize, numAttentionHeads, numKeyValueHeads, intermediateSize, numLayers } = this.modelArgs;
+        const { batchSize, sequenceLength } = this.dataArgs;
+        const headDim = hiddenSize / numAttentionHeads;
+        const attentionBlock = 2 * this.bytesPerParam * batchSize * sequenceLength * hiddenSize;
+        const mlpBlock = this.bytesPerParam * batchSize * sequenceLength * intermediateSize;
+        const layerNorms = 2 * this.bytesPerParam * batchSize * sequenceLength * hiddenSize;
+        const layer = attentionBlock + mlpBlock + layerNorms;
+        return layer * numLayers;
+    }
+
+    private computeOutputs(): number {
+        return 4 * this.dataArgs.batchSize * this.dataArgs.sequenceLength * this.modelArgs.vocabSize * 2;
+    }
+
+    private computeGradients(): number {
+        return (4 * this.modelArgs.numParams * 1e9) / this.gpuDivisor;
+    }
+
+    private computeFirstMoments(): number | null {
+        const { optimizer, optimizerSgdMomentum } = this.optimizerArgs;
+        if (!((optimizer === 'SGD' && optimizerSgdMomentum) || optimizer === 'Adam')) {
+            return null;
+        }
+        return (4 * this.modelArgs.numParams * 1e9) / this.gpuDivisor;
+    }
+
+    private computeSecondMoments(): number | null {
+        if (this.optimizerArgs.optimizer !== 'Adam') {
+            return null;
+        }
+        return (4 * this.modelArgs.numParams * 1e9) / this.gpuDivisor;
+    }
+
+    private roundNum(num: number): number {
+        const divisor = this.unit === "MiB" ? 2 ** 20 : 2 ** 30;
+        return Math.round((num / divisor) * 1e3) / 1e3;
+    }
+
+    public estimate_result(): Record<string, number | null> {
+        return {
+            cudaKernels: this.roundNum(1e6 * 2 ** 20),
+            parameters: this.roundNum(this.computeParameters()),
+            activations: this.roundNum(this.computeActivations()),
+            outputs: this.roundNum(this.computeOutputs()),
+            gradients: this.roundNum(this.computeGradients()),
+            firstMoments: this.roundNum(this.computeFirstMoments() || 0),
+            secondMoments: this.roundNum(this.computeSecondMoments() || 0)
+        };
+    }
+}

frontend/src/backend/dataArgs.ts

@@ -0,0 +1,6 @@
+export class DataArgs {
+    constructor(
+        public batchSize: number = 4,
+        public sequenceLength: number = 512
+    ) {}
+}

frontend/src/backend/finetuningArgs.ts

@@ -0,0 +1,20 @@
+import { LoraArgs, QLoraArgs } from './loraArgs';
+
+export class FinetuningArgs {
+    public loraArgs: LoraArgs;
+    public qloraArgs: QLoraArgs;
+
+    constructor(
+        public trainingPrecision: string = 'mixed',
+        public isFsdp: boolean = true,
+        loraAlpha?: number,
+        loraDropout?: number,
+        loraRank?: number,
+        loraTarget?: string,
+        qloraAlpha?: number,
+        qloraDropout?: number
+    ) {
+        this.loraArgs = new LoraArgs(loraAlpha, loraDropout, loraRank, loraTarget);
+        this.qloraArgs = new QLoraArgs(qloraAlpha, qloraDropout);
+    }
+}

frontend/src/backend/loraArgs.ts

@@ -0,0 +1,15 @@
+export class LoraArgs {
+    constructor(
+        public loraAlpha?: number,
+        public loraDropout?: number,
+        public loraRank: number = 8,
+        public loraTarget?: string
+    ) {}
+}
+
+export class QLoraArgs {
+    constructor(
+        public qloraAlpha?: number,
+        public qloraDropout?: number
+    ) {}
+}

frontend/src/backend/modelArgs.ts

@@ -0,0 +1,11 @@
+export class ModelArgs {
+    constructor(
+        public numParams: number = 1,
+        public vocabSize: number = 1,
+        public hiddenSize: number = 1,
+        public numAttentionHeads: number = 1,
+        public numKeyValueHeads: number = 1,
+        public intermediateSize: number = 1,
+        public numLayers: number = 1
+    ) {}
+}

frontend/src/backend/optimizerArgs.ts

@@ -0,0 +1,6 @@
+export class OptimizerArgs {
+    constructor(
+        public optimizer: string = "adam",
+        public optimizerSgdMomentum?: number
+    ) {}
+}