aws-cdk-neuronx-patterns/lib/base/neuronx/calculator.test.js

> [!WARNING] > This library is experimental module.

"use strict";
Object.defineProperty(exports, "__esModule", { value: true });
const aws_cdk_lib_1 = require("aws-cdk-lib");
const calculator_1 = require("./calculator");
const model_1 = require("./model");
const neuronx_instance_type_1 = require("./neuronx-instance-type");
describe("When the attention head is divisible", () => {
    test("Inf2", () => {
        const tensorParallel = (0, calculator_1.calcTensorParallel)(neuronx_instance_type_1.NeuronxInstanceType.INF2_48XLARGE, aws_cdk_lib_1.Size.gibibytes(16), 32);
        expect(tensorParallel).toEqual([
            {
                tp: 1,
                workers: 24,
                usage: 1,
            },
            {
                tp: 2,
                workers: 12,
                usage: 1,
            },
            {
                tp: 4,
                workers: 6,
                usage: 1,
            },
            {
                tp: 8,
                workers: 3,
                usage: 1,
            },
        ]);
    });
    test("Trn1", () => {
        const tensorParallel = (0, calculator_1.calcTensorParallel)(neuronx_instance_type_1.NeuronxInstanceType.TRN1_32XLARGE, aws_cdk_lib_1.Size.gibibytes(16), 48);
        expect(tensorParallel).toEqual([
            {
                tp: 1,
                workers: 32,
                usage: 1,
            },
            {
                tp: 2,
                workers: 16,
                usage: 1,
            },
            {
                tp: 8,
                workers: 4,
                usage: 1,
            },
            {
                tp: 16,
                workers: 2,
                usage: 1,
            },
        ]);
    });
});
describe("The one with low usage is shown last", () => {
    test("Inf2", () => {
        const tensorParallel = (0, calculator_1.calcTensorParallel)(neuronx_instance_type_1.NeuronxInstanceType.INF2_24XLARGE, aws_cdk_lib_1.Size.gibibytes(16), 32);
        expect(tensorParallel).toEqual([
            {
                tp: 1,
                workers: 12,
                usage: 1,
            },
            {
                tp: 2,
                workers: 6,
                usage: 1,
            },
            {
                tp: 4,
                workers: 3,
                usage: 1,
            },
            {
                tp: 8,
                workers: 1,
                usage: 0.6666666666666666,
            },
        ]);
    });
    test("Trn1", () => {
        const tensorParallel = (0, calculator_1.calcTensorParallel)(neuronx_instance_type_1.NeuronxInstanceType.TRN1_32XLARGE, aws_cdk_lib_1.Size.gibibytes(16), 32);
        expect(tensorParallel).toEqual([
            {
                tp: 1,
                workers: 32,
                usage: 1,
            },
            {
                tp: 2,
                workers: 16,
                usage: 1,
            },
            {
                tp: 8,
                workers: 4,
                usage: 1,
            },
            {
                tp: 16,
                workers: 2,
                usage: 1,
            },
            {
                tp: 32,
                workers: 1,
                usage: 1,
            },
        ]);
    });
});
describe("When the memory footprint is calculated", () => {
    test("Million size models are calculatable", () => {
        const memoryFootprint = (0, calculator_1.calcMemoryFootprint)(576, 30, model_1.DataTypeBits.BF16_OR_FP16, 128, 1);
        expect(memoryFootprint.toGibibytes()).toEqual(1);
    });
    test("Billion size models are calculatable", () => {
        const memoryFootprint = (0, calculator_1.calcMemoryFootprint)(6144, 48, model_1.DataTypeBits.BF16_OR_FP16, 1024, 1);
        expect(memoryFootprint.toGibibytes()).toEqual(42);
    });
});
//# sourceMappingURL=data:application/json;base64,eyJ2ZXJzaW9uIjozLCJmaWxlIjoiY2FsY3VsYXRvci50ZXN0LmpzIiwic291cmNlUm9vdCI6IiIsInNvdXJjZXMiOlsiLi4vLi4vLi4vc3JjL2Jhc2UvbmV1cm9ueC9jYWxjdWxhdG9yLnRlc3QudHMiXSwibmFtZXMiOltdLCJtYXBwaW5ncyI6Ijs7QUFBQSw2Q0FBbUM7QUFDbkMsNkNBQXVFO0FBQ3ZFLG1DQUF1QztBQUN2QyxtRUFBOEQ7QUFFOUQsUUFBUSxDQUFDLHNDQUFzQyxFQUFFLEdBQUcsRUFBRTtJQUNwRCxJQUFJLENBQUMsTUFBTSxFQUFFLEdBQUcsRUFBRTtRQUNoQixNQUFNLGNBQWMsR0FBRyxJQUFBLCtCQUFrQixFQUN2QywyQ0FBbUIsQ0FBQyxhQUFhLEVBQ2pDLGtCQUFJLENBQUMsU0FBUyxDQUFDLEVBQUUsQ0FBQyxFQUNsQixFQUFFLENBQ0gsQ0FBQztRQUNGLE1BQU0sQ0FBQyxjQUFjLENBQUMsQ0FBQyxPQUFPLENBQUM7WUFDN0I7Z0JBQ0UsRUFBRSxFQUFFLENBQUM7Z0JBQ0wsT0FBTyxFQUFFLEVBQUU7Z0JBQ1gsS0FBSyxFQUFFLENBQUM7YUFDVDtZQUNEO2dCQUNFLEVBQUUsRUFBRSxDQUFDO2dCQUNMLE9BQU8sRUFBRSxFQUFFO2dCQUNYLEtBQUssRUFBRSxDQUFDO2FBQ1Q7WUFDRDtnQkFDRSxFQUFFLEVBQUUsQ0FBQztnQkFDTCxPQUFPLEVBQUUsQ0FBQztnQkFDVixLQUFLLEVBQUUsQ0FBQzthQUNUO1lBQ0Q7Z0JBQ0UsRUFBRSxFQUFFLENBQUM7Z0JBQ0wsT0FBTyxFQUFFLENBQUM7Z0JBQ1YsS0FBSyxFQUFFLENBQUM7YUFDVDtTQUNGLENBQUMsQ0FBQztJQUNMLENBQUMsQ0FBQyxDQUFDO0lBQ0gsSUFBSSxDQUFDLE1BQU0sRUFBRSxHQUFHLEVBQUU7UUFDaEIsTUFBTSxjQUFjLEdBQUcsSUFBQSwrQkFBa0IsRUFDdkMsMkNBQW1CLENBQUMsYUFBYSxFQUNqQyxrQkFBSSxDQUFDLFNBQVMsQ0FBQyxFQUFFLENBQUMsRUFDbEIsRUFBRSxDQUNILENBQUM7UUFDRixNQUFNLENBQUMsY0FBYyxDQUFDLENBQUMsT0FBTyxDQUFDO1lBQzdCO2dCQUNFLEVBQUUsRUFBRSxDQUFDO2dCQUNMLE9BQU8sRUFBRSxFQUFFO2dCQUNYLEtBQUssRUFBRSxDQUFDO2FBQ1Q7WUFDRDtnQkFDRSxFQUFFLEVBQUUsQ0FBQztnQkFDTCxPQUFPLEVBQUUsRUFBRTtnQkFDWCxLQUFLLEVBQUUsQ0FBQzthQUNUO1lBQ0Q7Z0JBQ0UsRUFBRSxFQUFFLENBQUM7Z0JBQ0wsT0FBTyxFQUFFLENBQUM7Z0JBQ1YsS0FBSyxFQUFFLENBQUM7YUFDVDtZQUNEO2dCQUNFLEVBQUUsRUFBRSxFQUFFO2dCQUNOLE9BQU8sRUFBRSxDQUFDO2dCQUNWLEtBQUssRUFBRSxDQUFDO2FBQ1Q7U0FDRixDQUFDLENBQUM7SUFDTCxDQUFDLENBQUMsQ0FBQztBQUNMLENBQUMsQ0FBQyxDQUFDO0FBQ0gsUUFBUSxDQUFDLHNDQUFzQyxFQUFFLEdBQUcsRUFBRTtJQUNwRCxJQUFJLENBQUMsTUFBTSxFQUFFLEdBQUcsRUFBRTtRQUNoQixNQUFNLGNBQWMsR0FBRyxJQUFBLCtCQUFrQixFQUN2QywyQ0FBbUIsQ0FBQyxhQUFhLEVBQ2pDLGtCQUFJLENBQUMsU0FBUyxDQUFDLEVBQUUsQ0FBQyxFQUNsQixFQUFFLENBQ0gsQ0FBQztRQUNGLE1BQU0sQ0FBQyxjQUFjLENBQUMsQ0FBQyxPQUFPLENBQUM7WUFDN0I7Z0JBQ0UsRUFBRSxFQUFFLENBQUM7Z0JBQ0wsT0FBTyxFQUFFLEVBQUU7Z0JBQ1gsS0FBSyxFQUFFLENBQUM7YUFDVDtZQUNEO2dCQUNFLEVBQUUsRUFBRSxDQUFDO2dCQUNMLE9BQU8sRUFBRSxDQUFDO2dCQUNWLEtBQUssRUFBRSxDQUFDO2FBQ1Q7WUFDRDtnQkFDRSxFQUFFLEVBQUUsQ0FBQztnQkFDTCxPQUFPLEVBQUUsQ0FBQztnQkFDVixLQUFLLEVBQUUsQ0FBQzthQUNUO1lBQ0Q7Z0JBQ0UsRUFBRSxFQUFFLENBQUM7Z0JBQ0wsT0FBTyxFQUFFLENBQUM7Z0JBQ1YsS0FBSyxFQUFFLGtCQUFrQjthQUMxQjtTQUNGLENBQUMsQ0FBQztJQUNMLENBQUMsQ0FBQyxDQUFDO0lBQ0gsSUFBSSxDQUFDLE1BQU0sRUFBRSxHQUFHLEVBQUU7UUFDaEIsTUFBTSxjQUFjLEdBQUcsSUFBQSwrQkFBa0IsRUFDdkMsMkNBQW1CLENBQUMsYUFBYSxFQUNqQyxrQkFBSSxDQUFDLFNBQVMsQ0FBQyxFQUFFLENBQUMsRUFDbEIsRUFBRSxDQUNILENBQUM7UUFDRixNQUFNLENBQUMsY0FBYyxDQUFDLENBQUMsT0FBTyxDQUFDO1lBQzdCO2dCQUNFLEVBQUUsRUFBRSxDQUFDO2dCQUNMLE9BQU8sRUFBRSxFQUFFO2dCQUNYLEtBQUssRUFBRSxDQUFDO2FBQ1Q7WUFDRDtnQkFDRSxFQUFFLEVBQUUsQ0FBQztnQkFDTCxPQUFPLEVBQUUsRUFBRTtnQkFDWCxLQUFLLEVBQUUsQ0FBQzthQUNUO1lBQ0Q7Z0JBQ0UsRUFBRSxFQUFFLENBQUM7Z0JBQ0wsT0FBTyxFQUFFLENBQUM7Z0JBQ1YsS0FBSyxFQUFFLENBQUM7YUFDVDtZQUNEO2dCQUNFLEVBQUUsRUFBRSxFQUFFO2dCQUNOLE9BQU8sRUFBRSxDQUFDO2dCQUNWLEtBQUssRUFBRSxDQUFDO2FBQ1Q7WUFDRDtnQkFDRSxFQUFFLEVBQUUsRUFBRTtnQkFDTixPQUFPLEVBQUUsQ0FBQztnQkFDVixLQUFLLEVBQUUsQ0FBQzthQUNUO1NBQ0YsQ0FBQyxDQUFDO0lBQ0wsQ0FBQyxDQUFDLENBQUM7QUFDTCxDQUFDLENBQUMsQ0FBQztBQUVILFFBQVEsQ0FBQyx5Q0FBeUMsRUFBRSxHQUFHLEVBQUU7SUFDdkQsSUFBSSxDQUFDLHNDQUFzQyxFQUFFLEdBQUcsRUFBRTtRQUNoRCxNQUFNLGVBQWUsR0FBRyxJQUFBLGdDQUFtQixFQUN6QyxHQUFHLEVBQ0gsRUFBRSxFQUNGLG9CQUFZLENBQUMsWUFBWSxFQUN6QixHQUFHLEVBQ0gsQ0FBQyxDQUNGLENBQUM7UUFDRixNQUFNLENBQUMsZUFBZSxDQUFDLFdBQVcsRUFBRSxDQUFDLENBQUMsT0FBTyxDQUFDLENBQUMsQ0FBQyxDQUFDO0lBQ25ELENBQUMsQ0FBQyxDQUFDO0lBQ0gsSUFBSSxDQUFDLHNDQUFzQyxFQUFFLEdBQUcsRUFBRTtRQUNoRCxNQUFNLGVBQWUsR0FBRyxJQUFBLGdDQUFtQixFQUN6QyxJQUFJLEVBQ0osRUFBRSxFQUNGLG9CQUFZLENBQUMsWUFBWSxFQUN6QixJQUFJLEVBQ0osQ0FBQyxDQUNGLENBQUM7UUFDRixNQUFNLENBQUMsZUFBZSxDQUFDLFdBQVcsRUFBRSxDQUFDLENBQUMsT0FBTyxDQUFDLEVBQUUsQ0FBQyxDQUFDO0lBQ3BELENBQUMsQ0FBQyxDQUFDO0FBQ0wsQ0FBQyxDQUFDLENBQUMiLCJzb3VyY2VzQ29udGVudCI6WyJpbXBvcnQgeyBTaXplIH0gZnJvbSBcImF3cy1jZGstbGliXCI7XG5pbXBvcnQgeyBjYWxjTWVtb3J5Rm9vdHByaW50LCBjYWxjVGVuc29yUGFyYWxsZWwgfSBmcm9tIFwiLi9jYWxjdWxhdG9yXCI7XG5pbXBvcnQgeyBEYXRhVHlwZUJpdHMgfSBmcm9tIFwiLi9tb2RlbFwiO1xuaW1wb3J0IHsgTmV1cm9ueEluc3RhbmNlVHlwZSB9IGZyb20gXCIuL25ldXJvbngtaW5zdGFuY2UtdHlwZVwiO1xuXG5kZXNjcmliZShcIldoZW4gdGhlIGF0dGVudGlvbiBoZWFkIGlzIGRpdmlzaWJsZVwiLCAoKSA9PiB7XG4gIHRlc3QoXCJJbmYyXCIsICgpID0+IHtcbiAgICBjb25zdCB0ZW5zb3JQYXJhbGxlbCA9IGNhbGNUZW5zb3JQYXJhbGxlbChcbiAgICAgIE5ldXJvbnhJbnN0YW5jZVR5cGUuSU5GMl80OFhMQVJHRSxcbiAgICAgIFNpemUuZ2liaWJ5dGVzKDE2KSxcbiAgICAgIDMyLFxuICAgICk7XG4gICAgZXhwZWN0KHRlbnNvclBhcmFsbGVsKS50b0VxdWFsKFtcbiAgICAgIHtcbiAgICAgICAgdHA6IDEsXG4gICAgICAgIHdvcmtlcnM6IDI0LFxuICAgICAgICB1c2FnZTogMSxcbiAgICAgIH0sXG4gICAgICB7XG4gICAgICAgIHRwOiAyLFxuICAgICAgICB3b3JrZXJzOiAxMixcbiAgICAgICAgdXNhZ2U6IDEsXG4gICAgICB9LFxuICAgICAge1xuICAgICAgICB0cDogNCxcbiAgICAgICAgd29ya2VyczogNixcbiAgICAgICAgdXNhZ2U6IDEsXG4gICAgICB9LFxuICAgICAge1xuICAgICAgICB0cDogOCxcbiAgICAgICAgd29ya2VyczogMyxcbiAgICAgICAgdXNhZ2U6IDEsXG4gICAgICB9LFxuICAgIF0pO1xuICB9KTtcbiAgdGVzdChcIlRybjFcIiwgKCkgPT4ge1xuICAgIGNvbnN0IHRlbnNvclBhcmFsbGVsID0gY2FsY1RlbnNvclBhcmFsbGVsKFxuICAgICAgTmV1cm9ueEluc3RhbmNlVHlwZS5UUk4xXzMyWExBUkdFLFxuICAgICAgU2l6ZS5naWJpYnl0ZXMoMTYpLFxuICAgICAgNDgsXG4gICAgKTtcbiAgICBleHBlY3QodGVuc29yUGFyYWxsZWwpLnRvRXF1YWwoW1xuICAgICAge1xuICAgICAgICB0cDogMSxcbiAgICAgICAgd29ya2VyczogMzIsXG4gICAgICAgIHVzYWdlOiAxLFxuICAgICAgfSxcbiAgICAgIHtcbiAgICAgICAgdHA6IDIsXG4gICAgICAgIHdvcmtlcnM6IDE2LFxuICAgICAgICB1c2FnZTogMSxcbiAgICAgIH0sXG4gICAgICB7XG4gICAgICAgIHRwOiA4LFxuICAgICAgICB3b3JrZXJzOiA0LFxuICAgICAgICB1c2FnZTogMSxcbiAgICAgIH0sXG4gICAgICB7XG4gICAgICAgIHRwOiAxNixcbiAgICAgICAgd29ya2VyczogMixcbiAgICAgICAgdXNhZ2U6IDEsXG4gICAgICB9LFxuICAgIF0pO1xuICB9KTtcbn0pO1xuZGVzY3JpYmUoXCJUaGUgb25lIHdpdGggbG93IHVzYWdlIGlzIHNob3duIGxhc3RcIiwgKCkgPT4ge1xuICB0ZXN0KFwiSW5mMlwiLCAoKSA9PiB7XG4gICAgY29uc3QgdGVuc29yUGFyYWxsZWwgPSBjYWxjVGVuc29yUGFyYWxsZWwoXG4gICAgICBOZXVyb254SW5zdGFuY2VUeXBlLklORjJfMjRYTEFSR0UsXG4gICAgICBTaXplLmdpYmlieXRlcygxNiksXG4gICAgICAzMixcbiAgICApO1xuICAgIGV4cGVjdCh0ZW5zb3JQYXJhbGxlbCkudG9FcXVhbChbXG4gICAgICB7XG4gICAgICAgIHRwOiAxLFxuICAgICAgICB3b3JrZXJzOiAxMixcbiAgICAgICAgdXNhZ2U6IDEsXG4gICAgICB9LFxuICAgICAge1xuICAgICAgICB0cDogMixcbiAgICAgICAgd29ya2VyczogNixcbiAgICAgICAgdXNhZ2U6IDEsXG4gICAgICB9LFxuICAgICAge1xuICAgICAgICB0cDogNCxcbiAgICAgICAgd29ya2VyczogMyxcbiAgICAgICAgdXNhZ2U6IDEsXG4gICAgICB9LFxuICAgICAge1xuICAgICAgICB0cDogOCxcbiAgICAgICAgd29ya2VyczogMSxcbiAgICAgICAgdXNhZ2U6IDAuNjY2NjY2NjY2NjY2NjY2NixcbiAgICAgIH0sXG4gICAgXSk7XG4gIH0pO1xuICB0ZXN0KFwiVHJuMVwiLCAoKSA9PiB7XG4gICAgY29uc3QgdGVuc29yUGFyYWxsZWwgPSBjYWxjVGVuc29yUGFyYWxsZWwoXG4gICAgICBOZXVyb254SW5zdGFuY2VUeXBlLlRSTjFfMzJYTEFSR0UsXG4gICAgICBTaXplLmdpYmlieXRlcygxNiksXG4gICAgICAzMixcbiAgICApO1xuICAgIGV4cGVjdCh0ZW5zb3JQYXJhbGxlbCkudG9FcXVhbChbXG4gICAgICB7XG4gICAgICAgIHRwOiAxLFxuICAgICAgICB3b3JrZXJzOiAzMixcbiAgICAgICAgdXNhZ2U6IDEsXG4gICAgICB9LFxuICAgICAge1xuICAgICAgICB0cDogMixcbiAgICAgICAgd29ya2VyczogMTYsXG4gICAgICAgIHVzYWdlOiAxLFxuICAgICAgfSxcbiAgICAgIHtcbiAgICAgICAgdHA6IDgsXG4gICAgICAgIHdvcmtlcnM6IDQsXG4gICAgICAgIHVzYWdlOiAxLFxuICAgICAgfSxcbiAgICAgIHtcbiAgICAgICAgdHA6IDE2LFxuICAgICAgICB3b3JrZXJzOiAyLFxuICAgICAgICB1c2FnZTogMSxcbiAgICAgIH0sXG4gICAgICB7XG4gICAgICAgIHRwOiAzMixcbiAgICAgICAgd29ya2VyczogMSxcbiAgICAgICAgdXNhZ2U6IDEsXG4gICAgICB9LFxuICAgIF0pO1xuICB9KTtcbn0pO1xuXG5kZXNjcmliZShcIldoZW4gdGhlIG1lbW9yeSBmb290cHJpbnQgaXMgY2FsY3VsYXRlZFwiLCAoKSA9PiB7XG4gIHRlc3QoXCJNaWxsaW9uIHNpemUgbW9kZWxzIGFyZSBjYWxjdWxhdGFibGVcIiwgKCkgPT4ge1xuICAgIGNvbnN0IG1lbW9yeUZvb3RwcmludCA9IGNhbGNNZW1vcnlGb290cHJpbnQoXG4gICAgICA1NzYsXG4gICAgICAzMCxcbiAgICAgIERhdGFUeXBlQml0cy5CRjE2X09SX0ZQMTYsXG4gICAgICAxMjgsXG4gICAgICAxLFxuICAgICk7XG4gICAgZXhwZWN0KG1lbW9yeUZvb3RwcmludC50b0dpYmlieXRlcygpKS50b0VxdWFsKDEpO1xuICB9KTtcbiAgdGVzdChcIkJpbGxpb24gc2l6ZSBtb2RlbHMgYXJlIGNhbGN1bGF0YWJsZVwiLCAoKSA9PiB7XG4gICAgY29uc3QgbWVtb3J5Rm9vdHByaW50ID0gY2FsY01lbW9yeUZvb3RwcmludChcbiAgICAgIDYxNDQsXG4gICAgICA0OCxcbiAgICAgIERhdGFUeXBlQml0cy5CRjE2X09SX0ZQMTYsXG4gICAgICAxMDI0LFxuICAgICAgMSxcbiAgICApO1xuICAgIGV4cGVjdChtZW1vcnlGb290cHJpbnQudG9HaWJpYnl0ZXMoKSkudG9FcXVhbCg0Mik7XG4gIH0pO1xufSk7XG4iXX0=