{ "cells": [ { "cell_type": "code", "execution_count": 3, "metadata": {}, "outputs": [], "source": [ "from utils import activation_memory, param_grads_opt" ] }, { "cell_type": "code", "execution_count": null, "metadata": {}, "outputs": [], "source": [ "tpl = (\n", " a, # attention heads\n", " b, # micro batch size\n", " h, # hidden dimension size\n", " L, # number of layers\n", " s, # sequence length\n", " mixed=True,\n", " recomputation=None\n", ")" ] } ], "metadata": { "kernelspec": { "display_name": "jupyter", "language": "python", "name": "python3" }, "language_info": { "codemirror_mode": { "name": "ipython", "version": 3 }, "file_extension": ".py", "mimetype": "text/x-python", "name": "python", "nbconvert_exporter": "python", "pygments_lexer": "ipython3", "version": "3.10.14" } }, "nbformat": 4, "nbformat_minor": 2 }