aws-cdk-neuronx-patterns
Version:
> [!WARNING] > This library is experimental module.
61 lines • 12 kB
JavaScript
;
Object.defineProperty(exports, "__esModule", { value: true });
exports.calcMemoryFootprint = calcMemoryFootprint;
exports.calcTensorParallel = calcTensorParallel;
exports.inferMemoryFootprintFromParameters = inferMemoryFootprintFromParameters;
const aws_cdk_lib_1 = require("aws-cdk-lib");
const model_1 = require("./model");
function calcDataSize(dataTypeBits) {
return dataTypeBits / 8;
}
function calcWeightMemoryFootprint(layers, embeddingDimension, dataTypeBits) {
const dataSize = calcDataSize(dataTypeBits);
return aws_cdk_lib_1.Size.bytes(12 * layers * embeddingDimension ** 2 * dataSize);
}
function calcKvCacheMemoryFootprint(layers, embeddingDimension, dataTypeBits, maxSequenceLength, batchSize) {
const dataSize = calcDataSize(dataTypeBits);
return aws_cdk_lib_1.Size.bytes(2 * layers * embeddingDimension * dataSize * maxSequenceLength * batchSize);
}
function calcMemoryFootprint(embeddingDimension, layers, dataTypeBits, maxSequenceLength, batchSize) {
// # OPT-66B example (BF16, Inf2)
// # n_layer=64, n_ctx=2048, n_embd=9216, batch=16
// weight_mem_footprint = 12 x 64 x 9216^2 x 2 = 121.5 GiB
// KV_cache_mem_footprint = 16 x 64 x 2048 x 9216 x 2 x 2 = 72 GiB
// mem_footprint = 121.5GiB + 72GiB = 193.5 GiB
// num_neuron_cores = ceil_to_closest_supported_size (193.5GiB / 16GiB, Inf2)
// = ceil_to_closest_supported_size (12.1) = 24
// ## Currently, the Neuron runtime supports tensor-parallelism degrees 2, 8, and 32 on Trn1
// ## and supports tensor-parallelism degrees 2, 4, 8, 12 and 24 on Inf2.
const weightMemFootprint = calcWeightMemoryFootprint(layers, embeddingDimension, dataTypeBits);
const kvCacheMemFootprint = calcKvCacheMemoryFootprint(layers, embeddingDimension, dataTypeBits, maxSequenceLength, batchSize);
const bytes = aws_cdk_lib_1.Size.bytes(weightMemFootprint.toBytes() + kvCacheMemFootprint.toBytes());
return aws_cdk_lib_1.Size.gibibytes(Math.ceil(bytes.toGibibytes({ rounding: aws_cdk_lib_1.SizeRoundingBehavior.NONE })));
}
function calcTensorParallel(neuronxInstanceType, memoryFootprint, attentionHeads) {
const toGibibytesCeil = (size) => Math.ceil(size.toGibibytes({
rounding: aws_cdk_lib_1.SizeRoundingBehavior.NONE,
}));
return (neuronxInstanceType.supportedTensorParallelism
// Attention heads need to be divisible by TP
.filter((tp) => !attentionHeads || attentionHeads % tp === 0)
// Fits into Neuron Core's memory capacity
.filter((tp) => tp *
toGibibytesCeil(neuronxInstanceType.acceleratorChips.acceleratorMemory) >=
toGibibytesCeil(memoryFootprint))
.map((tp) => {
const workers = Math.floor(neuronxInstanceType.acceleratorChips.neuronxCores / tp);
return {
tp,
workers,
usage: (tp * workers) / neuronxInstanceType.acceleratorChips.neuronxCores,
};
})
.sort((a, b) => b.usage - a.usage));
}
function inferMemoryFootprintFromParameters(parameters, maxSequenceLength, batchSize) {
const dataSize = calcDataSize(model_1.DataTypeBits.BF16_OR_FP16);
const base = aws_cdk_lib_1.Size.gibibytes(parameters.toBillion() * dataSize * 1.2);
const additional = aws_cdk_lib_1.Size.gibibytes(((base.toGibibytes() * (maxSequenceLength / 1000)) / 10) * batchSize);
return aws_cdk_lib_1.Size.gibibytes(base.toGibibytes() + additional.toGibibytes());
}
//# sourceMappingURL=data:application/json;base64,eyJ2ZXJzaW9uIjozLCJmaWxlIjoiY2FsY3VsYXRvci5qcyIsInNvdXJjZVJvb3QiOiIiLCJzb3VyY2VzIjpbIi4uLy4uLy4uL3NyYy9iYXNlL25ldXJvbngvY2FsY3VsYXRvci50cyJdLCJuYW1lcyI6W10sIm1hcHBpbmdzIjoiOztBQTRCQSxrREFrQ0M7QUFFRCxnREFxQ0M7QUFFRCxnRkFXQztBQWxIRCw2Q0FBeUQ7QUFDekQsbUNBQW1EO0FBR25ELFNBQVMsWUFBWSxDQUFDLFlBQTBCO0lBQzlDLE9BQU8sWUFBWSxHQUFHLENBQUMsQ0FBQztBQUMxQixDQUFDO0FBQ0QsU0FBUyx5QkFBeUIsQ0FDaEMsTUFBYyxFQUNkLGtCQUEwQixFQUMxQixZQUEwQjtJQUUxQixNQUFNLFFBQVEsR0FBRyxZQUFZLENBQUMsWUFBWSxDQUFDLENBQUM7SUFDNUMsT0FBTyxrQkFBSSxDQUFDLEtBQUssQ0FBQyxFQUFFLEdBQUcsTUFBTSxHQUFHLGtCQUFrQixJQUFJLENBQUMsR0FBRyxRQUFRLENBQUMsQ0FBQztBQUN0RSxDQUFDO0FBQ0QsU0FBUywwQkFBMEIsQ0FDakMsTUFBYyxFQUNkLGtCQUEwQixFQUMxQixZQUEwQixFQUMxQixpQkFBeUIsRUFDekIsU0FBaUI7SUFFakIsTUFBTSxRQUFRLEdBQUcsWUFBWSxDQUFDLFlBQVksQ0FBQyxDQUFDO0lBQzVDLE9BQU8sa0JBQUksQ0FBQyxLQUFLLENBQ2YsQ0FBQyxHQUFHLE1BQU0sR0FBRyxrQkFBa0IsR0FBRyxRQUFRLEdBQUcsaUJBQWlCLEdBQUcsU0FBUyxDQUMzRSxDQUFDO0FBQ0osQ0FBQztBQUVELFNBQWdCLG1CQUFtQixDQUNqQyxrQkFBMEIsRUFDMUIsTUFBYyxFQUNkLFlBQTBCLEVBQzFCLGlCQUF5QixFQUN6QixTQUFpQjtJQUVqQixpQ0FBaUM7SUFDakMsa0RBQWtEO0lBQ2xELDBEQUEwRDtJQUMxRCxrRUFBa0U7SUFDbEUsK0NBQStDO0lBQy9DLDZFQUE2RTtJQUM3RSxnRUFBZ0U7SUFDaEUsNkdBQTZHO0lBQzdHLDBGQUEwRjtJQUMxRixNQUFNLGtCQUFrQixHQUFHLHlCQUF5QixDQUNsRCxNQUFNLEVBQ04sa0JBQWtCLEVBQ2xCLFlBQVksQ0FDYixDQUFDO0lBQ0YsTUFBTSxtQkFBbUIsR0FBRywwQkFBMEIsQ0FDcEQsTUFBTSxFQUNOLGtCQUFrQixFQUNsQixZQUFZLEVBQ1osaUJBQWlCLEVBQ2pCLFNBQVMsQ0FDVixDQUFDO0lBQ0YsTUFBTSxLQUFLLEdBQUcsa0JBQUksQ0FBQyxLQUFLLENBQ3RCLGtCQUFrQixDQUFDLE9BQU8sRUFBRSxHQUFHLG1CQUFtQixDQUFDLE9BQU8sRUFBRSxDQUM3RCxDQUFDO0lBQ0YsT0FBTyxrQkFBSSxDQUFDLFNBQVMsQ0FDbkIsSUFBSSxDQUFDLElBQUksQ0FBQyxLQUFLLENBQUMsV0FBVyxDQUFDLEVBQUUsUUFBUSxFQUFFLGtDQUFvQixDQUFDLElBQUksRUFBRSxDQUFDLENBQUMsQ0FDdEUsQ0FBQztBQUNKLENBQUM7QUFFRCxTQUFnQixrQkFBa0IsQ0FDaEMsbUJBQXlDLEVBQ3pDLGVBQXFCLEVBQ3JCLGNBQXVCO0lBRXZCLE1BQU0sZUFBZSxHQUFHLENBQUMsSUFBVSxFQUFFLEVBQUUsQ0FDckMsSUFBSSxDQUFDLElBQUksQ0FDUCxJQUFJLENBQUMsV0FBVyxDQUFDO1FBQ2YsUUFBUSxFQUFFLGtDQUFvQixDQUFDLElBQUk7S0FDcEMsQ0FBQyxDQUNILENBQUM7SUFDSixPQUFPLENBQ0wsbUJBQW1CLENBQUMsMEJBQTBCO1FBQzVDLDZDQUE2QztTQUM1QyxNQUFNLENBQUMsQ0FBQyxFQUFFLEVBQUUsRUFBRSxDQUFDLENBQUMsY0FBYyxJQUFJLGNBQWMsR0FBRyxFQUFFLEtBQUssQ0FBQyxDQUFDO1FBQzdELDBDQUEwQztTQUN6QyxNQUFNLENBQ0wsQ0FBQyxFQUFFLEVBQUUsRUFBRSxDQUNMLEVBQUU7UUFDQSxlQUFlLENBQ2IsbUJBQW1CLENBQUMsZ0JBQWdCLENBQUMsaUJBQWlCLENBQ3ZEO1FBQ0gsZUFBZSxDQUFDLGVBQWUsQ0FBQyxDQUNuQztTQUNBLEdBQUcsQ0FBQyxDQUFDLEVBQUUsRUFBRSxFQUFFO1FBQ1YsTUFBTSxPQUFPLEdBQUcsSUFBSSxDQUFDLEtBQUssQ0FDeEIsbUJBQW1CLENBQUMsZ0JBQWdCLENBQUMsWUFBWSxHQUFHLEVBQUUsQ0FDdkQsQ0FBQztRQUNGLE9BQU87WUFDTCxFQUFFO1lBQ0YsT0FBTztZQUNQLEtBQUssRUFDSCxDQUFDLEVBQUUsR0FBRyxPQUFPLENBQUMsR0FBRyxtQkFBbUIsQ0FBQyxnQkFBZ0IsQ0FBQyxZQUFZO1NBQ3JFLENBQUM7SUFDSixDQUFDLENBQUM7U0FDRCxJQUFJLENBQUMsQ0FBQyxDQUFDLEVBQUUsQ0FBQyxFQUFFLEVBQUUsQ0FBQyxDQUFDLENBQUMsS0FBSyxHQUFHLENBQUMsQ0FBQyxLQUFLLENBQUMsQ0FDckMsQ0FBQztBQUNKLENBQUM7QUFFRCxTQUFnQixrQ0FBa0MsQ0FDaEQsVUFBc0IsRUFDdEIsaUJBQXlCLEVBQ3pCLFNBQWlCO0lBRWpCLE1BQU0sUUFBUSxHQUFHLFlBQVksQ0FBQyxvQkFBWSxDQUFDLFlBQVksQ0FBQyxDQUFDO0lBQ3pELE1BQU0sSUFBSSxHQUFHLGtCQUFJLENBQUMsU0FBUyxDQUFDLFVBQVUsQ0FBQyxTQUFTLEVBQUUsR0FBRyxRQUFRLEdBQUcsR0FBRyxDQUFDLENBQUM7SUFDckUsTUFBTSxVQUFVLEdBQUcsa0JBQUksQ0FBQyxTQUFTLENBQy9CLENBQUMsQ0FBQyxJQUFJLENBQUMsV0FBVyxFQUFFLEdBQUcsQ0FBQyxpQkFBaUIsR0FBRyxJQUFJLENBQUMsQ0FBQyxHQUFHLEVBQUUsQ0FBQyxHQUFHLFNBQVMsQ0FDckUsQ0FBQztJQUNGLE9BQU8sa0JBQUksQ0FBQyxTQUFTLENBQUMsSUFBSSxDQUFDLFdBQVcsRUFBRSxHQUFHLFVBQVUsQ0FBQyxXQUFXLEVBQUUsQ0FBQyxDQUFDO0FBQ3ZFLENBQUMiLCJzb3VyY2VzQ29udGVudCI6WyJpbXBvcnQgeyBTaXplLCBTaXplUm91bmRpbmdCZWhhdmlvciB9IGZyb20gXCJhd3MtY2RrLWxpYlwiO1xuaW1wb3J0IHsgRGF0YVR5cGVCaXRzLCBQYXJhbWV0ZXJzIH0gZnJvbSBcIi4vbW9kZWxcIjtcbmltcG9ydCB7IElOZXVyb254SW5zdGFuY2VUeXBlIH0gZnJvbSBcIi4vbmV1cm9ueC1pbnN0YW5jZS10eXBlXCI7XG5cbmZ1bmN0aW9uIGNhbGNEYXRhU2l6ZShkYXRhVHlwZUJpdHM6IERhdGFUeXBlQml0cykge1xuICByZXR1cm4gZGF0YVR5cGVCaXRzIC8gODtcbn1cbmZ1bmN0aW9uIGNhbGNXZWlnaHRNZW1vcnlGb290cHJpbnQoXG4gIGxheWVyczogbnVtYmVyLFxuICBlbWJlZGRpbmdEaW1lbnNpb246IG51bWJlcixcbiAgZGF0YVR5cGVCaXRzOiBEYXRhVHlwZUJpdHMsXG4pIHtcbiAgY29uc3QgZGF0YVNpemUgPSBjYWxjRGF0YVNpemUoZGF0YVR5cGVCaXRzKTtcbiAgcmV0dXJuIFNpemUuYnl0ZXMoMTIgKiBsYXllcnMgKiBlbWJlZGRpbmdEaW1lbnNpb24gKiogMiAqIGRhdGFTaXplKTtcbn1cbmZ1bmN0aW9uIGNhbGNLdkNhY2hlTWVtb3J5Rm9vdHByaW50KFxuICBsYXllcnM6IG51bWJlcixcbiAgZW1iZWRkaW5nRGltZW5zaW9uOiBudW1iZXIsXG4gIGRhdGFUeXBlQml0czogRGF0YVR5cGVCaXRzLFxuICBtYXhTZXF1ZW5jZUxlbmd0aDogbnVtYmVyLFxuICBiYXRjaFNpemU6IG51bWJlcixcbikge1xuICBjb25zdCBkYXRhU2l6ZSA9IGNhbGNEYXRhU2l6ZShkYXRhVHlwZUJpdHMpO1xuICByZXR1cm4gU2l6ZS5ieXRlcyhcbiAgICAyICogbGF5ZXJzICogZW1iZWRkaW5nRGltZW5zaW9uICogZGF0YVNpemUgKiBtYXhTZXF1ZW5jZUxlbmd0aCAqIGJhdGNoU2l6ZSxcbiAgKTtcbn1cblxuZXhwb3J0IGZ1bmN0aW9uIGNhbGNNZW1vcnlGb290cHJpbnQoXG4gIGVtYmVkZGluZ0RpbWVuc2lvbjogbnVtYmVyLFxuICBsYXllcnM6IG51bWJlcixcbiAgZGF0YVR5cGVCaXRzOiBEYXRhVHlwZUJpdHMsXG4gIG1heFNlcXVlbmNlTGVuZ3RoOiBudW1iZXIsXG4gIGJhdGNoU2l6ZTogbnVtYmVyLFxuKSB7XG4gIC8vICMgT1BULTY2QiBleGFtcGxlIChCRjE2LCBJbmYyKVxuICAvLyAjIG5fbGF5ZXI9NjQsIG5fY3R4PTIwNDgsIG5fZW1iZD05MjE2LCBiYXRjaD0xNlxuICAvLyB3ZWlnaHRfbWVtX2Zvb3RwcmludCA9IDEyIHggNjQgeCA5MjE2XjIgeCAyID0gMTIxLjUgR2lCXG4gIC8vIEtWX2NhY2hlX21lbV9mb290cHJpbnQgPSAxNiB4IDY0IHggMjA0OCB4IDkyMTYgeCAyIHggMiA9IDcyIEdpQlxuICAvLyBtZW1fZm9vdHByaW50ID0gMTIxLjVHaUIgKyA3MkdpQiA9IDE5My41IEdpQlxuICAvLyBudW1fbmV1cm9uX2NvcmVzID0gY2VpbF90b19jbG9zZXN0X3N1cHBvcnRlZF9zaXplICgxOTMuNUdpQiAvIDE2R2lCLCBJbmYyKVxuICAvLyAgICAgICAgICAgICAgICAgID0gY2VpbF90b19jbG9zZXN0X3N1cHBvcnRlZF9zaXplICgxMi4xKSA9IDI0XG4gIC8vICAgICAgICAgICAgICAgICAgIyMgQ3VycmVudGx5LCB0aGUgTmV1cm9uIHJ1bnRpbWUgc3VwcG9ydHMgdGVuc29yLXBhcmFsbGVsaXNtIGRlZ3JlZXMgMiwgOCwgYW5kIDMyIG9uIFRybjFcbiAgLy8gICAgICAgICAgICAgICAgICAjIyBhbmQgc3VwcG9ydHMgdGVuc29yLXBhcmFsbGVsaXNtIGRlZ3JlZXMgMiwgNCwgOCwgMTIgYW5kIDI0IG9uIEluZjIuXG4gIGNvbnN0IHdlaWdodE1lbUZvb3RwcmludCA9IGNhbGNXZWlnaHRNZW1vcnlGb290cHJpbnQoXG4gICAgbGF5ZXJzLFxuICAgIGVtYmVkZGluZ0RpbWVuc2lvbixcbiAgICBkYXRhVHlwZUJpdHMsXG4gICk7XG4gIGNvbnN0IGt2Q2FjaGVNZW1Gb290cHJpbnQgPSBjYWxjS3ZDYWNoZU1lbW9yeUZvb3RwcmludChcbiAgICBsYXllcnMsXG4gICAgZW1iZWRkaW5nRGltZW5zaW9uLFxuICAgIGRhdGFUeXBlQml0cyxcbiAgICBtYXhTZXF1ZW5jZUxlbmd0aCxcbiAgICBiYXRjaFNpemUsXG4gICk7XG4gIGNvbnN0IGJ5dGVzID0gU2l6ZS5ieXRlcyhcbiAgICB3ZWlnaHRNZW1Gb290cHJpbnQudG9CeXRlcygpICsga3ZDYWNoZU1lbUZvb3RwcmludC50b0J5dGVzKCksXG4gICk7XG4gIHJldHVybiBTaXplLmdpYmlieXRlcyhcbiAgICBNYXRoLmNlaWwoYnl0ZXMudG9HaWJpYnl0ZXMoeyByb3VuZGluZzogU2l6ZVJvdW5kaW5nQmVoYXZpb3IuTk9ORSB9KSksXG4gICk7XG59XG5cbmV4cG9ydCBmdW5jdGlvbiBjYWxjVGVuc29yUGFyYWxsZWwoXG4gIG5ldXJvbnhJbnN0YW5jZVR5cGU6IElOZXVyb254SW5zdGFuY2VUeXBlLFxuICBtZW1vcnlGb290cHJpbnQ6IFNpemUsXG4gIGF0dGVudGlvbkhlYWRzPzogbnVtYmVyLFxuKSB7XG4gIGNvbnN0IHRvR2liaWJ5dGVzQ2VpbCA9IChzaXplOiBTaXplKSA9PlxuICAgIE1hdGguY2VpbChcbiAgICAgIHNpemUudG9HaWJpYnl0ZXMoe1xuICAgICAgICByb3VuZGluZzogU2l6ZVJvdW5kaW5nQmVoYXZpb3IuTk9ORSxcbiAgICAgIH0pLFxuICAgICk7XG4gIHJldHVybiAoXG4gICAgbmV1cm9ueEluc3RhbmNlVHlwZS5zdXBwb3J0ZWRUZW5zb3JQYXJhbGxlbGlzbVxuICAgICAgLy8gQXR0ZW50aW9uIGhlYWRzIG5lZWQgdG8gYmUgZGl2aXNpYmxlIGJ5IFRQXG4gICAgICAuZmlsdGVyKCh0cCkgPT4gIWF0dGVudGlvbkhlYWRzIHx8IGF0dGVudGlvbkhlYWRzICUgdHAgPT09IDApXG4gICAgICAvLyBGaXRzIGludG8gTmV1cm9uIENvcmUncyBtZW1vcnkgY2FwYWNpdHlcbiAgICAgIC5maWx0ZXIoXG4gICAgICAgICh0cCkgPT5cbiAgICAgICAgICB0cCAqXG4gICAgICAgICAgICB0b0dpYmlieXRlc0NlaWwoXG4gICAgICAgICAgICAgIG5ldXJvbnhJbnN0YW5jZVR5cGUuYWNjZWxlcmF0b3JDaGlwcy5hY2NlbGVyYXRvck1lbW9yeSxcbiAgICAgICAgICAgICkgPj1cbiAgICAgICAgICB0b0dpYmlieXRlc0NlaWwobWVtb3J5Rm9vdHByaW50KSxcbiAgICAgIClcbiAgICAgIC5tYXAoKHRwKSA9PiB7XG4gICAgICAgIGNvbnN0IHdvcmtlcnMgPSBNYXRoLmZsb29yKFxuICAgICAgICAgIG5ldXJvbnhJbnN0YW5jZVR5cGUuYWNjZWxlcmF0b3JDaGlwcy5uZXVyb254Q29yZXMgLyB0cCxcbiAgICAgICAgKTtcbiAgICAgICAgcmV0dXJuIHtcbiAgICAgICAgICB0cCxcbiAgICAgICAgICB3b3JrZXJzLFxuICAgICAgICAgIHVzYWdlOlxuICAgICAgICAgICAgKHRwICogd29ya2VycykgLyBuZXVyb254SW5zdGFuY2VUeXBlLmFjY2VsZXJhdG9yQ2hpcHMubmV1cm9ueENvcmVzLFxuICAgICAgICB9O1xuICAgICAgfSlcbiAgICAgIC5zb3J0KChhLCBiKSA9PiBiLnVzYWdlIC0gYS51c2FnZSlcbiAgKTtcbn1cblxuZXhwb3J0IGZ1bmN0aW9uIGluZmVyTWVtb3J5Rm9vdHByaW50RnJvbVBhcmFtZXRlcnMoXG4gIHBhcmFtZXRlcnM6IFBhcmFtZXRlcnMsXG4gIG1heFNlcXVlbmNlTGVuZ3RoOiBudW1iZXIsXG4gIGJhdGNoU2l6ZTogbnVtYmVyLFxuKSB7XG4gIGNvbnN0IGRhdGFTaXplID0gY2FsY0RhdGFTaXplKERhdGFUeXBlQml0cy5CRjE2X09SX0ZQMTYpO1xuICBjb25zdCBiYXNlID0gU2l6ZS5naWJpYnl0ZXMocGFyYW1ldGVycy50b0JpbGxpb24oKSAqIGRhdGFTaXplICogMS4yKTtcbiAgY29uc3QgYWRkaXRpb25hbCA9IFNpemUuZ2liaWJ5dGVzKFxuICAgICgoYmFzZS50b0dpYmlieXRlcygpICogKG1heFNlcXVlbmNlTGVuZ3RoIC8gMTAwMCkpIC8gMTApICogYmF0Y2hTaXplLFxuICApO1xuICByZXR1cm4gU2l6ZS5naWJpYnl0ZXMoYmFzZS50b0dpYmlieXRlcygpICsgYWRkaXRpb25hbC50b0dpYmlieXRlcygpKTtcbn1cbiJdfQ==