wehub/lightning-ai--litgpt
1
0
Fork 0
Code Issues Pull Requests Actions Deployments Models Agent Notes Packages Projects Releases Wiki Activity

chore: import upstream snapshot with attribution

Browse Source
This commit is contained in:
wehub-resource-sync
2026-07-13 12:47:19 +08:00
commit d5f64af28f
253 changed files with 46818 additions and 0 deletions
Download Patch File Download Diff File Expand all files Collapse all files
.devcontainer/Dockerfile
+9
View File
@@ -0,0 +1,9 @@
# See here for image contents: https://github.com/devcontainers/images/blob/main/src/python/.devcontainer/Dockerfile
# [Choice] Python version (use -bookworm or -bullseye variants on local arm64/Apple Silicon): 3, 3.12, 3.11, 3.10, 3.9, 3.8, 3-bookworm, 3.12-bookworm, 3.11-bookworm, 3.10-bookworm, 3.9-bookworm, 3.8-bookworm, 3-bullseye, 3.12-bullseye, 3.11-bullseye, 3.10-bullseye, 3.9-bullseye, 3.8-bullseye, 3-buster, 3.12-buster, 3.11-buster, 3.10-buster, 3.9-buster, 3.8-buster
ARG VARIANT=3-bookworm
FROM mcr.microsoft.com/devcontainers/python:1-${VARIANT}
# Temporary: Upgrade python packages due to https://cve.mitre.org/cgi-bin/cvename.cgi?name=CVE-2022-40897
# They are installed by the base image (python) which does not have the patch.
RUN python3 -m pip install --upgrade pip setuptools
.devcontainer/devcontainer.json
+105
View File
@@ -0,0 +1,105 @@
// For format details, see https://aka.ms/devcontainer.json. For config options, see the README at:
// https://github.com/microsoft/vscode-dev-containers/tree/v0.194.0/containers/python-3
{
"name": "Python 3 (litgpt)",
"build": {
"dockerfile": "Dockerfile",
"context": "..",
"args": {
"VARIANT": "3.11-bookworm"
}
},
"runArgs": [
// Enable GPU passthrough, requires WSL2 on Windows
//"--gpus=all",
// One of the following options is required for torch multiprocessing
//"--ipc=host",
//"--shm-size=4gb",
],
// Features to add to the dev container. More info: https://containers.dev/features.
"features": {
"ghcr.io/devcontainers/features/git:1": {},
"ghcr.io/devcontainers/features/git-lfs:1": {},
//"ghcr.io/devcontainers/features/nvidia-cuda:1": {},
"ghcr.io/devcontainers-extra/features/actionlint:1": {},
"ghcr.io/devcontainers-extra/features/pre-commit:2": {},
"ghcr.io/dhoeric/features/act:1": {},
"ghcr.io/devcontainers/features/docker-in-docker:2": {
"version": "latest",
"moby": true
}
},
// Set *default* container specific settings.json values on container create.
"customizations": {
"vscode": {
"settings": {
"editor.tabSize": 4,
"editor.renderWhitespace": "all",
"editor.formatOnSave": true,
"editor.rulers": [120],
"files.exclude": {
"**/__pycache__": true
},
"python.pythonPath": "/usr/local/bin/python",
"python.defaultInterpreterPath": "/usr/local/bin/python",
"python.languageServer": "Pylance",
"python.analysis.autoImportCompletions": true,
"python.analysis.completeFunctionParens": true,
"python.analysis.autoSearchPaths": true,
"python.testing.pytestArgs": ["tests"],
"python.testing.unittestEnabled": false,
"python.testing.pytestEnabled": true,
"code-eol.highlightNonDefault": true,
"code-eol.highlightExtraWhitespace": true,
"autoDocstring.docstringFormat": "google-notypes",
"autoDocstring.guessTypes": true,
"autoDocstring.generateDocstringOnEnter": true,
"autoDocstring.startOnNewLine": true,
"telemetry.telemetryLevel": "off",
"[python]": {
"editor.formatOnSave": true,
"editor.defaultFormatter": "charliermarsh.ruff",
"editor.codeActionsOnSave": {
"source.organizeImports": "always",
"source.fixAll": "always"
}
}
},
// Add the IDs of extensions you want installed when the container is created.
"extensions": [
"ms-python.python",
"ms-python.vscode-pylance",
"ms-toolsai.jupyter",
"GitHub.copilot",
"GitHub.copilot-chat",
"github.vscode-github-actions",
"SanjulaGanepola.github-local-actions",
"charliermarsh.ruff",
"esbenp.prettier-vscode",
"ms-vscode.test-adapter-converter",
"njqdev.vscode-python-typehint",
"KevinRose.vsc-python-indent",
"medo64.render-crlf",
"shardulm94.trailing-spaces",
"nhoizey.gremlins",
"wayou.vscode-todo-highlight",
"Gruntfuggly.todo-tree",
"njpwerner.autodocstring",
"rodolphebarbanneau.python-docstring-highlighter",
"mechatroner.rainbow-csv",
"uctakeoff.vscode-counter",
"bierner.github-markdown-preview",
"yahyabatulu.vscode-markdown-alert",
"ms-vscode-remote.vscode-remote-extensionpack",
"ms-azuretools.vscode-docker",
"redhat.vscode-yaml"
]
}
},
// Use 'forwardPorts' to make a list of ports inside the container available locally.
// "forwardPorts": [],
// Use 'postCreateCommand' to run commands after the container is created.
"postCreateCommand": "pre-commit install && pip install '.[extra,compiler,test]' -U",
// Comment out connect as root instead. More info: https://aka.ms/vscode-remote/containers/non-root.
"remoteUser": "vscode"
}
.github/CODEOWNERS
+21
View File
@@ -0,0 +1,21 @@
# Each line is a file pattern followed by one or more owners.
# These owners will be the default owners for everything in the repo. Unless a later match takes precedence,
# @global-owner1 and @global-owner2 will be requested for review when someone opens a pull request.
* @lianakoleva @k223kim @andyland @t-vi
# Core source
/litgpt/ @lianakoleva @k223kim @andyland
# CI/CD and configs
/.github/ @lianakoleva @k223kim
*.yml @lianakoleva @k223kim
# Docs
/README.md @williamfalcon @lianakoleva
# Retired committers
# @lantiga (Luca Antiga)
# @rasbt (Sebastian Raschka)
# @awaelchli (Adrian Wälchli)
# @borda (Jirka Borovec)
.github/ISSUE_TEMPLATE/ask-a-question.md
+9
View File
@@ -0,0 +1,9 @@
---
name: Ask a Question
about: Ask and answer questions related to LitGPT
title: ''
labels: question
---
Please describe your question here.
.github/ISSUE_TEMPLATE/bug-report.yaml
+56
View File
@@ -0,0 +1,56 @@
name: Bug Report
description: Report errors related to LitGPT
title: "Description"
labels: bug
body:
- type: markdown
attributes:
value: |
Thank you for taking the time to report an issue. Please fill out the details below to help us resolve it.
- type: textarea
id: bug_description
attributes:
label: Bug description
description: A description of the issue.
placeholder: |
Please provide a description of what the bug or issue is.
validations:
required: true
- type: input
attributes:
label: Reproduced in studio
description: >
Create a new Lightning Studio with code that reproduces the issue and share the link.
Also include all the relevant files and data required to reproduce shared issue.
In case the code does not crash, please add assert statements to show what is the real and expected output.
A simple guide on how to create such a studio can be found [here](https://www.youtube.com/watch?v=YcW-2Zt_bFg&ab_channel=LightningAI).
placeholder: https://lightning.ai/...
validations:
required: false
- type: dropdown
id: operating_system
attributes:
label: What operating system are you using?
description: If applicable, please select the operating system where you experienced this issue.
options:
- "Unknown"
- "macOS"
- "Linux"
- "Windows"
validations:
required: true
- type: textarea
id: version
attributes:
label: LitGPT Version
description: |
Please provide details about your LitGPT version by running the following code in your terminal:
```
pip show litgpt | grep Version:
```
validations:
required: false
.github/ISSUE_TEMPLATE/feature-request.md
+9
View File
@@ -0,0 +1,9 @@
---
name: Suggest a Feature
about: Propose a new feature or enhancement
title: ''
labels: enhancement
---
Please describe the feature or enhancement along with the intended usecase.
.github/dependabot.yml
+41
View File
@@ -0,0 +1,41 @@
# Basic dependabot.yml file with
# minimum configuration for two package managers
version: 2
updates:
# Enable version updates for python
- package-ecosystem: "pip"
# Look for a `requirements` in the `root` directory
directory: "/"
# Check for updates once a week
schedule:
interval: "monthly"
# Labels on pull requests for version updates only
labels:
- "dependencies"
pull-request-branch-name:
# Separate sections of the branch name with a hyphen
# for example, `dependabot-npm_and_yarn-next_js-acorn-6.4.1`
separator: "-"
# Allow up to 5 open pull requests for pip dependencies
open-pull-requests-limit: 3
# Enable version updates for GitHub Actions
- package-ecosystem: "github-actions"
directory: "/"
# Check for updates once a week
schedule:
interval: "weekly"
# Labels on pull requests for version updates only
labels:
- "CI / actions"
pull-request-branch-name:
# Separate sections of the branch name with a hyphen
# for example, `dependabot-npm_and_yarn-next_js-acorn-6.4.1`
separator: "-"
# Allow up to 5 open pull requests for GitHub Actions
open-pull-requests-limit: 1
groups:
GHA-updates:
patterns:
- "*"
.github/workflows/check-links.yml
+34
View File
@@ -0,0 +1,34 @@
name: Check hyperlinks
on:
push:
branches:
- main
pull_request:
branches:
- main
jobs:
test:
runs-on: ubuntu-latest
steps:
- uses: actions/checkout@de0fac2e4500dabe0009e67214ff5f5447ce83dd # v6.0.2
- name: Install uv
uses: astral-sh/setup-uv@08807647e7069bb48b6ef5acd8ec9567f424441b # v8.1.0
with:
activate-environment: true
python-version: "3.10"
enable-cache: true
- name: Install dependencies
# a newer version of mistune is incompatible with nbconvert
# pytest>=9 removed the `path` arg from pytest_collect_file; pytest-check-links still uses it
run: uv pip install "mistune<3.1" "pytest<9" pytest-check-links
- name: Check links
run: pytest --check-links README.md tutorials --check-links-ignore "http*"
- name: Minimize uv cache
run: uv cache prune --ci
.github/workflows/cpu-tests.yml
+143
View File
@@ -0,0 +1,143 @@
name: CPU tests
on:
push:
branches: [main]
# Note: using `pull_request` (not `pull_request_target`) for security reasons.
# This means PRs from external forks will NOT have access to secrets (e.g. HF_TOKEN)
# and some tests may fail or be skipped on fork PRs until we find a better solution.
pull_request:
branches: [main]
types: [opened, reopened, ready_for_review, labeled, synchronize]
workflow_dispatch: {}
# lock down all permissions by default
permissions:
contents: read # needed to check out code
checks: write # needed for test results
pull-requests: read # needed for PR metadata
actions: read # needed to use actions
security-events: none
statuses: write # needed to update commit status
concurrency:
group: ${{ github.workflow }}-${{ github.ref }}-${{ github.head_ref }}
cancel-in-progress: ${{ startsWith(github.event_name, 'pull_request') }}
defaults:
run:
shell: bash
env:
HF_HOME: .cache-HF # Define HF_HOME for caching
TRANSFORMERS_CACHE: .cache-HF/transformers
DATASETS_CACHE: .cache-HF/datasets
HF_DATASETS_CACHE: .cache-HF/datasets
UV_TORCH_BACKEND: cpu
jobs:
testing-imports:
runs-on: ${{ matrix.os }}
strategy:
fail-fast: false
matrix:
os: ["ubuntu-22.04", "ubuntu-24.04", "macOS-14", "windows-2022"]
python-version: ["3.10"]
timeout-minutes: 10
steps:
- name: Checkout generic
uses: actions/checkout@de0fac2e4500dabe0009e67214ff5f5447ce83dd # v6.0.2
- name: Install uv
uses: astral-sh/setup-uv@08807647e7069bb48b6ef5acd8ec9567f424441b # v8.1.0
with:
activate-environment: true
python-version: ${{ matrix.python-version }}
enable-cache: true
- name: Install minimal dependencies
run: |
uv sync --no-dev
uv pip list
- name: Testing package imports
# make sure all modules are still importable with only the minimal dependencies available
run: |
modules=$(
find litgpt -type f -name "*.py" | \
sed 's/\.py$//' | sed 's/\//./g' | \
sed 's/.__init__//g' | xargs -I {} echo "import {};"
)
echo "$modules"
python -c "$modules"
- name: Minimize uv cache
run: uv cache prune --ci
pytester:
runs-on: ${{ matrix.os }}
strategy:
fail-fast: false
matrix:
os: ["ubuntu-22.04"]
python-version: ["3.10", "3.11", "3.12", "3.13"]
requires: ["latest"]
include:
- { os: "ubuntu-22.04", python-version: "3.10", requires: "oldest" }
- { os: "windows-2022", python-version: "3.10", requires: "latest" }
- { os: "macOS-14", python-version: "3.10", requires: "latest" }
timeout-minutes: 35
steps:
- name: Checkout generic
uses: actions/checkout@de0fac2e4500dabe0009e67214ff5f5447ce83dd # v6.0.2
- name: Install uv
uses: astral-sh/setup-uv@08807647e7069bb48b6ef5acd8ec9567f424441b # v8.1.0
with:
activate-environment: true
python-version: ${{ matrix.python-version }}
enable-cache: true
# Add caching for HF models and tokenizers
- name: HF cache
uses: actions/cache@27d5ce7f107fe9357f9df03efb73ab90386fccae # v5
continue-on-error: true
with:
path: .cache-HF
key: hf-cache_${{ runner.os }}-py${{ matrix.python-version }}
restore-keys: |
hf-cache_${{ runner.os }}-py${{ matrix.python-version }}
hf-cache_${{ runner.os }}-
hf-cache_
- name: Set min. dependencies
if: matrix.requires == 'oldest'
run: uv run --no-project --with 'lightning-utilities[cli]>=0.15.1' python -m lightning_utilities.cli requirements set-oldest --req_files=pyproject.toml
- name: Install dependencies
run: |
uv sync --all-extras
uv pip list
- name: Run tests
env:
HF_TOKEN: ${{ secrets.HF_TOKEN }}
run: pytest -v litgpt/ tests/ --timeout=180 --durations=100
- name: Show cache
run: uvx py-tree -d 1 .cache-HF
- name: Minimize uv cache
run: uv cache prune --ci
testing-guardian:
runs-on: ubuntu-latest
needs: [pytester, testing-imports]
if: github.event_name == 'pull_request'
steps:
- run: echo "${{ needs.pytester.result }}"
- name: failing...
if: needs.pytester.result == 'failure'
run: exit 1
- name: cancelled or skipped...
if: contains(fromJSON('["cancelled", "skipped"]'), needs.pytester.result)
timeout-minutes: 1
run: sleep 90
.github/workflows/mkdocs-deploy.yml
+38
View File
@@ -0,0 +1,38 @@
name: Deploy MkDocs
on:
push:
branches: [main]
permissions:
contents: write
jobs:
deploy:
runs-on: ubuntu-24.04
steps:
# Step 1: Checkout the repository
- uses: actions/checkout@de0fac2e4500dabe0009e67214ff5f5447ce83dd # v6.0.2
# Step 2: Install uv
- name: Install uv
uses: astral-sh/setup-uv@08807647e7069bb48b6ef5acd8ec9567f424441b # v8.1.0
with:
activate-environment: true
python-version: "3.10"
enable-cache: true
# Step 3: Install MkDocs and dependencies
- run: uv pip install mkdocs mkdocs-material mkdocs-pagetree-plugin
# Step 4: Deploy to GitHub Pages
- run: |
mkdir -p gh-pages/docs
cp -r tutorials/* gh-pages/docs
cd gh-pages
mv docs/mkdocs.yml mkdocs.yml
echo "{{ pagetree }}" > docs/index.md
mkdocs gh-deploy --force
- name: Minimize uv cache
run: uv cache prune --ci
.github/workflows/publish-pkg.yml
+78
View File
@@ -0,0 +1,78 @@
# To create a release, create a tag and push it to GitHub:
#git tag -a "v0.0.1-beta" -m "beta version testing"
#git push --tags
# https://dev.to/iamtekson/publish-package-to-pypi-and-release-new-version-using-github-actions-108k
name: Publish LitGPT to PyPI
on:
push:
tags:
- "v*"
release:
types: [published]
jobs:
build:
name: Build source and wheel distributions
runs-on: ubuntu-latest
steps:
- name: Checkout source
uses: actions/checkout@de0fac2e4500dabe0009e67214ff5f5447ce83dd # v6.0.2
- name: Set up Python
uses: actions/setup-python@a309ff8b426b58ec0e2a45f0f869d46889d02405 # v6.2.0
with:
python-version: "3.x"
cache: "pip"
- name: Build source and wheel distributions
run: |
python -m pip install --upgrade build twine
pip install importlib_metadata==7.2.1
python -m build
twine check --strict dist/*
- uses: actions/upload-artifact@043fb46d1a93c77aae656e7c1c64a875d1fc6a0a # v7.0.1
with:
name: pypi-packages-${{ github.sha }}
path: dist
upload-release-assets:
needs: build
if: github.event_name == 'release'
runs-on: ubuntu-latest
timeout-minutes: 5
steps:
- uses: actions/download-artifact@3e5f45b2cfb9172054b4087a40e8e0b5a5461e7c # v8.0.1
with:
name: pypi-packages-${{ github.sha }}
path: dist
- run: ls -lh dist/
- name: Upload to release
uses: AButler/upload-release-assets@34491005a5d7ec239a784e460807ce844fde7962 # v4.0.0
with:
files: "dist/*"
repo-token: ${{ secrets.GITHUB_TOKEN }}
publish-pypi:
needs: build
if: startsWith(github.event.ref, 'refs/tags') || github.event_name == 'release'
runs-on: ubuntu-latest
timeout-minutes: 5
permissions:
id-token: write
steps:
- uses: actions/download-artifact@3e5f45b2cfb9172054b4087a40e8e0b5a5461e7c # v8.0.1
with:
name: pypi-packages-${{ github.sha }}
path: dist
- run: ls -lh dist/
- name: Publish distribution to PyPI
uses: pypa/gh-action-pypi-publish@cef221092ed1bacb1cc03d23a2d87d1d172e277b # v1.14.0
.gitignore
+26
View File
@@ -0,0 +1,26 @@
.ipynb_checkpoints/
__pycache__
.idea
.DS_Store
*.egg-info
build
dist
.venv
.venv/
.vscode
uv.lock
# data
data
datasets
!litgpt/data
!tests/data
checkpoints
out
wandb
events.out.tfevents*
# test artifacts from tests/test_readme.py
**/custom_finetuning_dataset.json
client.py
**/custom_texts/
.lightning/workflows/tests.yaml
+78
View File
@@ -0,0 +1,78 @@
trigger:
push:
branches: ["main"]
pull_request:
branches: ["main"]
image: "pytorchlightning/lightning-thunder:ubuntu24.04-cuda12.8.1-cudnn-fe1.15.0-py3.12-pt_2.8.0-dev"
machine: "L4_X_2"
interruptible: "true"
timeout: "45" # minutes
parametrize:
matrix:
dependency: ["", "compiler"]
include: []
exclude: []
env:
SKIP_WITH_CI: "1" # skip single tests with CI
NCCL_DEBUG: "INFO"
CUBLAS_WORKSPACE_CONFIG: ":4096:8"
NCCL_IGNORE_DISABLED_P2P: "1"
TORCH_VERSION: "2.8.0"
RUN_ONLY_CUDA_TESTS: "1" # run CUDA tests only
run: |
whereis nvidia
nvidia-smi
python --version
pip --version
pip list
set -ex
echo "Install uv and create virtual environment"
curl -LsSf https://astral.sh/uv/install.sh | sh
[ -f "$HOME/.local/bin/env" ] && . "$HOME/.local/bin/env"
export PATH="$HOME/.local/bin:$PATH"
uv venv .venv --system-site-packages
. .venv/bin/activate
hash -r
uv pip install -q '.[extra,test]' "torch==${TORCH_VERSION}" cffi -U
if [ "${dependency}" == "compiler" ]; then
uv pip uninstall torchvision torchaudio
uv pip install -q '.[compiler,extra,test]' "torch==${TORCH_VERSION}"
python -c "from thunder.executors import nvfuser_available ; assert nvfuser_available(), 'nvFuser is missing!'"
python -c "from thunder.executors.triton_utils import triton_version ; assert triton_version() is not None, 'triton is missing!'"
fi
uv pip list
python -c "import torch ; gpus = torch.cuda.device_count() ; assert gpus >= 2, f'GPU: {gpus}'"
python -c "from torch import __version__ as ver ; assert str(ver).split('+')[0] == '${TORCH_VERSION}', f'PyTorch: installed {ver} but expected ${TORCH_VERSION}'"
pytest -v --durations=100
wget https://raw.githubusercontent.com/Lightning-AI/utilities/main/scripts/run_standalone_tests.sh
PL_RUN_STANDALONE_TESTS=1 bash run_standalone_tests.sh "tests"
if [ "${dependency}" == "compiler" ]; then
uv pip uninstall lightning-thunder transformers
# install thunder from source, so that, thunder.tests will be available
uv pip install -U "lightning-thunder[test] @ git+https://github.com/Lightning-AI/lightning-thunder.git" "torch==${TORCH_VERSION}"
# Pin transformers to match thunder's test_networks.py requirements
# See: https://github.com/Lightning-AI/lightning-thunder/blob/main/requirements/test.txt
# Get transformers version from thunder requirements
TRANSFORMERS_VERSION=$(curl -fsSL https://raw.githubusercontent.com/Lightning-AI/lightning-thunder/main/requirements/test.txt \
| grep '^transformers==' \
| cut -d'=' -f3 \
| cut -d'#' -f1 \
| xargs)
if [ -z "${TRANSFORMERS_VERSION}" ]; then
echo "Error: Could not determine transformers version from lightning-thunder requirements"
exit 1
fi
uv pip install transformers==${TRANSFORMERS_VERSION}
# without env var, it filters out all tests
RUN_ONLY_CUDA_TESTS=0 pytest tests/ext_thunder/test_thunder_networks.py -v
fi
.pre-commit-config.yaml
+89
View File
@@ -0,0 +1,89 @@
# Copyright The Lightning team.
#
# Licensed under the Apache License, Version 2.0 (the "License");
# you may not use this file except in compliance with the License.
# You may obtain a copy of the License at
#
# http://www.apache.org/licenses/LICENSE-2.0
#
# Unless required by applicable law or agreed to in writing, software
# distributed under the License is distributed on an "AS IS" BASIS,
# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
# See the License for the specific language governing permissions and
# limitations under the License.
default_language_version:
python: python3
ci:
autofix_prs: true
autoupdate_commit_msg: "[pre-commit.ci] pre-commit suggestions"
autoupdate_schedule: quarterly
# submodules: true
repos:
- repo: https://github.com/pre-commit/pre-commit-hooks
rev: v6.0.0
hooks:
- id: end-of-file-fixer
- id: trailing-whitespace
exclude: README.md
- id: check-yaml
- id: check-toml
#- id: check-docstring-first
#- id: check-executables-have-shebangs
- id: check-case-conflict
- id: check-added-large-files
args: ["--maxkb=250", "--enforce-all"]
- id: detect-private-key
- repo: https://github.com/codespell-project/codespell
rev: v2.4.2
hooks:
- id: codespell
additional_dependencies: [tomli]
args: ["--write-changes"]
exclude: pyproject.toml
#- repo: https://github.com/crate-ci/typos
# rev: dictgen-v0.3.1
# hooks:
# - id: typos
# args: [] # empty to do not write fixes
# exclude: pyproject.toml
#- repo: https://github.com/executablebooks/mdformat
# rev: 0.7.21
# hooks:
# - id: mdformat
# args: ["--number"]
# additional_dependencies:
# - mdformat-gfm
# - mdformat-black
# - mdformat_frontmatter
- repo: https://github.com/pre-commit/mirrors-prettier
rev: v4.0.0-alpha.8
hooks:
- id: prettier
files: \.(json|yml|yaml|toml)
# https://prettier.io/docs/en/options.html#print-width
args: ["--print-width=140"]
- repo: https://github.com/astral-sh/ruff-pre-commit
rev: v0.15.9
hooks:
- id: ruff
args: ["--fix"]
- id: ruff-format
- id: ruff
- repo: https://github.com/tox-dev/pyproject-fmt
rev: v2.21.0
hooks:
- id: pyproject-fmt
additional_dependencies: [tox]
- repo: https://github.com/abravalheri/validate-pyproject
rev: v0.25
hooks:
- id: validate-pyproject
CITATION.cff
+9
View File
@@ -0,0 +1,9 @@
cff-version: 1.2.0
message: "If you use this software, you can cite it as shown below."
title: "LitGPT"
abstract: "20+ high-performance LLMs with recipes to pretrain, finetune and deploy at scale."
date-released: 2023-03-22
authors:
- name: "The Lightning AI team"
license: "Apache-2.0"
url: "https://github.com/Lightning-AI/litgpt"
CODE_OF_CONDUCT.md
+76
View File
@@ -0,0 +1,76 @@
# Contributor Covenant Code of Conduct
## Our Pledge
In the interest of fostering an open and welcoming environment, we as
contributors and maintainers pledge to making participation in our project and
our community a harassment-free experience for everyone, regardless of age, body
size, disability, ethnicity, sex characteristics, gender identity and expression,
level of experience, education, socioeconomic status, nationality, personal
appearance, race, religion, or sexual identity and orientation.
## Our Standards
Examples of behavior that contributes to creating a positive environment
include:
- Using welcoming and inclusive language
- Being respectful of differing viewpoints and experiences
- Gracefully accepting constructive criticism
- Focusing on what is best for the community
- Showing empathy towards other community members
Examples of unacceptable behavior by participants include:
- The use of sexualized language or imagery and unwelcome sexual attention or
advances
- Trolling, insulting/derogatory comments, and personal or political attacks
- Public or private harassment
- Publishing others' private information, such as a physical or electronic
address, without explicit permission
- Other conduct which could reasonably be considered inappropriate in a
professional setting
## Our Responsibilities
Project maintainers are responsible for clarifying the standards of acceptable
behavior and are expected to take appropriate and fair corrective action in
response to any instances of unacceptable behavior.
Project maintainers have the right and responsibility to remove, edit, or
reject comments, commits, code, wiki edits, issues, and other contributions
that are not aligned to this Code of Conduct, or to ban temporarily or
permanently any contributor for other behaviors that they deem inappropriate,
threatening, offensive, or harmful.
## Scope
This Code of Conduct applies both within project spaces and in public spaces
when an individual is representing the project or its community. Examples of
representing a project or community include using an official project e-mail
address, posting via an official social media account, or acting as an appointed
representative at an online or offline event. Representation of a project may be
further defined and clarified by project maintainers.
## Enforcement
Instances of abusive, harassing, or otherwise unacceptable behavior may be
reported by contacting the project team at community@lightning.ai. All
complaints will be reviewed and investigated and will result in a response that
is deemed necessary and appropriate to the circumstances. The project team is
obligated to maintain confidentiality with regard to the reporter of an incident.
Further details of specific enforcement policies may be posted separately.
Project maintainers who do not follow or enforce the Code of Conduct in good
faith may face temporary or permanent repercussions as determined by other
members of the project's leadership.
## Attribution
This Code of Conduct is adapted from the [Contributor Covenant][homepage], version 1.4,
available at https://www.contributor-covenant.org/version/1/4/code-of-conduct.html
For answers to common questions about this code of conduct, see
https://www.contributor-covenant.org/faq
[homepage]: https://www.contributor-covenant.org
CONTRIBUTING.md
+57
View File
@@ -0,0 +1,57 @@
# Contributing to LitGPT
We welcome all contributions, regardless of your level of experience or hardware. Whether it's a bug fix, a new feature, or an improvement to the docs — we appreciate your help!
## How to contribute
1. **Open an issue** — describe the bug or feature before writing code. This helps us align on scope early.
2. **Fork the repo** and create a branch from `main`.
3. **Make your changes** and add or update relevant tests.
4. **Open a pull request** against `main`. Include a clear description of what changed and why.
## Development setup
```bash
git clone https://github.com/<your-username>/litgpt
cd litgpt
```
```bash
# using uv (recommended)
uv sync --all-extras
# using pip
pip install -e ".[extra,compiler,test]"
```
Install pre-commit hooks to catch style issues before pushing:
```bash
# using uvx
uvx pre-commit install # install hooks
uvx pre-commit run --all-files # run manually
# using pip
pip install pre-commit
pre-commit install # install hooks
pre-commit run --all-files # run manually
```
## Running tests
```bash
pytest tests/
```
## Guidelines
- Keep pull requests focused — one logical change per PR.
- Write tests for new functionality.
- Follow the existing code style (enforced via [ruff](https://docs.astral.sh/ruff/) and pre-commit).
- All code should be your own original work; third-party snippets must be attributed.
## Community
- [Request a feature or report a bug](https://github.com/Lightning-AI/litgpt/issues)
- [Contribution tutorial](https://lightning.ai/pages/community/tutorial/how-to-contribute-to-litgpt/)
- [Join our Discord](https://discord.gg/VptPCZkGNa)
LICENSE.md
+201
View File
@@ -0,0 +1,201 @@
Apache License
Version 2.0, January 2004
http://www.apache.org/licenses/
TERMS AND CONDITIONS FOR USE, REPRODUCTION, AND DISTRIBUTION
1. Definitions.
"License" shall mean the terms and conditions for use, reproduction,
and distribution as defined by Sections 1 through 9 of this document.
"Licensor" shall mean the copyright owner or entity authorized by
the copyright owner that is granting the License.
"Legal Entity" shall mean the union of the acting entity and all
other entities that control, are controlled by, or are under common
control with that entity. For the purposes of this definition,
"control" means (i) the power, direct or indirect, to cause the
direction or management of such entity, whether by contract or
otherwise, or (ii) ownership of fifty percent (50%) or more of the
outstanding shares, or (iii) beneficial ownership of such entity.
"You" (or "Your") shall mean an individual or Legal Entity
exercising permissions granted by this License.
"Source" form shall mean the preferred form for making modifications,
including but not limited to software source code, documentation
source, and configuration files.
"Object" form shall mean any form resulting from mechanical
transformation or translation of a Source form, including but
not limited to compiled object code, generated documentation,
and conversions to other media types.
"Work" shall mean the work of authorship, whether in Source or
Object form, made available under the License, as indicated by a
copyright notice that is included in or attached to the work
(an example is provided in the Appendix below).
"Derivative Works" shall mean any work, whether in Source or Object
form, that is based on (or derived from) the Work and for which the
editorial revisions, annotations, elaborations, or other modifications
represent, as a whole, an original work of authorship. For the purposes
of this License, Derivative Works shall not include works that remain
separable from, or merely link (or bind by name) to the interfaces of,
the Work and Derivative Works thereof.
"Contribution" shall mean any work of authorship, including
the original version of the Work and any modifications or additions
to that Work or Derivative Works thereof, that is intentionally
submitted to Licensor for inclusion in the Work by the copyright owner
or by an individual or Legal Entity authorized to submit on behalf of
the copyright owner. For the purposes of this definition, "submitted"
means any form of electronic, verbal, or written communication sent
to the Licensor or its representatives, including but not limited to
communication on electronic mailing lists, source code control systems,
and issue tracking systems that are managed by, or on behalf of, the
Licensor for the purpose of discussing and improving the Work, but
excluding communication that is conspicuously marked or otherwise
designated in writing by the copyright owner as "Not a Contribution."
"Contributor" shall mean Licensor and any individual or Legal Entity
on behalf of whom a Contribution has been received by Licensor and
subsequently incorporated within the Work.
2. Grant of Copyright License. Subject to the terms and conditions of
this License, each Contributor hereby grants to You a perpetual,
worldwide, non-exclusive, no-charge, royalty-free, irrevocable
copyright license to reproduce, prepare Derivative Works of,
publicly display, publicly perform, sublicense, and distribute the
Work and such Derivative Works in Source or Object form.
3. Grant of Patent License. Subject to the terms and conditions of
this License, each Contributor hereby grants to You a perpetual,
worldwide, non-exclusive, no-charge, royalty-free, irrevocable
(except as stated in this section) patent license to make, have made,
use, offer to sell, sell, import, and otherwise transfer the Work,
where such license applies only to those patent claims licensable
by such Contributor that are necessarily infringed by their
Contribution(s) alone or by combination of their Contribution(s)
with the Work to which such Contribution(s) was submitted. If You
institute patent litigation against any entity (including a
cross-claim or counterclaim in a lawsuit) alleging that the Work
or a Contribution incorporated within the Work constitutes direct
or contributory patent infringement, then any patent licenses
granted to You under this License for that Work shall terminate
as of the date such litigation is filed.
4. Redistribution. You may reproduce and distribute copies of the
Work or Derivative Works thereof in any medium, with or without
modifications, and in Source or Object form, provided that You
meet the following conditions:
(a) You must give any other recipients of the Work or
Derivative Works a copy of this License; and
(b) You must cause any modified files to carry prominent notices
stating that You changed the files; and
(c) You must retain, in the Source form of any Derivative Works
that You distribute, all copyright, patent, trademark, and
attribution notices from the Source form of the Work,
excluding those notices that do not pertain to any part of
the Derivative Works; and
(d) If the Work includes a "NOTICE" text file as part of its
distribution, then any Derivative Works that You distribute must
include a readable copy of the attribution notices contained
within such NOTICE file, excluding those notices that do not
pertain to any part of the Derivative Works, in at least one
of the following places: within a NOTICE text file distributed
as part of the Derivative Works; within the Source form or
documentation, if provided along with the Derivative Works; or,
within a display generated by the Derivative Works, if and
wherever such third-party notices normally appear. The contents
of the NOTICE file are for informational purposes only and
do not modify the License. You may add Your own attribution
notices within Derivative Works that You distribute, alongside
or as an addendum to the NOTICE text from the Work, provided
that such additional attribution notices cannot be construed
as modifying the License.
You may add Your own copyright statement to Your modifications and
may provide additional or different license terms and conditions
for use, reproduction, or distribution of Your modifications, or
for any such Derivative Works as a whole, provided Your use,
reproduction, and distribution of the Work otherwise complies with
the conditions stated in this License.
5. Submission of Contributions. Unless You explicitly state otherwise,
any Contribution intentionally submitted for inclusion in the Work
by You to the Licensor shall be under the terms and conditions of
this License, without any additional terms or conditions.
Notwithstanding the above, nothing herein shall supersede or modify
the terms of any separate license agreement you may have executed
with Licensor regarding such Contributions.
6. Trademarks. This License does not grant permission to use the trade
names, trademarks, service marks, or product names of the Licensor,
except as required for reasonable and customary use in describing the
origin of the Work and reproducing the content of the NOTICE file.
7. Disclaimer of Warranty. Unless required by applicable law or
agreed to in writing, Licensor provides the Work (and each
Contributor provides its Contributions) on an "AS IS" BASIS,
WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or
implied, including, without limitation, any warranties or conditions
of TITLE, NON-INFRINGEMENT, MERCHANTABILITY, or FITNESS FOR A
PARTICULAR PURPOSE. You are solely responsible for determining the
appropriateness of using or redistributing the Work and assume any
risks associated with Your exercise of permissions under this License.
8. Limitation of Liability. In no event and under no legal theory,
whether in tort (including negligence), contract, or otherwise,
unless required by applicable law (such as deliberate and grossly
negligent acts) or agreed to in writing, shall any Contributor be
liable to You for damages, including any direct, indirect, special,
incidental, or consequential damages of any character arising as a
result of this License or out of the use or inability to use the
Work (including but not limited to damages for loss of goodwill,
work stoppage, computer failure or malfunction, or any and all
other commercial damages or losses), even if such Contributor
has been advised of the possibility of such damages.
9. Accepting Warranty or Additional Liability. While redistributing
the Work or Derivative Works thereof, You may choose to offer,
and charge a fee for, acceptance of support, warranty, indemnity,
or other liability obligations and/or rights consistent with this
License. However, in accepting such obligations, You may act only
on Your own behalf and on Your sole responsibility, not on behalf
of any other Contributor, and only if You agree to indemnify,
defend, and hold each Contributor harmless for any liability
incurred by, or claims asserted against, such Contributor by reason
of your accepting any such warranty or additional liability.
END OF TERMS AND CONDITIONS
APPENDIX: How to apply the Apache License to your work.
To apply the Apache License to your work, attach the following
boilerplate notice, with the fields enclosed by brackets "[]"
replaced with your own identifying information. (Don't include
the brackets!) The text should be enclosed in the appropriate
comment syntax for the file format. We also recommend that a
file or class name and description of purpose be included on the
same "printed page" as the copyright notice for easier
identification within third-party archives.
Copyright [2023] Lightning AI
Licensed under the Apache License, Version 2.0 (the "License");
you may not use this file except in compliance with the License.
You may obtain a copy of the License at
http://www.apache.org/licenses/LICENSE-2.0
Unless required by applicable law or agreed to in writing, software
distributed under the License is distributed on an "AS IS" BASIS,
WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
See the License for the specific language governing permissions and
limitations under the License.
README.md
+734
View File
@@ -0,0 +1,734 @@
<div align="center">
# ⚡ LitGPT
**20+ high-performance LLMs with recipes to pretrain, finetune, and deploy at scale.**
<pre>
✅ From scratch implementations ✅ No abstractions ✅ Beginner friendly
✅ Flash attention ✅ FSDP ✅ LoRA, QLoRA, Adapter
✅ Reduce GPU memory (fp4/8/16/32) ✅ 1-1000+ GPUs/TPUs ✅ 20+ LLMs
</pre>
---
![PyPI - Python Version](https://img.shields.io/pypi/pyversions/pytorch-lightning)
![cpu-tests](https://github.com/Lightning-AI/litgpt/actions/workflows/cpu-tests.yml/badge.svg) [![license](https://img.shields.io/badge/License-Apache%202.0-blue.svg)](https://github.com/Lightning-AI/litgpt/blob/main/LICENSE.md) [![Discord](https://img.shields.io/discord/1077906959069626439)](https://discord.gg/VptPCZkGNa)
<p align="center">
<a href="#quick-start">Quick start</a> •
<a href="#choose-from-20-llms">Models</a> •
<a href="#finetune-an-llm">Finetune</a> •
<a href="#deploy-an-llm">Deploy</a> •
<a href="#all-workflows">All workflows</a> •
<a href="#state-of-the-art-features">Features</a> •
<a href="#training-recipes">Recipes (YAML)</a> •
<a href="https://lightning.ai/">Lightning AI</a> •
<a href="#tutorials">Tutorials</a>
</p>
&nbsp;
<a target="_blank" href="https://lightning.ai/lightning-ai/studios/litgpt-quick-start">
<img src="https://pl-bolts-doc-images.s3.us-east-2.amazonaws.com/app-2/get-started-badge.svg" height="36px" alt="Get started"/>
</a>
&nbsp;
</div>
# Looking for GPUs?
Over 340,000 developers use [Lightning Cloud](https://lightning.ai/?utm_source=litgpt_readme&utm_medium=referral&utm_campaign=litgpt_readme) - purpose-built for PyTorch and PyTorch Lightning.
- [GPUs](https://lightning.ai/pricing?utm_source=litgpt_readme&utm_medium=referral&utm_campaign=litgpt_readme) from $0.19.
- [Clusters](https://lightning.ai/clusters?utm_source=litgpt_readme&utm_medium=referral&utm_campaign=litgpt_readme): frontier-grade training/inference clusters.
- [AI Studio (vibe train)](https://lightning.ai/studios?utm_source=litgpt_readme&utm_medium=referral&utm_campaign=litgpt_readme): workspaces where AI helps you debug, tune and vibe train.
- [AI Studio (vibe deploy)](https://lightning.ai/studios?utm_source=litgpt_readme&utm_medium=referral&utm_campaign=litgpt_readme): workspaces where AI helps you optimize, and deploy models.
- [Notebooks](https://lightning.ai/notebooks?utm_source=litgpt_readme&utm_medium=referral&utm_campaign=litgpt_readme): Persistent GPU workspaces where AI helps you code and analyze.
- [Inference](https://lightning.ai/deploy?utm_source=litgpt_readme&utm_medium=referral&utm_campaign=litgpt_readme): Deploy models as inference APIs.
# Finetune, pretrain, and inference LLMs Lightning fast ⚡⚡
Every LLM is implemented from scratch with **no abstractions** and **full control**, making them blazing fast, minimal, and performant at enterprise scale.
✅ **Enterprise ready -** Apache 2.0 for unlimited enterprise use.</br>
✅ **Developer friendly -** Easy debugging with no abstraction layers and single file implementations.</br>
✅ **Optimized performance -** Models designed to maximize performance, reduce costs, and speed up training.</br>
✅ **Proven recipes -** Highly-optimized training/finetuning recipes tested at enterprise scale.</br>
&nbsp;
# Quick start
Install LitGPT
```
pip install 'litgpt[extra]'
```
Load and use any of the [20+ LLMs](#choose-from-20-llms):
```python
from litgpt import LLM
llm = LLM.load("microsoft/phi-2")
text = llm.generate("Fix the spelling: Every fall, the family goes to the mountains.")
print(text)
# Corrected Sentence: Every fall, the family goes to the mountains.
```
&nbsp;
✅ Optimized for fast inference</br>
✅ Quantization</br>
✅ Runs on low-memory GPUs</br>
✅ No layers of internal abstractions</br>
✅ Optimized for production scale</br>
<details>
<summary>Advanced install options</summary>
Install from source:
```bash
git clone https://github.com/Lightning-AI/litgpt
cd litgpt
# if using uv
uv sync --all-extras
# if using pip
pip install -e ".[extra,compiler,test]"
```
</details>
[Explore the full Python API docs](tutorials/python-api.md).
&nbsp;
---
# Choose from 20+ LLMs
Every model is written from scratch to maximize performance and remove layers of abstraction:
| Model | Model size | Author | Reference |
|----|----|----|----|
| Llama 3, 3.1, 3.2, 3.3 | 1B, 3B, 8B, 70B, 405B | Meta AI | [Meta AI 2024](https://github.com/meta-llama/llama3) |
| Code Llama | 7B, 13B, 34B, 70B | Meta AI | [Rozière et al. 2023](https://arxiv.org/abs/2308.12950) |
| CodeGemma | 7B | Google | [Google Team, Google Deepmind](https://ai.google.dev/gemma/docs/codegemma) |
| Gemma 2 | 2B, 9B, 27B | Google | [Google Team, Google Deepmind](https://storage.googleapis.com/deepmind-media/gemma/gemma-2-report.pdf) |
| Phi 4 | 14B | Microsoft Research | [Abdin et al. 2024](https://arxiv.org/abs/2412.08905) |
| Qwen2.5 | 0.5B, 1.5B, 3B, 7B, 14B, 32B, 72B | Alibaba Group | [Qwen Team 2024](https://qwenlm.github.io/blog/qwen2.5/) |
| Qwen2.5 Coder | 0.5B, 1.5B, 3B, 7B, 14B, 32B | Alibaba Group | [Hui, Binyuan et al. 2024](https://arxiv.org/abs/2409.12186) |
| R1 Distill Llama | 8B, 70B | DeepSeek AI | [DeepSeek AI 2025](https://github.com/deepseek-ai/DeepSeek-R1/blob/main/DeepSeek_R1.pdf) |
| ... | ... | ... | ... |
<details>
<summary>See full list of 20+ LLMs</summary>
&nbsp;
#### All models
| Model | Model size | Author | Reference |
|----|----|----|----|
| CodeGemma | 7B | Google | [Google Team, Google Deepmind](https://ai.google.dev/gemma/docs/codegemma) |
| Code Llama | 7B, 13B, 34B, 70B | Meta AI | [Rozière et al. 2023](https://arxiv.org/abs/2308.12950) |
| Falcon | 7B, 40B, 180B | TII UAE | [TII 2023](https://falconllm.tii.ae) |
| Falcon 3 | 1B, 3B, 7B, 10B | TII UAE | [TII 2024](https://huggingface.co/blog/falcon3) |
| FreeWilly2 (Stable Beluga 2) | 70B | Stability AI | [Stability AI 2023](https://stability.ai/blog/stable-beluga-large-instruction-fine-tuned-models) |
| Function Calling Llama 2 | 7B | Trelis | [Trelis et al. 2023](https://huggingface.co/Trelis/Llama-2-7b-chat-hf-function-calling-v2) |
| Gemma | 2B, 7B | Google | [Google Team, Google Deepmind](https://storage.googleapis.com/deepmind-media/gemma/gemma-report.pdf) |
| Gemma 2 | 9B, 27B | Google | [Google Team, Google Deepmind](https://storage.googleapis.com/deepmind-media/gemma/gemma-2-report.pdf) |
| Gemma 3 | 1B, 4B, 12B, 27B | Google | [Google Team, Google Deepmind](https://arxiv.org/pdf/2503.19786) |
| Llama 2 | 7B, 13B, 70B | Meta AI | [Touvron et al. 2023](https://arxiv.org/abs/2307.09288) |
| Llama 3.1 | 8B, 70B | Meta AI | [Meta AI 2024](https://github.com/meta-llama/llama3) |
| Llama 3.2 | 1B, 3B | Meta AI | [Meta AI 2024](https://ai.meta.com/blog/llama-3-2-connect-2024-vision-edge-mobile-devices/) |
| Llama 3.3 | 70B | Meta AI | [Meta AI 2024](https://huggingface.co/meta-llama/Llama-3.3-70B-Instruct) |
| Mathstral | 7B | Mistral AI | [Mistral AI 2024](https://mistral.ai/news/mathstral/) |
| MicroLlama | 300M | Ken Wang | [MicroLlama repo](https://github.com/keeeeenw/MicroLlama) |
| Mixtral MoE | 8x7B | Mistral AI | [Mistral AI 2023](https://mistral.ai/news/mixtral-of-experts/) |
| Mistral | 7B, 123B | Mistral AI | [Mistral AI 2023](https://mistral.ai/news/announcing-mistral-7b/) |
| Mixtral MoE | 8x22B | Mistral AI | [Mistral AI 2024](https://mistral.ai/news/mixtral-8x22b/) |
| OLMo | 1B, 7B | Allen Institute for AI (AI2) | [Groeneveld et al. 2024](https://aclanthology.org/2024.acl-long.841/) |
| OpenLLaMA | 3B, 7B, 13B | OpenLM Research | [Geng & Liu 2023](https://github.com/openlm-research/open_llama) |
| Phi 1.5 & 2 | 1.3B, 2.7B | Microsoft Research | [Li et al. 2023](https://arxiv.org/abs/2309.05463) |
| Phi 3 | 3.8B | Microsoft Research | [Abdin et al. 2024](https://arxiv.org/abs/2404.14219) |
| Phi 4 | 14B | Microsoft Research | [Abdin et al. 2024](https://arxiv.org/abs/2412.08905) |
| Phi 4 Mini Instruct | 3.8B | Microsoft Research | [Microsoft 2025](https://arxiv.org/abs/2503.01743) |
| Phi 4 Mini Reasoning | 3.8B | Microsoft Research | [Xu, Peng et al. 2025](https://arxiv.org/abs/2504.21233) |
| Phi 4 Reasoning | 3.8B | Microsoft Research | [Abdin et al. 2025](https://arxiv.org/abs/2504.21318) |
| Phi 4 Reasoning Plus | 3.8B | Microsoft Research | [Abdin et al. 2025](https://arxiv.org/abs/2504.21318) |
| Platypus | 7B, 13B, 70B | Lee et al. | [Lee, Hunter, and Ruiz 2023](https://arxiv.org/abs/2308.07317) |
| Pythia | {14,31,70,160,410}M, {1,1.4,2.8,6.9,12}B | EleutherAI | [Biderman et al. 2023](https://arxiv.org/abs/2304.01373) |
| Qwen2.5 | 0.5B, 1.5B, 3B, 7B, 14B, 32B, 72B | Alibaba Group | [Qwen Team 2024](https://qwenlm.github.io/blog/qwen2.5/) |
| Qwen2.5 Coder | 0.5B, 1.5B, 3B, 7B, 14B, 32B | Alibaba Group | [Hui, Binyuan et al. 2024](https://arxiv.org/abs/2409.12186) |
| Qwen2.5 1M (Long Context) | 7B, 14B | Alibaba Group | [Qwen Team 2025](https://qwenlm.github.io/blog/qwen2.5-1m/) |
| Qwen2.5 Math | 1.5B, 7B, 72B | Alibaba Group | [An, Yang et al. 2024](https://arxiv.org/abs/2409.12122) |
| QwQ | 32B | Alibaba Group | [Qwen Team 2025](https://qwenlm.github.io/blog/qwq-32b/) |
| QwQ-Preview | 32B | Alibaba Group | [Qwen Team 2024](https://qwenlm.github.io/blog/qwq-32b-preview/) |
| Qwen3 | 0.6B, 1.7B, 4B{Hybrid, Thinking-2507, Instruct-2507}, 8B, 14B, 32B | Alibaba Group | [Qwen Team 2025](https://arxiv.org/abs/2505.09388/) |
| Qwen3 MoE | 30B{Hybrid, Thinking-2507, Instruct-2507}, 235B{Hybrid, Thinking-2507, Instruct-2507} | Alibaba Group | [Qwen Team 2025](https://arxiv.org/abs/2505.09388/) |
| R1 Distill Llama | 8B, 70B | DeepSeek AI | [DeepSeek AI 2025](https://github.com/deepseek-ai/DeepSeek-R1/blob/main/DeepSeek_R1.pdf) |
| SmolLM2 | 135M, 360M, 1.7B | Hugging Face | [Hugging Face 2024](https://github.com/huggingface/smollm) |
| Salamandra | 2B, 7B | Barcelona Supercomputing Centre | [BSC-LTC 2024](https://github.com/BSC-LTC/salamandra) |
| StableCode | 3B | Stability AI | [Stability AI 2023](https://stability.ai/blog/stablecode-llm-generative-ai-coding) |
| StableLM | 3B, 7B | Stability AI | [Stability AI 2023](https://github.com/Stability-AI/StableLM) |
| StableLM Zephyr | 3B | Stability AI | [Stability AI 2023](https://stability.ai/blog/stablecode-llm-generative-ai-coding) |
| TinyLlama | 1.1B | Zhang et al. | [Zhang et al. 2023](https://github.com/jzhang38/TinyLlama) |
**Tip**: You can list all available models by running the `litgpt download list` command.
</details>
&nbsp;
---
# Workflows
<p align="center">
<a href="#finetune-an-llm">Finetune</a> •
<a href="#pretrain-an-llm">Pretrain</a> •
<a href="#continue-pretraining-an-llm">Continued pretraining</a> •
<a href="#evaluate-an-llm">Evaluate</a> •
<a href="#deploy-an-llm">Deploy</a> •
<a href="#test-an-llm">Test</a>
</p>
&nbsp;
Use the command line interface to run advanced workflows such as pretraining or finetuning on your own data.
## All workflows
After installing LitGPT, select the model and workflow to run (finetune, pretrain, evaluate, deploy, etc...):
```bash
# litgpt [action] [model]
litgpt serve meta-llama/Llama-3.2-3B-Instruct
litgpt finetune meta-llama/Llama-3.2-3B-Instruct
litgpt pretrain meta-llama/Llama-3.2-3B-Instruct
litgpt chat meta-llama/Llama-3.2-3B-Instruct
litgpt evaluate meta-llama/Llama-3.2-3B-Instruct
```
&nbsp;
----
## Finetune an LLM
<div align="center">
<a target="_blank" href="https://lightning.ai/lightning-ai/studios/litgpt-finetune">
<img src="https://pl-bolts-doc-images.s3.us-east-2.amazonaws.com/app-2/run-on-studio.svg" height="36px" alt="Run on Studios"/>
</a>
</div>
&nbsp;
Finetuning is the process of taking a pretrained AI model and further training it on a smaller, specialized dataset tailored to a specific task or application.
&nbsp;
```bash
# 0) setup your dataset
curl -L https://huggingface.co/datasets/ksaw008/finance_alpaca/resolve/main/finance_alpaca.json -o my_custom_dataset.json
# 1) Finetune a model (auto downloads weights)
litgpt finetune microsoft/phi-2 \
--data JSON \
--data.json_path my_custom_dataset.json \
--data.val_split_fraction 0.1 \
--out_dir out/custom-model
# 2) Test the model
litgpt chat out/custom-model/final
# 3) Deploy the model
litgpt serve out/custom-model/final
```
[Read the full finetuning docs](tutorials/finetune.md)
&nbsp;
----
## Deploy an LLM
<div align="center">
<a target="_blank" href="https://lightning.ai/lightning-ai/studios/litgpt-serve">
<img src="https://pl-bolts-doc-images.s3.us-east-2.amazonaws.com/app-2/deploy-on-studios.svg" height="36px" alt="Deploy on Studios"/>
</a>
</div>
&nbsp;
Deploy a pretrained or finetune LLM to use it in real-world applications. Deploy, automatically sets up a web server that can be accessed by a website or app.
```bash
# deploy an out-of-the-box LLM
litgpt serve microsoft/phi-2
# deploy your own trained model
litgpt serve path/to/microsoft/phi-2/checkpoint
```
<details>
<summary>Show code to query server:</summary>
&nbsp;
Test the server in a separate terminal and integrate the model API into your AI product:
```python
# 3) Use the server (in a separate Python session)
import requests, json
response = requests.post(
"http://127.0.0.1:8000/predict",
json={"prompt": "Fix typos in the following sentence: Example input"}
)
print(response.json()["output"])
```
</details>
[Read the full deploy docs](tutorials/deploy.md).
&nbsp;
----
## Evaluate an LLM
Evaluate an LLM to test its performance on various tasks to see how well it understands and generates text. Simply put, we can evaluate things like how well would it do in college-level chemistry, coding, etc... (MMLU, Truthful QA, etc...)
```bash
litgpt evaluate microsoft/phi-2 --tasks 'truthfulqa_mc2,mmlu'
```
[Read the full evaluation docs](tutorials/evaluation.md).
&nbsp;
----
## Test an LLM
<div align="center">
<a target="_blank" href="https://lightning.ai/lightning-ai/studios/litgpt-chat">
<img src="https://pl-bolts-doc-images.s3.us-east-2.amazonaws.com/app-2/run-on-studio.svg" height="36px" alt="Run on Studios"/>
</a>
</div>
&nbsp;
Test how well the model works via an interactive chat. Use the `chat` command to chat, extract embeddings, etc...
Here's an example showing how to use the Phi-2 LLM:
```bash
litgpt chat microsoft/phi-2
>> Prompt: What do Llamas eat?
```
<details>
<summary>Full code:</summary>
&nbsp;
```bash
# 1) List all supported LLMs
litgpt download list
# 2) Use a model (auto downloads weights)
litgpt chat microsoft/phi-2
>> Prompt: What do Llamas eat?
```
The download of certain models requires an additional access token. You can read more about this in the [download](tutorials/download_model_weights.md#specific-models-and-access-tokens) documentation.
</details>
[Read the full chat docs](tutorials/inference.md).
&nbsp;
----
## Pretrain an LLM
<div align="center">
<a target="_blank" href="https://lightning.ai/lightning-ai/studios/litgpt-pretrain">
<img src="https://pl-bolts-doc-images.s3.us-east-2.amazonaws.com/app-2/run-on-studio.svg" height="36px" alt="Run on Studios"/>
</a>
</div>
&nbsp;
Pretraining is the process of teaching an AI model by exposing it to a large amount of data before it is fine-tuned for specific tasks.
<details>
<summary>Show code:</summary>
&nbsp;
```bash
mkdir -p custom_texts
curl https://www.gutenberg.org/cache/epub/24440/pg24440.txt --output custom_texts/book1.txt
curl https://www.gutenberg.org/cache/epub/26393/pg26393.txt --output custom_texts/book2.txt
# 1) Download a tokenizer
litgpt download EleutherAI/pythia-160m \
--tokenizer_only True
# 2) Pretrain the model
litgpt pretrain EleutherAI/pythia-160m \
--tokenizer_dir EleutherAI/pythia-160m \
--data TextFiles \
--data.train_data_path "custom_texts/" \
--train.max_tokens 10_000_000 \
--out_dir out/custom-model
# 3) Test the model
litgpt chat out/custom-model/final
```
</details>
[Read the full pretraining docs](tutorials/pretrain.md)
&nbsp;
----
## Continue pretraining an LLM
<div align="center">
<a target="_blank" href="https://lightning.ai/lightning-ai/studios/litgpt-continue-pretraining">
<img src="https://pl-bolts-doc-images.s3.us-east-2.amazonaws.com/app-2/run-on-studio.svg" height="36px" alt="Run on Studios"/>
</a>
</div>
&nbsp;
Continued pretraining is another way of finetuning that specializes an already pretrained model by training on custom data:
<details>
<summary>Show code:</summary>
&nbsp;
```bash
mkdir -p custom_texts
curl https://www.gutenberg.org/cache/epub/24440/pg24440.txt --output custom_texts/book1.txt
curl https://www.gutenberg.org/cache/epub/26393/pg26393.txt --output custom_texts/book2.txt
# 1) Continue pretraining a model (auto downloads weights)
litgpt pretrain EleutherAI/pythia-160m \
--tokenizer_dir EleutherAI/pythia-160m \
--initial_checkpoint_dir EleutherAI/pythia-160m \
--data TextFiles \
--data.train_data_path "custom_texts/" \
--train.max_tokens 10_000_000 \
--out_dir out/custom-model
# 2) Test the model
litgpt chat out/custom-model/final
```
</details>
[Read the full continued pretraining docs](tutorials/pretrain.md#continued-pretraining-on-custom-data)
&nbsp;
----
# State-of-the-art features
✅ State-of-the-art optimizations: Flash Attention v2, multi-GPU support via fully-sharded data parallelism, [optional CPU offloading](tutorials/oom.md#do-sharding-across-multiple-gpus), and [TPU and XLA support](extensions/xla).</br>
✅ [Pretrain](tutorials/pretrain.md), [finetune](tutorials/finetune.md), and [deploy](tutorials/inference.md)</br>
✅ Reduce compute requirements with low-precision settings: FP16, BF16, and FP16/FP32 mixed.</br>
✅ Lower memory requirements with [quantization](tutorials/quantize.md): 4-bit floats, 8-bit integers, and double quantization.</br>
✅ [Configuration files](config_hub) for great out-of-the-box performance.</br>
✅ Parameter-efficient finetuning: [LoRA](tutorials/finetune_lora.md), [QLoRA](tutorials/finetune_lora.md), [Adapter](tutorials/finetune_adapter.md), and [Adapter v2](tutorials/finetune_adapter.md).</br>
✅ [Exporting](tutorials/convert_lit_models.md) to other popular model weight formats.</br>
✅ Many popular datasets for [pretraining](tutorials/pretrain.md) and [finetuning](tutorials/prepare_dataset.md), and [support for custom datasets](tutorials/prepare_dataset.md#preparing-custom-datasets-for-instruction-finetuning).</br>
✅ Readable and easy-to-modify code to experiment with the latest research ideas.</br>
&nbsp;
---
# Training recipes
LitGPT comes with validated recipes (YAML configs) to train models under different conditions. We've generated these recipes based on the parameters we found to perform the best for different training conditions.
Browse all training recipes [here](config_hub).
### Example
```bash
litgpt finetune \
--config https://raw.githubusercontent.com/Lightning-AI/litgpt/main/config_hub/finetune/llama-2-7b/lora.yaml
```
<details>
<summary>✅ Use configs to customize training</summary>
Configs let you customize training for all granular parameters like:
```yaml
# The path to the base model's checkpoint directory to load for finetuning. (type: <class 'Path'>, default: checkpoints/stabilityai/stablelm-base-alpha-3b)
checkpoint_dir: checkpoints/meta-llama/Llama-2-7b-hf
# Directory in which to save checkpoints and logs. (type: <class 'Path'>, default: out/lora)
out_dir: out/finetune/qlora-llama2-7b
# The precision to use for finetuning. Possible choices: "bf16-true", "bf16-mixed", "32-true". (type: Optional[str], default: null)
precision: bf16-true
...
```
</details>
<details>
<summary>✅ Example: LoRA finetuning config</summary>
&nbsp;
```yaml
# The path to the base model's checkpoint directory to load for finetuning. (type: <class 'Path'>, default: checkpoints/stabilityai/stablelm-base-alpha-3b)
checkpoint_dir: checkpoints/meta-llama/Llama-2-7b-hf
# Directory in which to save checkpoints and logs. (type: <class 'Path'>, default: out/lora)
out_dir: out/finetune/qlora-llama2-7b
# The precision to use for finetuning. Possible choices: "bf16-true", "bf16-mixed", "32-true". (type: Optional[str], default: null)
precision: bf16-true
# If set, quantize the model with this algorithm. See ``tutorials/quantize.md`` for more information. (type: Optional[Literal['nf4', 'nf4-dq', 'fp4', 'fp4-dq', 'int8-training']], default: null)
quantize: bnb.nf4
# How many devices/GPUs to use. (type: Union[int, str], default: 1)
devices: 1
# How many nodes to use. (type: int, default: 1)
num_nodes: 1
# The LoRA rank. (type: int, default: 8)
lora_r: 32
# The LoRA alpha. (type: int, default: 16)
lora_alpha: 16
# The LoRA dropout value. (type: float, default: 0.05)
lora_dropout: 0.05
# Whether to apply LoRA to the query weights in attention. (type: bool, default: True)
lora_query: true
# Whether to apply LoRA to the key weights in attention. (type: bool, default: False)
lora_key: false
# Whether to apply LoRA to the value weights in attention. (type: bool, default: True)
lora_value: true
# Whether to apply LoRA to the output projection in the attention block. (type: bool, default: False)
lora_projection: false
# Whether to apply LoRA to the weights of the MLP in the attention block. (type: bool, default: False)
lora_mlp: false
# Whether to apply LoRA to output head in GPT. (type: bool, default: False)
lora_head: false
# Data-related arguments. If not provided, the default is ``litgpt.data.Alpaca``.
data:
class_path: litgpt.data.Alpaca2k
init_args:
mask_prompt: false
val_split_fraction: 0.05
prompt_style: alpaca
ignore_index: -100
seed: 42
num_workers: 4
download_dir: data/alpaca2k
# Training-related arguments. See ``litgpt.args.TrainArgs`` for details
train:
# Number of optimizer steps between saving checkpoints (type: Optional[int], default: 1000)
save_interval: 200
# Number of iterations between logging calls (type: int, default: 1)
log_interval: 1
# Number of samples between optimizer steps across data-parallel ranks (type: int, default: 128)
global_batch_size: 8
# Number of samples per data-parallel rank (type: int, default: 4)
micro_batch_size: 2
# Number of iterations with learning rate warmup active (type: int, default: 100)
lr_warmup_steps: 10
# Number of epochs to train on (type: Optional[int], default: 5)
epochs: 4
# Total number of tokens to train on (type: Optional[int], default: null)
max_tokens:
# Limits the number of optimizer steps to run (type: Optional[int], default: null)
max_steps:
# Limits the length of samples (type: Optional[int], default: null)
max_seq_length: 512
# Whether to tie the embedding weights with the language modeling head weights (type: Optional[bool], default: null)
tie_embeddings:
# (type: float, default: 0.0003)
learning_rate: 0.0002
# (type: float, default: 0.02)
weight_decay: 0.0
# (type: float, default: 0.9)
beta1: 0.9
# (type: float, default: 0.95)
beta2: 0.95
# (type: Optional[float], default: null)
max_norm:
# (type: float, default: 6e-05)
min_lr: 6.0e-05
# Evaluation-related arguments. See ``litgpt.args.EvalArgs`` for details
eval:
# Number of optimizer steps between evaluation calls (type: int, default: 100)
interval: 100
# Number of tokens to generate (type: Optional[int], default: 100)
max_new_tokens: 100
# Number of iterations (type: int, default: 100)
max_iters: 100
# The name of the logger to send metrics to. (type: Literal['wandb', 'tensorboard', 'csv'], default: csv)
logger_name: csv
# The random seed to use for reproducibility. (type: int, default: 1337)
seed: 1337
```
</details>
<details>
<summary>✅ Override any parameter in the CLI:</summary>
```bash
litgpt finetune \
--config https://raw.githubusercontent.com/Lightning-AI/litgpt/main/config_hub/finetune/llama-2-7b/lora.yaml \
--lora_r 4
```
</details>
&nbsp;
----
# Project highlights
LitGPT powers many great AI projects, initiatives, challenges and of course enterprises. Please submit a pull request to be considered for a feature.
<details>
<summary>📊 SAMBA: Simple Hybrid State Space Models for Efficient Unlimited Context Language Modeling</summary>
The [Samba](https://github.com/microsoft/Samba) project by researchers at Microsoft is built on top of the LitGPT code base and combines state space models with sliding window attention, which outperforms pure state space models.
</details>
<details>
<summary>🏆 NeurIPS 2023 Large Language Model Efficiency Challenge: 1 LLM + 1 GPU + 1 Day</summary>
The LitGPT repository was the official starter kit for the [NeurIPS 2023 LLM Efficiency Challenge](https://llm-efficiency-challenge.github.io), which is a competition focused on finetuning an existing non-instruction tuned LLM for 24 hours on a single GPU.
</details>
<details>
<summary>🦙 TinyLlama: An Open-Source Small Language Model</summary>
LitGPT powered the [TinyLlama project](https://github.com/jzhang38/TinyLlama) and [TinyLlama: An Open-Source Small Language Model](https://arxiv.org/abs/2401.02385) research paper.
</details>
<details>
<summary>🍪 MicroLlama: MicroLlama-300M</summary>
[MicroLlama](https://github.com/keeeeenw/MicroLlama) is a 300M Llama model pretrained on 50B tokens powered by TinyLlama and LitGPT.
</details>
<details>
<summary>🔬 Pre-training Small Base LMs with Fewer Tokens</summary>
The research paper ["Pre-training Small Base LMs with Fewer Tokens"](https://arxiv.org/abs/2404.08634), which utilizes LitGPT, develops smaller base language models by inheriting a few transformer blocks from larger models and training on a tiny fraction of the data used by the larger models. It demonstrates that these smaller models can perform comparably to larger models despite using significantly less training data and resources.
</details>
&nbsp;
----
# Community
We welcome all individual contributors, regardless of their level of experience or hardware. Your contributions are valuable, and we are excited to see what you can accomplish in this collaborative and supportive environment.
- [Request a feature](https://github.com/Lightning-AI/litgpt/issues)
- [Submit your first contribution](https://lightning.ai/pages/community/tutorial/how-to-contribute-to-litgpt/)
- [Join our Discord](https://discord.gg/VptPCZkGNa)
&nbsp;
# Tutorials
🚀 [Get started](tutorials/0_to_litgpt.md)</br>
⚡️ [Finetuning, incl. LoRA, QLoRA, and Adapters](tutorials/finetune.md)</br>
🤖 [Pretraining](tutorials/pretrain.md)</br>
💬 [Model evaluation](tutorials/evaluation.md)</br>
📘 [Supported and custom datasets](tutorials/prepare_dataset.md)</br>
🧹 [Quantization](tutorials/quantize.md)</br>
🤯 [Tips for dealing with out-of-memory (OOM) errors](tutorials/oom.md)</br>
🧑🏽‍💻 [Using cloud TPUs](extensions/xla)</br>
&nbsp;
----
### Acknowledgments
This implementation extends on [Lit-LLaMA](https://github.com/lightning-AI/lit-llama) and [nanoGPT](https://github.com/karpathy/nanoGPT), and it's **powered by [Lightning Fabric](https://lightning.ai/docs/fabric/stable/) ⚡**.
- [@karpathy](https://github.com/karpathy) for [nanoGPT](https://github.com/karpathy/nanoGPT)
- [@EleutherAI](https://github.com/EleutherAI) for [GPT-NeoX](https://github.com/EleutherAI/gpt-neox) and the [Evaluation Harness](https://github.com/EleutherAI/lm-evaluation-harness)
- [@TimDettmers](https://github.com/TimDettmers) for [bitsandbytes](https://github.com/TimDettmers/bitsandbytes)
- [@Microsoft](https://github.com/microsoft) for [LoRA](https://github.com/microsoft/LoRA)
- [@tridao](https://github.com/tridao) for [Flash Attention 2](https://github.com/Dao-AILab/flash-attention)
### License
LitGPT is released under the [Apache 2.0](https://github.com/Lightning-AI/litgpt/blob/main/LICENSE.md) license.
### Citation
If you use LitGPT in your research, please cite the following work:
```bibtex
@misc{litgpt-2023,
author = {Lightning AI},
title = {LitGPT},
howpublished = {\url{https://github.com/Lightning-AI/litgpt}},
year = {2023},
}
```
&nbsp;
README.wehub.md
+7
View File
@@ -0,0 +1,7 @@
# WeHub 来源说明
- 原始项目：`Lightning-AI/litgpt`
- 原始仓库：https://github.com/Lightning-AI/litgpt
- 导入方式：上游默认分支的最新快照
- 原作者、版权和许可证信息以原始仓库及本仓库 LICENSE 为准
- 本文件仅用于记录来源，不代表 WeHub 是原项目作者
config_hub/finetune/README.md
+119
View File
@@ -0,0 +1,119 @@
## Config files
The table below lists the performances you can expect from the provided config files. Note that you can achieve lower memory consumption by lowering the micro batch size as needed. In addition, you can lower the rank (`lora_r`) in the LoRA configuration files and disable LoRA for certain layers (for example, setting `lora_projection` and other LoRA layer-specific parameters to `false`).
For more information, see the [Dealing with out-of-memory (OOM) errors](../../tutorials/oom.md) on lowering the memory requirements.
The "Cost" column refers to the on-demand compute cost on [Lightning AI Studios where these benchmarks were executed](https://lightning.ai/lightning-ai/studios/automated-benchmarks-for-litgpt).
All experiments were conducted using bfloat-16 precision on the Alpaca2k dataset. The "Multitask score" refers to [MMLU](https://arxiv.org/abs/2009.03300).
&nbsp;
| Config | Model | Epochs | Max seq length | Micro batch size | Machine | Training runtime | Cost | Peak memory | Validation loss | Validation perplexity | Multitask score (MMLU) |
| --------------------------------- | ---------------------- | ------ | -------------- | ---------------- | ------- | ---------------- | ---- | ----------- | --------------- | --------------------- | --------------- |
| falcon-7b/lora.yaml | falcon-7b | 4 | 512 | 1 | 1xA10G | 24.84 min | $0.7 | 16.69 GB | 0.945 | 2.573 | 26.2% |
| falcon-7b/lora.yaml | falcon-7b | 4 | 512 | 1 | 4xA10G | 24.94 min | $2.0 | 16.69 GB | 0.945 | 2.573 | 26.4% |
| falcon-7b/qlora.yaml | falcon-7b | 4 | 512 | 1 | 1xA10G | 50.85 min | $1.5 | 9.44 GB | 0.993 | 2.699 | 26.3% |
| | | | | | | | | | | | |
| gemma-2b/full.yaml | gemma-2b | 1 | 512 | 1 | 4xA10G | 14.06 min | $1.1 | 17.43 GB | 1.021 | 2.777 | 32.4% |
| gemma-2b/lora.yaml | gemma-2b | 2 | 512 | 2 | 1xA10G | 9.41 min | $0.3 | 12.62 GB | 0.981 | 2.666 | 34.4% |
| gemma-2b/lora.yaml | gemma-2b | 2 | 512 | 2 | 4xA10G | 9.41 min | $0.8 | 12.62 GB | 0.981 | 2.667 | 34.0% |
| gemma-2b/qlora.yaml | gemma-2b | 2 | 512 | 2 | 1xA10G | 12.91 min | $0.4 | 11.58 GB | 1.085 | 2.959 | 36.4% |
| | | | | | | | | | | | |
| gemma-7b/lora.yaml | gemma-7b | 2 | 512 | 1 | 1xA10G | OOM | OOM | OOM | OOM | OOM | |
| gemma-7b/lora.yaml | gemma-7b | 2 | 512 | 1 | 4xA10G | OOM | OOM | OOM | OOM | OOM | |
| gemma-7b/qlora.yaml | gemma-7b | 2 | 512 | 1 | 1xA10G | 43.58 min | $1.3 | 17.18 GB | 0.973 | 2.646 | 62.45% |
| | | | | | | | | | | | |
| gemma2-2b/lora.yaml | gemma-2b | 2 | 512 | 2 | 1xA10G | 11.96 min | $0.4 | 14.31 GB | 0.951 | 2.589 | 23.84% |
| gemma2b/qlora.yaml | gemma-2b | 2 | 512 | 2 | 1xA10G | 16.06 min | $0.5 | 13.52 GB | 0.983 | 2.673 | 24.12% |
| | | | | | | | | | | | |
| gemma2-9b/lora.yaml | gemma-2-9b | 2 | 512 | 1 | 1xA10G | OOM | OOM | OOM | OOM | OOM | |
| gemma2-9b/lora.yaml | gemma-2-9b | 2 | 512 | 1 | 4xA10G | OOM | OOM | OOM | OOM | OOM | |
| gemma2-9b/qlora.yaml | gemma-2-9b | 2 | 512 | 1 | 1xA10G | 50.01 min | $4.0 | 20.92 GB | 0.852 | 2.345 | 24.2% |
| | | | | | | | | | | | |
| llama-2-7b/full.yaml | llama-2-7b | 1 | 512 | 4 | 4xA10G | OOM | OOM | OOM | OOM | OOM | |
| llama-2-7b/lora.yaml | llama-2-7b | 4 | 512 | 2 | 1xA10G | 32.82 min | $1.0 | 19.77 GB | 0.802 | 2.230 | 40.3% |
| llama-2-7b/lora.yaml | llama-2-7b | 4 | 512 | 2 | 4xA10G | 32.83 min | $2.6 | 19.77 GB | 0.802 | 2.229 | 40.2% |
| llama-2-7b/qlora.yaml | llama-2-7b | 4 | 512 | 2 | 1xA10G | 45.67 min | $1.4 | 13.68 GB | 0.814 | 2.258 | 38.6% |
| | | | | | | | | | | | |
| llama-3-8b/full.yaml | llama-3-8b | 1 | 512 | 4 | 4xA10G | OOM | OOM | OOM | OOM | OOM | |
| llama-3-8b/lora.yaml | llama-3-8b | 2 | 512 | 1 | 1xA10G | 14.79 min | $0.4 | 19.73 GB | 0.888 | 2.431 | 62.4% |
| llama-3-8b/lora.yaml | llama-3-8b | 2 | 512 | 1 | 4xA10G | 14.88 min | $1.2 | 19.73 GB | 0.889 | 2.432 | 62.5% |
| llama-3-8b/qlora.yaml | llama-3-8b | 2 | 512 | 2 | 1xA10G | 22.24 min | $0.7 | 17.41 GB | 0.939 | 2.558 | 62.2% |
| | | | | | | | | | | | |
| llama-3.1-8b/full.yaml | llama-3.1-8b | 1 | 512 | 4 | 1xA10G | OOM | OOM | OOM | OOM | OOM | OOM |
| llama-3.1-8b/lora.yaml | llama-3.1-8b | 2 | 512 | 1 | 1xA10G | 13.36 min | $1.1 | 19.73 GB | 0.878 | 2.406 | xx.xx |
| llama-3.1-8b/qlora.yaml | llama-3.1-8b | 2 | 512 | 2 | 1xA10G | 21.81 min | $0.7 | 17.41 GB | 0.928 | 2.529 | xx.xx |
| | | | | | | | | | | | |
| llama-3.2-1b/full.yaml | llama-3.2-1b | 1 | 512 | 4 | 1xA10G | 2.01 min | $0.1 | 8.70 GB | 1.442 | 4.229 | 38.21% |
| llama-3.2-1b/lora.yaml | llama-3.2-1b | 2 | 512 | 1 | 1xA10G | 4.17 min | $0.4 | 4.49 GB | 1.114 | 3.046 | 36.87% |
| llama-3.2-1b/qlora.yaml | llama-3.2-1b | 2 | 512 | 2 | 1xA10G | 6.20 min | $0.6 | 5.53 GB | 1.201 | 3.322 | 36.49% |
| | | | | | | | | | | | |
| llama-3.2-3b/full.yaml | llama-3.2-3b | 1 | 512 | 4 | 1xA10G | 4.71 min | $0.4 | 16.51 GB | 1.255 | 3.509 | 54.69% |
| llama-3.2-3b/lora.yaml | llama-3.2-3b | 2 | 512 | 1 | 1xA10G | 8.31 min | $0.8 | 9.67 GB | 0.973 | 2.647 | 54.77% |
| llama-3.2-3b/qlora.yaml | llama-3.2-3b | 2 | 512 | 2 | 1xA10G | 14.89 min | $1.4 | 10.30 GB | 1.031 | 2.804 | 55.08% |
| | | | | | | | | | | | |
| mistral-7b-v0.2/lora.yaml | mistral-7b-v0.2 | 4 | 512 | 2 | 1xA10G | 31.00 min | $0.9 | 20.66 GB | 0.801 | 2.228 | 55.7% |
| mistral-7b-v0.2/lora.yaml | mistral-7b-v0.2 | 4 | 512 | 2 | 4xA10G | 31.00 min | $2.5 | 20.66 GB | 0.802 | 2.229 | 55.5% |
| mistral-7b-v0.2/qlora.yaml | mistral-7b-v0.2 | 4 | 512 | 2 | 1xA10G | 44.75 min | $1.3 | 14.29 GB | 0.813 | 2.255 | 56.5% |
| | | | | | | | | | | | |
| mistral-7b/lora.yaml | mistral-7b | 4 | 512 | 2 | 1xA10G | 31.01 min | $0.9 | 20.66 GB | 0.794 | 2.211 | 57.9% |
| mistral-7b/lora.yaml | mistral-7b | 4 | 512 | 2 | 4xA10G | 31.03 min | $2.5 | 20.66 GB | 0.796 | 2.218 | 57.9% |
| mistral-7b/qlora.yaml | mistral-7b | 4 | 512 | 2 | 1xA10G | 44.75 min | $1.3 | 14.29 GB | 0.803 | 2.231 | 57.9% |
| | | | | | | | | | | | |
| phi-2/full.yaml | phi-2 | 1 | 512 | 4 | 4xA10G | 11.87 min | $1.0 | 14.44 GB | 1.305 | 3.688 | 38.4% |
| phi-2/lora.yaml | phi-2 | 1 | 512 | 4 | 1xA10G | 3.78 min | $0.1 | 13.98 GB | 0.819 | 2.269 | 53.0% |
| phi-2/lora.yaml | phi-2 | 1 | 512 | 4 | 4xA10G | 3.78 min | $0.3 | 13.98 GB | 0.820 | 2.271 | 52.4% |
| phi-2/qlora.yaml | phi-2 | 1 | 512 | 4 | 1xA10G | 4.51 min | $0.1 | 14.27 GB | 0.837 | 2.310 | 52.3% |
| | | | | | | | | | | | |
| phi-3/full.yaml | Phi-3-mini-4k-instruct | 1 | 512 | 4 | 1xA10G | 6.93 min | $0.2 | 17.01 GB | 0.714 | 2.043 | 69.81% |
| phi-3/lora.yaml | Phi-3-mini-4k-instruct | 1 | 512 | 4 | 1xA10G | 6.46 min | $0.2 | 19.75 GB | 0.707 | 2.028 | 69.70% |
| phi-3/qlora.yaml | Phi-3-mini-4k-instruct | 1 | 512 | 4 | 1xA10G | 7.47 min | $0.2 | 19.13 GB | 0.729 | 2.074 | 68.96% |
| | | | | | | | | | | | |
| stablelm-base-alpha-3b/full.yaml | stablelm-base-alpha-3b | 1 | 512 | 1 | 4xA10G | 70.13 min | $5.6 | 21.23 GB | 1.513 | 4.540 | 23.2% |
| stablelm-base-alpha-3b/lora.yaml | stablelm-base-alpha-3b | 4 | 512 | 1 | 1xA10G | 13.07 min | $0.4 | 8.58 GB | 1.361 | 3.900 | 25.9% |
| stablelm-base-alpha-3b/lora.yaml | stablelm-base-alpha-3b | 4 | 512 | 1 | 4xA10G | 13.16 min | $1.1 | 8.58 GB | 1.362 | 3.906 | 25.9% |
| stablelm-base-alpha-3b/qlora.yaml | stablelm-base-alpha-3b | 4 | 512 | 1 | 1xA10G | 25.86 min | $0.8 | 5.24 GB | 1.388 | 4.009 | 26.1% |
| | | | | | | | | | | | |
| tiny-llama/full.yaml | tiny-llama | 1 | 512 | 4 | 1xA10G | 2.58 min | $0.1 | 14.10 GB | 1.088 | 2.968 | 24.6% |
| tiny-llama/full.yaml | tiny-llama | 1 | 512 | 4 | 4xA10G | 2.57 min | $0.2 | 14.10 GB | 1.088 | 2.968 | 24.5% |
| tiny-llama/lora.yaml | tiny-llama | 3 | 512 | 8 | 1xA10G | 8.09 min | $0.2 | 13.50 GB | 1.039 | 2.826 | 25.5% |
| tiny-llama/qlora.yaml | tiny-llama | 3 | 512 | 8 | 1xA10G | 8.70 min | $0.3 | 16.24 GB | 1.056 | 2.874 | 25.3% |
*OOM = Out of memory
&nbsp;
## Extending the context length
If you require a longer sequence length than the one used in a given config file, you can either edit the `max_seq_length` in the config file or pass an additional argument when running the finetuning command, for example, `--max_seq_length 4096` to override the sequence length provided in the config file.
&nbsp;
## Training on GPUs without bfloat16 support
If you are training on GPUs without bfloat-16 support, you need to change the `precision` option to `16-true` (16-bit floating point precision) or `16-mixed` (16/32-bit mixed precision) training:
```bash
litgpt finetune lora \
--config config_hub/finetune/phi-2/lora.yaml \
--precision 16-true
```
or
```bash
litgpt finetune lora \
--config config_hub/finetune/phi-2/lora.yaml \
--precision 16-mixed
```
Note that `16-true` is more compute and memory-efficient, but it can sometimes lead to training convergence issues. In this case, it's recommended to use `16-mixed`.
&nbsp;
## Multi-GPU experiments
All runs are single-GPU experiments, use `--devices 4` to utilize more than one GPU:
```bash
litgpt finetune lora \
--config config_hub/finetune/phi-2/lora.yaml \
--devices 4
```
config_hub/finetune/falcon-7b/lora.yaml
+131
View File
@@ -0,0 +1,131 @@
# The path to the base model's checkpoint directory to load for finetuning. (type: <class 'Path'>, default: checkpoints/stabilityai/stablelm-base-alpha-3b)
checkpoint_dir: checkpoints/tiiuae/falcon-7b
# Directory in which to save checkpoints and logs. (type: <class 'Path'>, default: out/lora)
out_dir: out/finetune/lora-falcon-7b
# The precision to use for finetuning. Possible choices: "bf16-true", "bf16-mixed", "32-true". (type: Optional[str], default: null)
precision: bf16-true
# If set, quantize the model with this algorithm. See ``tutorials/quantize.md`` for more information. (type: Optional[Literal['nf4', 'nf4-dq', 'fp4', 'fp4-dq', 'int8-training']], default: null)
quantize:
# How many devices/GPUs to use. (type: Union[int, str], default: 1)
devices: 1
# How many nodes to use. (type: int, default: 1)
num_nodes: 1
# The LoRA rank. (type: int, default: 8)
lora_r: 32
# The LoRA alpha. (type: int, default: 16)
lora_alpha: 16
# The LoRA dropout value. (type: float, default: 0.05)
lora_dropout: 0.05
# Whether to apply LoRA to the query weights in attention. (type: bool, default: True)
lora_query: true
# Whether to apply LoRA to the key weights in attention. (type: bool, default: False)
lora_key: false
# Whether to apply LoRA to the value weights in attention. (type: bool, default: True)
lora_value: true
# Whether to apply LoRA to the output projection in the attention block. (type: bool, default: False)
lora_projection: false
# Whether to apply LoRA to the weights of the MLP in the attention block. (type: bool, default: False)
lora_mlp: false
# Whether to apply LoRA to output head in GPT. (type: bool, default: False)
lora_head: false
# Data-related arguments. If not provided, the default is ``litgpt.data.Alpaca``.
data:
class_path: litgpt.data.Alpaca2k
init_args:
mask_prompt: false
prompt_style: alpaca
ignore_index: -100
seed: 42
num_workers: 4
# Training-related arguments. See ``litgpt.args.TrainArgs`` for details
train:
# Number of optimizer steps between saving checkpoints (type: Optional[int], default: 1000)
save_interval: 200
# Number of iterations between logging calls (type: int, default: 1)
log_interval: 1
# Number of samples between optimizer steps across data-parallel ranks (type: int, default: 128)
global_batch_size: 8
# Number of samples per data-parallel rank (type: int, default: 4)
micro_batch_size: 1
# Number of iterations with learning rate warmup active (type: int, default: 100)
lr_warmup_steps: 10
# Number of epochs to train on (type: Optional[int], default: 5)
epochs: 4
# Total number of tokens to train on (type: Optional[int], default: null)
max_tokens:
# Limits the number of optimizer steps to run. (type: Optional[int], default: null)
max_steps:
# Limits the length of samples. Off by default (type: Optional[int], default: null)
max_seq_length: 512
# Whether to tie the embedding weights with the language modeling head weights. (type: Optional[bool], default: null)
tie_embeddings:
# (type: Optional[float], default: null)
max_norm:
# (type: float, default: 6e-05)
min_lr: 6.0e-05
# Evaluation-related arguments. See ``litgpt.args.EvalArgs`` for details
eval:
# Number of optimizer steps between evaluation calls (type: int, default: 100)
interval: 100
# Number of tokens to generate (type: Optional[int], default: 100)
max_new_tokens: 100
# Number of iterations (type: int, default: 100)
max_iters: 100
# Whether to evaluate on the validation set at the beginning of the training
initial_validation: false
# Whether to evaluate on the validation set at the end the training
final_validation: true
# The name of the logger to send metrics to. (type: LoggerChoice, i.e. Literal['wandb', 'tensorboard', 'csv', 'mlflow', 'litlogger'], default: csv)
logger_name: csv
# The random seed to use for reproducibility. (type: int, default: 1337)
seed: 1337
# Optimizer-related arguments
optimizer:
class_path: torch.optim.AdamW
init_args:
# (type: float, default: 0.001)
lr: 0.0002
# (type: float, default: 0.01)
weight_decay: 0.0
# (type: tuple, default: (0.9,0.999))
betas:
- 0.9
- 0.95
config_hub/finetune/falcon-7b/qlora.yaml
+133
View File
@@ -0,0 +1,133 @@
# The path to the base model's checkpoint directory to load for finetuning. (type: <class 'Path'>, default: checkpoints/stabilityai/stablelm-base-alpha-3b)
checkpoint_dir: checkpoints/tiiuae/falcon-7b
# Directory in which to save checkpoints and logs. (type: <class 'Path'>, default: out/lora)
out_dir: out/finetune/qlora-falcon-7b
# The precision to use for finetuning. Possible choices: "bf16-true", "bf16-mixed", "32-true". (type: Optional[str], default: null)
precision: bf16-true
# If set, quantize the model with this algorithm. See ``tutorials/quantize.md`` for more information. (type: Optional[Literal['nf4', 'nf4-dq', 'fp4', 'fp4-dq', 'int8-training']], default: null)
quantize: bnb.nf4
# How many devices/GPUs to use. (type: Union[int, str], default: 1)
devices: 1
# How many nodes to use. (type: int, default: 1)
num_nodes: 1
# The LoRA rank. (type: int, default: 8)
lora_r: 32
# The LoRA alpha. (type: int, default: 16)
lora_alpha: 16
# The LoRA dropout value. (type: float, default: 0.05)
lora_dropout: 0.05
# Whether to apply LoRA to the query weights in attention. (type: bool, default: True)
lora_query: true
# Whether to apply LoRA to the key weights in attention. (type: bool, default: False)
lora_key: false
# Whether to apply LoRA to the value weights in attention. (type: bool, default: True)
lora_value: true
# Whether to apply LoRA to the output projection in the attention block. (type: bool, default: False)
lora_projection: false
# Whether to apply LoRA to the weights of the MLP in the attention block. (type: bool, default: False)
lora_mlp: false
# Whether to apply LoRA to output head in GPT. (type: bool, default: False)
lora_head: false
# Data-related arguments. If not provided, the default is ``litgpt.data.Alpaca``.
data:
class_path: litgpt.data.Alpaca2k
init_args:
mask_prompt: false
val_split_fraction: 0.05
prompt_style: alpaca
ignore_index: -100
seed: 42
num_workers: 4
download_dir: data/alpaca2k
# Training-related arguments. See ``litgpt.args.TrainArgs`` for details
train:
# Number of optimizer steps between saving checkpoints (type: Optional[int], default: 1000)
save_interval: 200
# Number of iterations between logging calls (type: int, default: 1)
log_interval: 1
# Number of samples between optimizer steps across data-parallel ranks (type: int, default: 128)
global_batch_size: 8
# Number of samples per data-parallel rank (type: int, default: 4)
micro_batch_size: 1
# Number of iterations with learning rate warmup active (type: int, default: 100)
lr_warmup_steps: 10
# Number of epochs to train on (type: Optional[int], default: 5)
epochs: 4
# Total number of tokens to train on (type: Optional[int], default: null)
max_tokens:
# Limits the number of optimizer steps to run (type: Optional[int], default: null)
max_steps:
# Limits the length of samples (type: Optional[int], default: null)
max_seq_length: 512
# Whether to tie the embedding weights with the language modeling head weights (type: Optional[bool], default: null)
tie_embeddings:
# (type: Optional[float], default: null)
max_norm:
# (type: float, default: 6e-05)
min_lr: 6.0e-05
# Evaluation-related arguments. See ``litgpt.args.EvalArgs`` for details
eval:
# Number of optimizer steps between evaluation calls (type: int, default: 100)
interval: 100
# Number of tokens to generate (type: Optional[int], default: 100)
max_new_tokens: 100
# Number of iterations (type: int, default: 100)
max_iters: 100
# Whether to evaluate on the validation set at the beginning of the training
initial_validation: false
# Whether to evaluate on the validation set at the end the training
final_validation: true
# The name of the logger to send metrics to. (type: LoggerChoice, i.e. Literal['wandb', 'tensorboard', 'csv', 'mlflow', 'litlogger'], default: csv)
logger_name: csv
# The random seed to use for reproducibility. (type: int, default: 1337)
seed: 1337
# Optimizer-related arguments
optimizer:
class_path: torch.optim.AdamW
init_args:
# (type: float, default: 0.001)
lr: 0.0002
# (type: float, default: 0.01)
weight_decay: 0.0
# (type: tuple, default: (0.9,0.999))
betas:
- 0.9
- 0.95
config_hub/finetune/gemma-2b/full.yaml
+102
View File
@@ -0,0 +1,102 @@
# The path to the base model's checkpoint directory to load for finetuning. (type: <class 'Path'>, default: checkpoints/stabilityai/stablelm-base-alpha-3b)
checkpoint_dir: checkpoints/google/gemma-2b
# Directory in which to save checkpoints and logs. (type: <class 'Path'>, default: out/lora)
out_dir: out/finetune/full-gemma-2b
# The precision to use for finetuning. Possible choices: "bf16-true", "bf16-mixed", "32-true". (type: Optional[str], default: null)
precision: bf16-true
# How many devices/GPUs to use. (type: Union[int, str], default: 1)
devices: 4
# How many nodes to use. (type: int, default: 1)
num_nodes: 1
# Data-related arguments. If not provided, the default is ``litgpt.data.Alpaca``.
data:
class_path: litgpt.data.Alpaca2k
init_args:
mask_prompt: false
val_split_fraction: 0.03847
prompt_style: alpaca
ignore_index: -100
seed: 42
num_workers: 4
# Training-related arguments. See ``litgpt.args.TrainArgs`` for details
train:
# Number of optimizer steps between saving checkpoints (type: Optional[int], default: 1000)
save_interval: 800
# Number of iterations between logging calls (type: int, default: 1)
log_interval: 1
# Number of samples between optimizer steps across data-parallel ranks (type: int, default: 128)
global_batch_size: 16
# Number of samples per data-parallel rank (type: int, default: 4)
micro_batch_size: 1
# Number of iterations with learning rate warmup active (type: int, default: 100)
lr_warmup_steps: 100
# Number of epochs to train on (type: Optional[int], default: 5)
epochs: 1
# Total number of tokens to train on (type: Optional[int], default: null)
max_tokens:
# Limits the number of optimizer steps to run. (type: Optional[int], default: null)
max_steps: 50
# Limits the length of samples. Off by default (type: Optional[int], default: null)
max_seq_length: 512
# Whether to tie the embedding weights with the language modeling head weights. (type: Optional[bool], default: null)
tie_embeddings:
# (type: Optional[float], default: null)
max_norm:
# (type: float, default: 6e-05)
min_lr: 6.0e-05
# Evaluation-related arguments. See ``litgpt.args.EvalArgs`` for details
eval:
# Number of optimizer steps between evaluation calls (type: int, default: 100)
interval: 25
# Number of tokens to generate (type: Optional[int], default: 100)
max_new_tokens: 100
# Number of iterations (type: int, default: 100)
max_iters: 100
# Whether to evaluate on the validation set at the beginning of the training
initial_validation: false
# Whether to evaluate on the validation set at the end the training
final_validation: true
# The name of the logger to send metrics to. (type: LoggerChoice, i.e. Literal['wandb', 'tensorboard', 'csv', 'mlflow', 'litlogger'], default: csv)
logger_name: csv
# The random seed to use for reproducibility. (type: int, default: 1337)
seed: 1337
# Optimizer-related arguments
optimizer:
class_path: torch.optim.AdamW
init_args:
# (type: float, default: 0.001)
lr: 0.0002
# (type: float, default: 0.01)
weight_decay: 0.0
# (type: tuple, default: (0.9,0.999))
betas:
- 0.9
- 0.95
config_hub/finetune/gemma-2b/lora.yaml
+132
View File
@@ -0,0 +1,132 @@
# The path to the base model's checkpoint directory to load for finetuning. (type: <class 'Path'>, default: checkpoints/stabilityai/stablelm-base-alpha-3b)
checkpoint_dir: checkpoints/google/gemma-2b
# Directory in which to save checkpoints and logs. (type: <class 'Path'>, default: out/lora)
out_dir: out/finetune/lora-gemma-2b
# The precision to use for finetuning. Possible choices: "bf16-true", "bf16-mixed", "32-true". (type: Optional[str], default: null)
precision: bf16-true
# If set, quantize the model with this algorithm. See ``tutorials/quantize.md`` for more information. (type: Optional[Literal['nf4', 'nf4-dq', 'fp4', 'fp4-dq', 'int8-training']], default: null)
quantize:
# How many devices/GPUs to use. (type: Union[int, str], default: 1)
devices: 1
# How many nodes to use. (type: int, default: 1)
num_nodes: 1
# The LoRA rank. (type: int, default: 8)
lora_r: 8
# The LoRA alpha. (type: int, default: 16)
lora_alpha: 16
# The LoRA dropout value. (type: float, default: 0.05)
lora_dropout: 0.1
# Whether to apply LoRA to the query weights in attention. (type: bool, default: True)
lora_query: true
# Whether to apply LoRA to the key weights in attention. (type: bool, default: False)
lora_key: true
# Whether to apply LoRA to the value weights in attention. (type: bool, default: True)
lora_value: true
# Whether to apply LoRA to the output projection in the attention block. (type: bool, default: False)
lora_projection: true
# Whether to apply LoRA to the weights of the MLP in the attention block. (type: bool, default: False)
lora_mlp: true
# Whether to apply LoRA to output head in GPT. (type: bool, default: False)
lora_head: true
# Data-related arguments. If not provided, the default is ``litgpt.data.Alpaca``.
data:
class_path: litgpt.data.Alpaca2k
init_args:
mask_prompt: false
val_split_fraction: 0.03847
prompt_style: alpaca
ignore_index: -100
seed: 42
num_workers: 4
# Training-related arguments. See ``litgpt.args.TrainArgs`` for details
train:
# Number of optimizer steps between saving checkpoints (type: Optional[int], default: 1000)
save_interval: 800
# Number of iterations between logging calls (type: int, default: 1)
log_interval: 1
# Number of samples between optimizer steps across data-parallel ranks (type: int, default: 128)
global_batch_size: 6
# Number of samples per data-parallel rank (type: int, default: 4)
micro_batch_size: 2
# Number of iterations with learning rate warmup active (type: int, default: 100)
lr_warmup_steps: 200
# Number of epochs to train on (type: Optional[int], default: 5)
epochs: 2
# Total number of tokens to train on (type: Optional[int], default: null)
max_tokens:
# Limits the number of optimizer steps to run. (type: Optional[int], default: null)
max_steps:
# Limits the length of samples. Off by default (type: Optional[int], default: null)
max_seq_length: 512
# Whether to tie the embedding weights with the language modeling head weights. (type: Optional[bool], default: null)
tie_embeddings:
# (type: Optional[float], default: null)
max_norm:
# (type: float, default: 6e-05)
min_lr: 6.0e-05
# Evaluation-related arguments. See ``litgpt.args.EvalArgs`` for details
eval:
# Number of optimizer steps between evaluation calls (type: int, default: 100)
interval: 25
# Number of tokens to generate (type: Optional[int], default: 100)
max_new_tokens: 100
# Number of iterations (type: int, default: 100)
max_iters: 100
# Whether to evaluate on the validation set at the beginning of the training
initial_validation: false
# Whether to evaluate on the validation set at the end the training
final_validation: true
# The name of the logger to send metrics to. (type: LoggerChoice, i.e. Literal['wandb', 'tensorboard', 'csv', 'mlflow', 'litlogger'], default: csv)
logger_name: csv
# The random seed to use for reproducibility. (type: int, default: 1337)
seed: 1337
# Optimizer-related arguments
optimizer:
class_path: torch.optim.AdamW
init_args:
# (type: float, default: 0.001)
lr: 0.0002
# (type: float, default: 0.01)
weight_decay: 0.0
# (type: tuple, default: (0.9,0.999))
betas:
- 0.9
- 0.95
config_hub/finetune/gemma-2b/qlora.yaml
+132
View File
@@ -0,0 +1,132 @@
# The path to the base model's checkpoint directory to load for finetuning. (type: <class 'Path'>, default: checkpoints/stabilityai/stablelm-base-alpha-3b)
checkpoint_dir: checkpoints/google/gemma-2b
# Directory in which to save checkpoints and logs. (type: <class 'Path'>, default: out/lora)
out_dir: out/finetune/qlora-gemma-2b
# The precision to use for finetuning. Possible choices: "bf16-true", "bf16-mixed", "32-true". (type: Optional[str], default: null)
precision: bf16-true
# If set, quantize the model with this algorithm. See ``tutorials/quantize.md`` for more information. (type: Optional[Literal['nf4', 'nf4-dq', 'fp4', 'fp4-dq', 'int8-training']], default: null)
quantize: bnb.nf4
# How many devices/GPUs to use. (type: Union[int, str], default: 1)
devices: 1
# How many nodes to use. (type: int, default: 1)
num_nodes: 1
# The LoRA rank. (type: int, default: 8)
lora_r: 16
# The LoRA alpha. (type: int, default: 16)
lora_alpha: 16
# The LoRA dropout value. (type: float, default: 0.05)
lora_dropout: 0.1
# Whether to apply LoRA to the query weights in attention. (type: bool, default: True)
lora_query: true
# Whether to apply LoRA to the key weights in attention. (type: bool, default: False)
lora_key: true
# Whether to apply LoRA to the value weights in attention. (type: bool, default: True)
lora_value: true
# Whether to apply LoRA to the output projection in the attention block. (type: bool, default: False)
lora_projection: true
# Whether to apply LoRA to the weights of the MLP in the attention block. (type: bool, default: False)
lora_mlp: true
# Whether to apply LoRA to output head in GPT. (type: bool, default: False)
lora_head: true
# Data-related arguments. If not provided, the default is ``litgpt.data.Alpaca``.
data:
class_path: litgpt.data.Alpaca2k
init_args:
mask_prompt: false
val_split_fraction: 0.03847
prompt_style: alpaca
ignore_index: -100
seed: 42
num_workers: 4
# Training-related arguments. See ``litgpt.args.TrainArgs`` for details
train:
# Number of optimizer steps between saving checkpoints (type: Optional[int], default: 1000)
save_interval: 800
# Number of iterations between logging calls (type: int, default: 1)
log_interval: 1
# Number of samples between optimizer steps across data-parallel ranks (type: int, default: 128)
global_batch_size: 6
# Number of samples per data-parallel rank (type: int, default: 4)
micro_batch_size: 2
# Number of iterations with learning rate warmup active (type: int, default: 100)
lr_warmup_steps: 200
# Number of epochs to train on (type: Optional[int], default: 5)
epochs: 2
# Total number of tokens to train on (type: Optional[int], default: null)
max_tokens:
# Limits the number of optimizer steps to run. (type: Optional[int], default: null)
max_steps:
# Limits the length of samples. Off by default (type: Optional[int], default: null)
max_seq_length: 512
# Whether to tie the embedding weights with the language modeling head weights. (type: Optional[bool], default: null)
tie_embeddings:
# (type: Optional[float], default: null)
max_norm:
# (type: float, default: 6e-05)
min_lr: 6.0e-05
# Evaluation-related arguments. See ``litgpt.args.EvalArgs`` for details
eval:
# Number of optimizer steps between evaluation calls (type: int, default: 100)
interval: 25
# Number of tokens to generate (type: Optional[int], default: 100)
max_new_tokens: 100
# Number of iterations (type: int, default: 100)
max_iters: 100
# Whether to evaluate on the validation set at the beginning of the training
initial_validation: false
# Whether to evaluate on the validation set at the end the training
final_validation: true
# The name of the logger to send metrics to. (type: LoggerChoice, i.e. Literal['wandb', 'tensorboard', 'csv', 'mlflow', 'litlogger'], default: csv)
logger_name: csv
# The random seed to use for reproducibility. (type: int, default: 1337)
seed: 1337
# Optimizer-related arguments
optimizer:
class_path: torch.optim.AdamW
init_args:
# (type: float, default: 0.001)
lr: 0.0002
# (type: float, default: 0.01)
weight_decay: 0.0
# (type: tuple, default: (0.9,0.999))
betas:
- 0.9
- 0.95
config_hub/finetune/gemma-7b/lora.yaml
+132
View File
@@ -0,0 +1,132 @@
# The path to the base model's checkpoint directory to load for finetuning. (type: <class 'Path'>, default: checkpoints/stabilityai/stablelm-base-alpha-3b)
checkpoint_dir: checkpoints/google/gemma-7b
# Directory in which to save checkpoints and logs. (type: <class 'Path'>, default: out/lora)
out_dir: out/finetune/qlora-gemma-7b
# The precision to use for finetuning. Possible choices: "bf16-true", "bf16-mixed", "32-true". (type: Optional[str], default: null)
precision: bf16-true
# If set, quantize the model with this algorithm. See ``tutorials/quantize.md`` for more information. (type: Optional[Literal['nf4', 'nf4-dq', 'fp4', 'fp4-dq', 'int8-training']], default: null)
quantize:
# How many devices/GPUs to use. (type: Union[int, str], default: 1)
devices: 1
# How many nodes to use. (type: int, default: 1)
num_nodes: 1
# The LoRA rank. (type: int, default: 8)
lora_r: 16
# The LoRA alpha. (type: int, default: 16)
lora_alpha: 16
# The LoRA dropout value. (type: float, default: 0.05)
lora_dropout: 0.1
# Whether to apply LoRA to the query weights in attention. (type: bool, default: True)
lora_query: true
# Whether to apply LoRA to the key weights in attention. (type: bool, default: False)
lora_key: true
# Whether to apply LoRA to the value weights in attention. (type: bool, default: True)
lora_value: true
# Whether to apply LoRA to the output projection in the attention block. (type: bool, default: False)
lora_projection: true
# Whether to apply LoRA to the weights of the MLP in the attention block. (type: bool, default: False)
lora_mlp: true
# Whether to apply LoRA to output head in GPT. (type: bool, default: False)
lora_head: true
# Data-related arguments. If not provided, the default is ``litgpt.data.Alpaca``.
data:
class_path: litgpt.data.Alpaca2k
init_args:
mask_prompt: false
val_split_fraction: 0.03847
prompt_style: alpaca
ignore_index: -100
seed: 42
num_workers: 4
# Training-related arguments. See ``litgpt.args.TrainArgs`` for details
train:
# Number of optimizer steps between saving checkpoints (type: Optional[int], default: 1000)
save_interval: 800
# Number of iterations between logging calls (type: int, default: 1)
log_interval: 1
# Number of samples between optimizer steps across data-parallel ranks (type: int, default: 128)
global_batch_size: 6
# Number of samples per data-parallel rank (type: int, default: 4)
micro_batch_size: 1
# Number of iterations with learning rate warmup active (type: int, default: 100)
lr_warmup_steps: 200
# Number of epochs to train on (type: Optional[int], default: 5)
epochs: 2
# Total number of tokens to train on (type: Optional[int], default: null)
max_tokens:
# Limits the number of optimizer steps to run. (type: Optional[int], default: null)
max_steps:
# Limits the length of samples. Off by default (type: Optional[int], default: null)
max_seq_length: 512
# Whether to tie the embedding weights with the language modeling head weights. (type: Optional[bool], default: null)
tie_embeddings:
# (type: Optional[float], default: null)
max_norm:
# (type: float, default: 6e-05)
min_lr: 6.0e-05
# Evaluation-related arguments. See ``litgpt.args.EvalArgs`` for details
eval:
# Number of optimizer steps between evaluation calls (type: int, default: 100)
interval: 25
# Number of tokens to generate (type: Optional[int], default: 100)
max_new_tokens: 100
# Number of iterations (type: int, default: 100)
max_iters: 100
# Whether to evaluate on the validation set at the beginning of the training
initial_validation: false
# Whether to evaluate on the validation set at the end the training
final_validation: true
# The name of the logger to send metrics to. (type: LoggerChoice, i.e. Literal['wandb', 'tensorboard', 'csv', 'mlflow', 'litlogger'], default: csv)
logger_name: csv
# The random seed to use for reproducibility. (type: int, default: 1337)
seed: 1337
# Optimizer-related arguments
optimizer:
class_path: torch.optim.AdamW
init_args:
# (type: float, default: 0.001)
lr: 0.0002
# (type: float, default: 0.01)
weight_decay: 0.0
# (type: tuple, default: (0.9,0.999))
betas:
- 0.9
- 0.95
config_hub/finetune/gemma-7b/qlora.yaml
+132
View File
@@ -0,0 +1,132 @@
# The path to the base model's checkpoint directory to load for finetuning. (type: <class 'Path'>, default: checkpoints/stabilityai/stablelm-base-alpha-3b)
checkpoint_dir: checkpoints/google/gemma-7b
# Directory in which to save checkpoints and logs. (type: <class 'Path'>, default: out/lora)
out_dir: out/finetune/qlora-gemma-7b
# The precision to use for finetuning. Possible choices: "bf16-true", "bf16-mixed", "32-true". (type: Optional[str], default: null)
precision: bf16-true
# If set, quantize the model with this algorithm. See ``tutorials/quantize.md`` for more information. (type: Optional[Literal['nf4', 'nf4-dq', 'fp4', 'fp4-dq', 'int8-training']], default: null)
quantize: bnb.nf4
# How many devices/GPUs to use. (type: Union[int, str], default: 1)
devices: 1
# How many nodes to use. (type: int, default: 1)
num_nodes: 1
# The LoRA rank. (type: int, default: 8)
lora_r: 16
# The LoRA alpha. (type: int, default: 16)
lora_alpha: 16
# The LoRA dropout value. (type: float, default: 0.05)
lora_dropout: 0.1
# Whether to apply LoRA to the query weights in attention. (type: bool, default: True)
lora_query: true
# Whether to apply LoRA to the key weights in attention. (type: bool, default: False)
lora_key: true
# Whether to apply LoRA to the value weights in attention. (type: bool, default: True)
lora_value: true
# Whether to apply LoRA to the output projection in the attention block. (type: bool, default: False)
lora_projection: true
# Whether to apply LoRA to the weights of the MLP in the attention block. (type: bool, default: False)
lora_mlp: true
# Whether to apply LoRA to output head in GPT. (type: bool, default: False)
lora_head: true
# Data-related arguments. If not provided, the default is ``litgpt.data.Alpaca``.
data:
class_path: litgpt.data.Alpaca2k
init_args:
mask_prompt: false
val_split_fraction: 0.03847
prompt_style: alpaca
ignore_index: -100
seed: 42
num_workers: 4
# Training-related arguments. See ``litgpt.args.TrainArgs`` for details
train:
# Number of optimizer steps between saving checkpoints (type: Optional[int], default: 1000)
save_interval: 800
# Number of iterations between logging calls (type: int, default: 1)
log_interval: 1
# Number of samples between optimizer steps across data-parallel ranks (type: int, default: 128)
global_batch_size: 6
# Number of samples per data-parallel rank (type: int, default: 4)
micro_batch_size: 1
# Number of iterations with learning rate warmup active (type: int, default: 100)
lr_warmup_steps: 200
# Number of epochs to train on (type: Optional[int], default: 5)
epochs: 2
# Total number of tokens to train on (type: Optional[int], default: null)
max_tokens:
# Limits the number of optimizer steps to run. (type: Optional[int], default: null)
max_steps:
# Limits the length of samples. Off by default (type: Optional[int], default: null)
max_seq_length: 512
# Whether to tie the embedding weights with the language modeling head weights. (type: Optional[bool], default: null)
tie_embeddings:
# (type: Optional[float], default: null)
max_norm:
# (type: float, default: 6e-05)
min_lr: 6.0e-05
# Evaluation-related arguments. See ``litgpt.args.EvalArgs`` for details
eval:
# Number of optimizer steps between evaluation calls (type: int, default: 100)
interval: 25
# Number of tokens to generate (type: Optional[int], default: 100)
max_new_tokens: 100
# Number of iterations (type: int, default: 100)
max_iters: 100
# Whether to evaluate on the validation set at the beginning of the training
initial_validation: false
# Whether to evaluate on the validation set at the end the training
final_validation: true
# The name of the logger to send metrics to. (type: LoggerChoice, i.e. Literal['wandb', 'tensorboard', 'csv', 'mlflow', 'litlogger'], default: csv)
logger_name: csv
# The random seed to use for reproducibility. (type: int, default: 1337)
seed: 1337
# Optimizer-related arguments
optimizer:
class_path: torch.optim.AdamW
init_args:
# (type: float, default: 0.001)
lr: 0.0002
# (type: float, default: 0.01)
weight_decay: 0.0
# (type: tuple, default: (0.9,0.999))
betas:
- 0.9
- 0.95
config_hub/finetune/gemma2-2b/lora.yaml
+132
View File
@@ -0,0 +1,132 @@
# The path to the base model's checkpoint directory to load for finetuning. (type: <class 'Path'>, default: checkpoints/stabilityai/stablelm-base-alpha-3b)
checkpoint_dir: checkpoints/google/gemma-2-2b
# Directory in which to save checkpoints and logs. (type: <class 'Path'>, default: out/lora)
out_dir: out/finetune/lora-gemma-2-2b
# The precision to use for finetuning. Possible choices: "bf16-true", "bf16-mixed", "32-true". (type: Optional[str], default: null)
precision: bf16-true
# If set, quantize the model with this algorithm. See ``tutorials/quantize.md`` for more information. (type: Optional[Literal['nf4', 'nf4-dq', 'fp4', 'fp4-dq', 'int8-training']], default: null)
quantize:
# How many devices/GPUs to use. (type: Union[int, str], default: 1)
devices: 1
# How many nodes to use. (type: int, default: 1)
num_nodes: 1
# The LoRA rank. (type: int, default: 8)
lora_r: 8
# The LoRA alpha. (type: int, default: 16)
lora_alpha: 16
# The LoRA dropout value. (type: float, default: 0.05)
lora_dropout: 0.1
# Whether to apply LoRA to the query weights in attention. (type: bool, default: True)
lora_query: true
# Whether to apply LoRA to the key weights in attention. (type: bool, default: False)
lora_key: true
# Whether to apply LoRA to the value weights in attention. (type: bool, default: True)
lora_value: true
# Whether to apply LoRA to the output projection in the attention block. (type: bool, default: False)
lora_projection: true
# Whether to apply LoRA to the weights of the MLP in the attention block. (type: bool, default: False)
lora_mlp: true
# Whether to apply LoRA to output head in GPT. (type: bool, default: False)
lora_head: true
# Data-related arguments. If not provided, the default is ``litgpt.data.Alpaca``.
data:
class_path: litgpt.data.Alpaca2k
init_args:
mask_prompt: false
val_split_fraction: 0.03847
prompt_style: alpaca
ignore_index: -100
seed: 42
num_workers: 4
# Training-related arguments. See ``litgpt.args.TrainArgs`` for details
train:
# Number of optimizer steps between saving checkpoints (type: Optional[int], default: 1000)
save_interval: 800
# Number of iterations between logging calls (type: int, default: 1)
log_interval: 1
# Number of samples between optimizer steps across data-parallel ranks (type: int, default: 128)
global_batch_size: 6
# Number of samples per data-parallel rank (type: int, default: 4)
micro_batch_size: 2
# Number of iterations with learning rate warmup active (type: int, default: 100)
lr_warmup_steps: 200
# Number of epochs to train on (type: Optional[int], default: 5)
epochs: 2
# Total number of tokens to train on (type: Optional[int], default: null)
max_tokens:
# Limits the number of optimizer steps to run. (type: Optional[int], default: null)
max_steps:
# Limits the length of samples. Off by default (type: Optional[int], default: null)
max_seq_length: 512
# Whether to tie the embedding weights with the language modeling head weights. (type: Optional[bool], default: null)
tie_embeddings:
# (type: Optional[float], default: null)
max_norm:
# (type: float, default: 6e-05)
min_lr: 6.0e-05
# Evaluation-related arguments. See ``litgpt.args.EvalArgs`` for details
eval:
# Number of optimizer steps between evaluation calls (type: int, default: 100)
interval: 25
# Number of tokens to generate (type: Optional[int], default: 100)
max_new_tokens: 100
# Number of iterations (type: int, default: 100)
max_iters: 100
# Whether to evaluate on the validation set at the beginning of the training
initial_validation: false
# Whether to evaluate on the validation set at the end the training
final_validation: true
# The name of the logger to send metrics to. (type: LoggerChoice, i.e. Literal['wandb', 'tensorboard', 'csv', 'mlflow', 'litlogger'], default: csv)
logger_name: csv
# The random seed to use for reproducibility. (type: int, default: 1337)
seed: 1337
# Optimizer-related arguments
optimizer:
class_path: torch.optim.AdamW
init_args:
# (type: float, default: 0.001)
lr: 0.0002
# (type: float, default: 0.01)
weight_decay: 0.0
# (type: tuple, default: (0.9,0.999))
betas:
- 0.9
- 0.95
config_hub/finetune/gemma2-2b/qlora.yaml
+132
View File
@@ -0,0 +1,132 @@
# The path to the base model's checkpoint directory to load for finetuning. (type: <class 'Path'>, default: checkpoints/stabilityai/stablelm-base-alpha-3b)
checkpoint_dir: checkpoints/google/gemma-2-2b
# Directory in which to save checkpoints and logs. (type: <class 'Path'>, default: out/lora)
out_dir: out/finetune/qlora-gemma-2-2b
# The precision to use for finetuning. Possible choices: "bf16-true", "bf16-mixed", "32-true". (type: Optional[str], default: null)
precision: bf16-true
# If set, quantize the model with this algorithm. See ``tutorials/quantize.md`` for more information. (type: Optional[Literal['nf4', 'nf4-dq', 'fp4', 'fp4-dq', 'int8-training']], default: null)
quantize: bnb.nf4
# How many devices/GPUs to use. (type: Union[int, str], default: 1)
devices: 1
# How many nodes to use. (type: int, default: 1)
num_nodes: 1
# The LoRA rank. (type: int, default: 8)
lora_r: 16
# The LoRA alpha. (type: int, default: 16)
lora_alpha: 16
# The LoRA dropout value. (type: float, default: 0.05)
lora_dropout: 0.1
# Whether to apply LoRA to the query weights in attention. (type: bool, default: True)
lora_query: true
# Whether to apply LoRA to the key weights in attention. (type: bool, default: False)
lora_key: true
# Whether to apply LoRA to the value weights in attention. (type: bool, default: True)
lora_value: true
# Whether to apply LoRA to the output projection in the attention block. (type: bool, default: False)
lora_projection: true
# Whether to apply LoRA to the weights of the MLP in the attention block. (type: bool, default: False)
lora_mlp: true
# Whether to apply LoRA to output head in GPT. (type: bool, default: False)
lora_head: true
# Data-related arguments. If not provided, the default is ``litgpt.data.Alpaca``.
data:
class_path: litgpt.data.Alpaca2k
init_args:
mask_prompt: false
val_split_fraction: 0.03847
prompt_style: alpaca
ignore_index: -100
seed: 42
num_workers: 4
# Training-related arguments. See ``litgpt.args.TrainArgs`` for details
train:
# Number of optimizer steps between saving checkpoints (type: Optional[int], default: 1000)
save_interval: 800
# Number of iterations between logging calls (type: int, default: 1)
log_interval: 1
# Number of samples between optimizer steps across data-parallel ranks (type: int, default: 128)
global_batch_size: 6
# Number of samples per data-parallel rank (type: int, default: 4)
micro_batch_size: 2
# Number of iterations with learning rate warmup active (type: int, default: 100)
lr_warmup_steps: 200
# Number of epochs to train on (type: Optional[int], default: 5)
epochs: 2
# Total number of tokens to train on (type: Optional[int], default: null)
max_tokens:
# Limits the number of optimizer steps to run. (type: Optional[int], default: null)
max_steps:
# Limits the length of samples. Off by default (type: Optional[int], default: null)
max_seq_length: 512
# Whether to tie the embedding weights with the language modeling head weights. (type: Optional[bool], default: null)
tie_embeddings:
# (type: Optional[float], default: null)
max_norm:
# (type: float, default: 6e-05)
min_lr: 6.0e-05
# Evaluation-related arguments. See ``litgpt.args.EvalArgs`` for details
eval:
# Number of optimizer steps between evaluation calls (type: int, default: 100)
interval: 25
# Number of tokens to generate (type: Optional[int], default: 100)
max_new_tokens: 100
# Number of iterations (type: int, default: 100)
max_iters: 100
# Whether to evaluate on the validation set at the beginning of the training
initial_validation: false
# Whether to evaluate on the validation set at the end the training
final_validation: true
# The name of the logger to send metrics to. (type: LoggerChoice, i.e. Literal['wandb', 'tensorboard', 'csv', 'mlflow', 'litlogger'], default: csv)
logger_name: csv
# The random seed to use for reproducibility. (type: int, default: 1337)
seed: 1337
# Optimizer-related arguments
optimizer:
class_path: torch.optim.AdamW
init_args:
# (type: float, default: 0.001)
lr: 0.0002
# (type: float, default: 0.01)
weight_decay: 0.0
# (type: tuple, default: (0.9,0.999))
betas:
- 0.9
- 0.95
config_hub/finetune/gemma2-9b/lora.yaml
+132
View File
@@ -0,0 +1,132 @@
# The path to the base model's checkpoint directory to load for finetuning. (type: <class 'Path'>, default: checkpoints/stabilityai/stablelm-base-alpha-3b)
checkpoint_dir: checkpoints/google/gemma-2-9b
# Directory in which to save checkpoints and logs. (type: <class 'Path'>, default: out/lora)
out_dir: out/finetune/lora-gemma-2-9b
# The precision to use for finetuning. Possible choices: "bf16-true", "bf16-mixed", "32-true". (type: Optional[str], default: null)
precision: bf16-true
# If set, quantize the model with this algorithm. See ``tutorials/quantize.md`` for more information. (type: Optional[Literal['nf4', 'nf4-dq', 'fp4', 'fp4-dq', 'int8-training']], default: null)
quantize:
# How many devices/GPUs to use. (type: Union[int, str], default: 1)
devices: 1
# How many nodes to use. (type: int, default: 1)
num_nodes: 1
# The LoRA rank. (type: int, default: 8)
lora_r: 16
# The LoRA alpha. (type: int, default: 16)
lora_alpha: 16
# The LoRA dropout value. (type: float, default: 0.05)
lora_dropout: 0.1
# Whether to apply LoRA to the query weights in attention. (type: bool, default: True)
lora_query: true
# Whether to apply LoRA to the key weights in attention. (type: bool, default: False)
lora_key: true
# Whether to apply LoRA to the value weights in attention. (type: bool, default: True)
lora_value: true
# Whether to apply LoRA to the output projection in the attention block. (type: bool, default: False)
lora_projection: true
# Whether to apply LoRA to the weights of the MLP in the attention block. (type: bool, default: False)
lora_mlp: true
# Whether to apply LoRA to output head in GPT. (type: bool, default: False)
lora_head: true
# Data-related arguments. If not provided, the default is ``litgpt.data.Alpaca``.
data:
class_path: litgpt.data.Alpaca2k
init_args:
mask_prompt: false
val_split_fraction: 0.03847
prompt_style: alpaca
ignore_index: -100
seed: 42
num_workers: 4
# Training-related arguments. See ``litgpt.args.TrainArgs`` for details
train:
# Number of optimizer steps between saving checkpoints (type: Optional[int], default: 1000)
save_interval: 800
# Number of iterations between logging calls (type: int, default: 1)
log_interval: 1
# Number of samples between optimizer steps across data-parallel ranks (type: int, default: 128)
global_batch_size: 6
# Number of samples per data-parallel rank (type: int, default: 4)
micro_batch_size: 1
# Number of iterations with learning rate warmup active (type: int, default: 100)
lr_warmup_steps: 200
# Number of epochs to train on (type: Optional[int], default: 5)
epochs: 2
# Total number of tokens to train on (type: Optional[int], default: null)
max_tokens:
# Limits the number of optimizer steps to run. (type: Optional[int], default: null)
max_steps:
# Limits the length of samples. Off by default (type: Optional[int], default: null)
max_seq_length: 512
# Whether to tie the embedding weights with the language modeling head weights. (type: Optional[bool], default: null)
tie_embeddings:
# (type: Optional[float], default: null)
max_norm:
# (type: float, default: 6e-05)
min_lr: 6.0e-05
# Evaluation-related arguments. See ``litgpt.args.EvalArgs`` for details
eval:
# Number of optimizer steps between evaluation calls (type: int, default: 100)
interval: 25
# Number of tokens to generate (type: Optional[int], default: 100)
max_new_tokens: 100
# Number of iterations (type: int, default: 100)
max_iters: 100
# Whether to evaluate on the validation set at the beginning of the training
initial_validation: false
# Whether to evaluate on the validation set at the end the training
final_validation: true
# The name of the logger to send metrics to. (type: LoggerChoice, i.e. Literal['wandb', 'tensorboard', 'csv', 'mlflow', 'litlogger'], default: csv)
logger_name: csv
# The random seed to use for reproducibility. (type: int, default: 1337)
seed: 1337
# Optimizer-related arguments
optimizer:
class_path: torch.optim.AdamW
init_args:
# (type: float, default: 0.001)
lr: 0.0002
# (type: float, default: 0.01)
weight_decay: 0.0
# (type: tuple, default: (0.9,0.999))
betas:
- 0.9
- 0.95
config_hub/finetune/gemma2-9b/qlora.yaml
+132
View File
@@ -0,0 +1,132 @@
# The path to the base model's checkpoint directory to load for finetuning. (type: <class 'Path'>, default: checkpoints/stabilityai/stablelm-base-alpha-3b)
checkpoint_dir: checkpoints/google/gemma-2-9b
# Directory in which to save checkpoints and logs. (type: <class 'Path'>, default: out/lora)
out_dir: out/finetune/qlora-gemma-2-9b
# The precision to use for finetuning. Possible choices: "bf16-true", "bf16-mixed", "32-true". (type: Optional[str], default: null)
precision: bf16-true
# If set, quantize the model with this algorithm. See ``tutorials/quantize.md`` for more information. (type: Optional[Literal['nf4', 'nf4-dq', 'fp4', 'fp4-dq', 'int8-training']], default: null)
quantize: bnb.nf4
# How many devices/GPUs to use. (type: Union[int, str], default: 1)
devices: 1
# How many nodes to use. (type: int, default: 1)
num_nodes: 1
# The LoRA rank. (type: int, default: 8)
lora_r: 16
# The LoRA alpha. (type: int, default: 16)
lora_alpha: 16
# The LoRA dropout value. (type: float, default: 0.05)
lora_dropout: 0.1
# Whether to apply LoRA to the query weights in attention. (type: bool, default: True)
lora_query: true
# Whether to apply LoRA to the key weights in attention. (type: bool, default: False)
lora_key: true
# Whether to apply LoRA to the value weights in attention. (type: bool, default: True)
lora_value: true
# Whether to apply LoRA to the output projection in the attention block. (type: bool, default: False)
lora_projection: true
# Whether to apply LoRA to the weights of the MLP in the attention block. (type: bool, default: False)
lora_mlp: true
# Whether to apply LoRA to output head in GPT. (type: bool, default: False)
lora_head: true
# Data-related arguments. If not provided, the default is ``litgpt.data.Alpaca``.
data:
class_path: litgpt.data.Alpaca2k
init_args:
mask_prompt: false
val_split_fraction: 0.03847
prompt_style: alpaca
ignore_index: -100
seed: 42
num_workers: 4
# Training-related arguments. See ``litgpt.args.TrainArgs`` for details
train:
# Number of optimizer steps between saving checkpoints (type: Optional[int], default: 1000)
save_interval: 800
# Number of iterations between logging calls (type: int, default: 1)
log_interval: 1
# Number of samples between optimizer steps across data-parallel ranks (type: int, default: 128)
global_batch_size: 6
# Number of samples per data-parallel rank (type: int, default: 4)
micro_batch_size: 1
# Number of iterations with learning rate warmup active (type: int, default: 100)
lr_warmup_steps: 200
# Number of epochs to train on (type: Optional[int], default: 5)
epochs: 2
# Total number of tokens to train on (type: Optional[int], default: null)
max_tokens:
# Limits the number of optimizer steps to run. (type: Optional[int], default: null)
max_steps:
# Limits the length of samples. Off by default (type: Optional[int], default: null)
max_seq_length: 512
# Whether to tie the embedding weights with the language modeling head weights. (type: Optional[bool], default: null)
tie_embeddings:
# (type: Optional[float], default: null)
max_norm:
# (type: float, default: 6e-05)
min_lr: 6.0e-05
# Evaluation-related arguments. See ``litgpt.args.EvalArgs`` for details
eval:
# Number of optimizer steps between evaluation calls (type: int, default: 100)
interval: 25
# Number of tokens to generate (type: Optional[int], default: 100)
max_new_tokens: 100
# Number of iterations (type: int, default: 100)
max_iters: 100
# Whether to evaluate on the validation set at the beginning of the training
initial_validation: false
# Whether to evaluate on the validation set at the end the training
final_validation: true
# The name of the logger to send metrics to. (type: LoggerChoice, i.e. Literal['wandb', 'tensorboard', 'csv', 'mlflow', 'litlogger'], default: csv)
logger_name: csv
# The random seed to use for reproducibility. (type: int, default: 1337)
seed: 1337
# Optimizer-related arguments
optimizer:
class_path: torch.optim.AdamW
init_args:
# (type: float, default: 0.001)
lr: 0.0002
# (type: float, default: 0.01)
weight_decay: 0.0
# (type: tuple, default: (0.9,0.999))
betas:
- 0.9
- 0.95
config_hub/finetune/llama-2-7b/full.yaml
+107
View File
@@ -0,0 +1,107 @@
# The path to the base model's checkpoint directory to load for finetuning. (type: <class 'Path'>, default: checkpoints/stabilityai/stablelm-base-alpha-3b)
checkpoint_dir: checkpoints/meta-llama/Llama-2-7b-hf
# Directory in which to save checkpoints and logs. (type: <class 'Path'>, default: out/finetune/full)
out_dir: out/finetune/full-llama2-7b
# The precision to use for finetuning. Possible choices: "bf16-true", "bf16-mixed", "32-true". (type: Optional[str], default: null)
precision: bf16-true
# How many devices/GPUs to use (type: Union[int, str], default: 1)
devices: 4
# How many nodes to use. (type: int, default: 1)
num_nodes: 1
# Path to a checkpoint directory to resume from in case training was interrupted, or ``True`` to resume
# from the latest checkpoint in ``out_dir``. An error will be raised if no checkpoint is found. Passing
# ``'auto'`` will resume from the latest checkpoint but not error if no checkpoint exists.
# (type: Union[bool, Literal["auto"], Path], default: False)
resume: false
# Data-related arguments. If not provided, the default is ``litgpt.data.Alpaca``.
data:
class_path: litgpt.data.Alpaca2k
init_args:
mask_prompt: false
prompt_style: alpaca
ignore_index: -100
seed: 42
num_workers: 4
# Training-related arguments. See ``litgpt.args.TrainArgs`` for details
train:
# Number of optimizer steps between saving checkpoints (type: Optional[int], default: 1000)
save_interval: 200
# Number of iterations between logging calls (type: int, default: 1)
log_interval: 1
# Number of samples between optimizer steps across data-parallel ranks (type: int, default: 64)
global_batch_size: 64
# Number of samples per data-parallel rank (type: int, default: 1)
micro_batch_size: 4
# Number of iterations with learning rate warmup active (type: int, default: 100)
lr_warmup_steps: 25
# Number of epochs to train on (type: Optional[int], default: 5)
epochs: 1
# Total number of tokens to train on (type: Optional[int], default: null)
max_tokens:
# Limits the number of optimizer steps to run. (type: Optional[int], default: null)
max_steps:
# Limits the length of samples. Off by default (type: Optional[int], default: null)
max_seq_length: 512
# Whether to tie the embedding weights with the language modeling head weights. (type: Optional[bool], default: null)
tie_embeddings:
# (type: Optional[float], default: null)
max_norm:
# (type: float, default: 6e-05)
min_lr: 6.0e-05
# Evaluation-related arguments. See ``litgpt.args.EvalArgs`` for details
eval:
# Number of optimizer steps between evaluation calls (type: int, default: 600)
interval: 25
# Number of tokens to generate (type: Optional[int], default: 100)
max_new_tokens: 100
# Number of iterations (type: int, default: 100)
max_iters: 100
# Whether to evaluate on the validation set at the beginning of the training
initial_validation: false
# Whether to evaluate on the validation set at the end the training
final_validation: true
# The name of the logger to send metrics to. (type: LoggerChoice, i.e. Literal['wandb', 'tensorboard', 'csv', 'mlflow', 'litlogger'], default: csv)
logger_name: csv
# The random seed to use for reproducibility. (type: int, default: 1337)
seed: 1337
# Optimizer-related arguments
optimizer:
class_path: torch.optim.AdamW
init_args:
# (type: float, default: 0.001)
lr: 0.0002
# (type: float, default: 0.01)
weight_decay: 0.0
# (type: tuple, default: (0.9,0.999))
betas:
- 0.9
- 0.95
config_hub/finetune/llama-2-7b/lora.yaml
+131
View File
@@ -0,0 +1,131 @@
# The path to the base model's checkpoint directory to load for finetuning. (type: <class 'Path'>, default: checkpoints/stabilityai/stablelm-base-alpha-3b)
checkpoint_dir: checkpoints/meta-llama/Llama-2-7b-hf
# Directory in which to save checkpoints and logs. (type: <class 'Path'>, default: out/lora)
out_dir: out/finetune/lora-llama2-7b
# The precision to use for finetuning. Possible choices: "bf16-true", "bf16-mixed", "32-true". (type: Optional[str], default: null)
precision: bf16-true
# If set, quantize the model with this algorithm. See ``tutorials/quantize.md`` for more information. (type: Optional[Literal['nf4', 'nf4-dq', 'fp4', 'fp4-dq', 'int8-training']], default: null)
quantize:
# How many devices/GPUs to use. (type: Union[int, str], default: 1)
devices: 1
# How many nodes to use. (type: int, default: 1)
num_nodes: 1
# The LoRA rank. (type: int, default: 8)
lora_r: 32
# The LoRA alpha. (type: int, default: 16)
lora_alpha: 16
# The LoRA dropout value. (type: float, default: 0.05)
lora_dropout: 0.05
# Whether to apply LoRA to the query weights in attention. (type: bool, default: True)
lora_query: true
# Whether to apply LoRA to the key weights in attention. (type: bool, default: False)
lora_key: false
# Whether to apply LoRA to the value weights in attention. (type: bool, default: True)
lora_value: true
# Whether to apply LoRA to the output projection in the attention block. (type: bool, default: False)
lora_projection: false
# Whether to apply LoRA to the weights of the MLP in the attention block. (type: bool, default: False)
lora_mlp: false
# Whether to apply LoRA to output head in GPT. (type: bool, default: False)
lora_head: false
# Data-related arguments. If not provided, the default is ``litgpt.data.Alpaca``.
data:
class_path: litgpt.data.Alpaca2k
init_args:
mask_prompt: false
prompt_style: alpaca
ignore_index: -100
seed: 42
num_workers: 4
# Training-related arguments. See ``litgpt.args.TrainArgs`` for details
train:
# Number of optimizer steps between saving checkpoints (type: Optional[int], default: 1000)
save_interval: 200
# Number of iterations between logging calls (type: int, default: 1)
log_interval: 1
# Number of samples between optimizer steps across data-parallel ranks (type: int, default: 128)
global_batch_size: 8
# Number of samples per data-parallel rank (type: int, default: 4)
micro_batch_size: 2
# Number of iterations with learning rate warmup active (type: int, default: 100)
lr_warmup_steps: 10
# Number of epochs to train on (type: Optional[int], default: 5)
epochs: 4
# Total number of tokens to train on (type: Optional[int], default: null)
max_tokens:
# Limits the number of optimizer steps to run. (type: Optional[int], default: null)
max_steps:
# Limits the length of samples. Off by default (type: Optional[int], default: null)
max_seq_length: 512
# Whether to tie the embedding weights with the language modeling head weights. (type: Optional[bool], default: null)
tie_embeddings:
# (type: Optional[float], default: null)
max_norm:
# (type: float, default: 6e-05)
min_lr: 6.0e-05
# Evaluation-related arguments. See ``litgpt.args.EvalArgs`` for details
eval:
# Number of optimizer steps between evaluation calls (type: int, default: 100)
interval: 100
# Number of tokens to generate (type: Optional[int], default: 100)
max_new_tokens: 100
# Number of iterations (type: int, default: 100)
max_iters: 100
# Whether to evaluate on the validation set at the beginning of the training
initial_validation: false
# Whether to evaluate on the validation set at the end the training
final_validation: true
# The name of the logger to send metrics to. (type: LoggerChoice, i.e. Literal['wandb', 'tensorboard', 'csv', 'mlflow', 'litlogger'], default: csv)
logger_name: csv
# The random seed to use for reproducibility. (type: int, default: 1337)
seed: 1337
# Optimizer-related arguments
optimizer:
class_path: torch.optim.AdamW
init_args:
# (type: float, default: 0.001)
lr: 0.0002
# (type: float, default: 0.01)
weight_decay: 0.0
# (type: tuple, default: (0.9,0.999))
betas:
- 0.9
- 0.95
config_hub/finetune/llama-2-7b/qlora.yaml
+133
View File
@@ -0,0 +1,133 @@
# The path to the base model's checkpoint directory to load for finetuning. (type: <class 'Path'>, default: checkpoints/stabilityai/stablelm-base-alpha-3b)
checkpoint_dir: checkpoints/meta-llama/Llama-2-7b-hf
# Directory in which to save checkpoints and logs. (type: <class 'Path'>, default: out/lora)
out_dir: out/finetune/qlora-llama2-7b
# The precision to use for finetuning. Possible choices: "bf16-true", "bf16-mixed", "32-true". (type: Optional[str], default: null)
precision: bf16-true
# If set, quantize the model with this algorithm. See ``tutorials/quantize.md`` for more information. (type: Optional[Literal['nf4', 'nf4-dq', 'fp4', 'fp4-dq', 'int8-training']], default: null)
quantize: bnb.nf4
# How many devices/GPUs to use. (type: Union[int, str], default: 1)
devices: 1
# How many nodes to use. (type: int, default: 1)
num_nodes: 1
# The LoRA rank. (type: int, default: 8)
lora_r: 32
# The LoRA alpha. (type: int, default: 16)
lora_alpha: 16
# The LoRA dropout value. (type: float, default: 0.05)
lora_dropout: 0.05
# Whether to apply LoRA to the query weights in attention. (type: bool, default: True)
lora_query: true
# Whether to apply LoRA to the key weights in attention. (type: bool, default: False)
lora_key: false
# Whether to apply LoRA to the value weights in attention. (type: bool, default: True)
lora_value: true
# Whether to apply LoRA to the output projection in the attention block. (type: bool, default: False)
lora_projection: false
# Whether to apply LoRA to the weights of the MLP in the attention block. (type: bool, default: False)
lora_mlp: false
# Whether to apply LoRA to output head in GPT. (type: bool, default: False)
lora_head: false
# Data-related arguments. If not provided, the default is ``litgpt.data.Alpaca``.
data:
class_path: litgpt.data.Alpaca2k
init_args:
mask_prompt: false
val_split_fraction: 0.05
prompt_style: alpaca
ignore_index: -100
seed: 42
num_workers: 4
download_dir: data/alpaca2k
# Training-related arguments. See ``litgpt.args.TrainArgs`` for details
train:
# Number of optimizer steps between saving checkpoints (type: Optional[int], default: 1000)
save_interval: 200
# Number of iterations between logging calls (type: int, default: 1)
log_interval: 1
# Number of samples between optimizer steps across data-parallel ranks (type: int, default: 128)
global_batch_size: 8
# Number of samples per data-parallel rank (type: int, default: 4)
micro_batch_size: 2
# Number of iterations with learning rate warmup active (type: int, default: 100)
lr_warmup_steps: 10
# Number of epochs to train on (type: Optional[int], default: 5)
epochs: 4
# Total number of tokens to train on (type: Optional[int], default: null)
max_tokens:
# Limits the number of optimizer steps to run (type: Optional[int], default: null)
max_steps:
# Limits the length of samples (type: Optional[int], default: null)
max_seq_length: 512
# Whether to tie the embedding weights with the language modeling head weights (type: Optional[bool], default: null)
tie_embeddings:
# (type: Optional[float], default: null)
max_norm:
# (type: float, default: 6e-05)
min_lr: 6.0e-05
# Evaluation-related arguments. See ``litgpt.args.EvalArgs`` for details
eval:
# Number of optimizer steps between evaluation calls (type: int, default: 100)
interval: 100
# Number of tokens to generate (type: Optional[int], default: 100)
max_new_tokens: 100
# Number of iterations (type: int, default: 100)
max_iters: 100
# Whether to evaluate on the validation set at the beginning of the training
initial_validation: false
# Whether to evaluate on the validation set at the end the training
final_validation: true
# The name of the logger to send metrics to. (type: LoggerChoice, i.e. Literal['wandb', 'tensorboard', 'csv', 'mlflow', 'litlogger'], default: csv)
logger_name: csv
# The random seed to use for reproducibility. (type: int, default: 1337)
seed: 1337
# Optimizer-related arguments
optimizer:
class_path: torch.optim.AdamW
init_args:
# (type: float, default: 0.001)
lr: 0.0002
# (type: float, default: 0.01)
weight_decay: 0.0
# (type: tuple, default: (0.9,0.999))
betas:
- 0.9
- 0.95
config_hub/finetune/llama-3-8b/full.yaml
+107
View File
@@ -0,0 +1,107 @@
# The path to the base model's checkpoint directory to load for finetuning. (type: <class 'Path'>, default: checkpoints/stabilityai/stablelm-base-alpha-3b)
checkpoint_dir: checkpoints/meta-llama/Meta-Llama-3-8B
# Directory in which to save checkpoints and logs. (type: <class 'Path'>, default: out/finetune/full)
out_dir: out/finetune/full-llama-3-8b
# The precision to use for finetuning. Possible choices: "bf16-true", "bf16-mixed", "32-true". (type: Optional[str], default: null)
precision: bf16-true
# How many devices/GPUs to use (type: Union[int, str], default: 1)
devices: 4
# How many nodes to use. (type: int, default: 1)
num_nodes: 1
# Path to a checkpoint directory to resume from in case training was interrupted, or ``True`` to resume
# from the latest checkpoint in ``out_dir``. An error will be raised if no checkpoint is found. Passing
# ``'auto'`` will resume from the latest checkpoint but not error if no checkpoint exists.
# (type: Union[bool, Literal["auto"], Path], default: False)
resume: false
# Data-related arguments. If not provided, the default is ``litgpt.data.Alpaca``.
data:
class_path: litgpt.data.Alpaca2k
init_args:
mask_prompt: false
prompt_style: alpaca
ignore_index: -100
seed: 42
num_workers: 4
# Training-related arguments. See ``litgpt.args.TrainArgs`` for details
train:
# Number of optimizer steps between saving checkpoints (type: Optional[int], default: 1000)
save_interval: 200
# Number of iterations between logging calls (type: int, default: 1)
log_interval: 1
# Number of samples between optimizer steps across data-parallel ranks (type: int, default: 64)
global_batch_size: 64
# Number of samples per data-parallel rank (type: int, default: 1)
micro_batch_size: 4
# Number of iterations with learning rate warmup active (type: int, default: 100)
lr_warmup_steps: 25
# Number of epochs to train on (type: Optional[int], default: 5)
epochs: 1
# Total number of tokens to train on (type: Optional[int], default: null)
max_tokens:
# Limits the number of optimizer steps to run. (type: Optional[int], default: null)
max_steps:
# Limits the length of samples. Off by default (type: Optional[int], default: null)
max_seq_length: 512
# Whether to tie the embedding weights with the language modeling head weights. (type: Optional[bool], default: null)
tie_embeddings:
# (type: Optional[float], default: null)
max_norm:
# (type: float, default: 6e-05)
min_lr: 6.0e-05
# Evaluation-related arguments. See ``litgpt.args.EvalArgs`` for details
eval:
# Number of optimizer steps between evaluation calls (type: int, default: 600)
interval: 25
# Number of tokens to generate (type: Optional[int], default: 100)
max_new_tokens: 100
# Number of iterations (type: int, default: 100)
max_iters: 100
# Whether to evaluate on the validation set at the beginning of the training
initial_validation: false
# Whether to evaluate on the validation set at the end the training
final_validation: true
# The name of the logger to send metrics to. (type: LoggerChoice, i.e. Literal['wandb', 'tensorboard', 'csv', 'mlflow', 'litlogger'], default: csv)
logger_name: csv
# The random seed to use for reproducibility. (type: int, default: 1337)
seed: 1337
# Optimizer-related arguments
optimizer:
class_path: torch.optim.AdamW
init_args:
# (type: float, default: 0.001)
lr: 0.0002
# (type: float, default: 0.01)
weight_decay: 0.1
# (type: tuple, default: (0.9,0.999))
betas:
- 0.9
- 0.95
config_hub/finetune/llama-3-8b/lora.yaml
+131
View File
@@ -0,0 +1,131 @@
# The path to the base model's checkpoint directory to load for finetuning. (type: <class 'Path'>, default: checkpoints/stabilityai/stablelm-base-alpha-3b)
checkpoint_dir: checkpoints/meta-llama/Meta-Llama-3-8B
# Directory in which to save checkpoints and logs. (type: <class 'Path'>, default: out/lora)
out_dir: out/finetune/lora-llama-3-8b
# The precision to use for finetuning. Possible choices: "bf16-true", "bf16-mixed", "32-true". (type: Optional[str], default: null)
precision: bf16-true
# If set, quantize the model with this algorithm. See ``tutorials/quantize.md`` for more information. (type: Optional[Literal['nf4', 'nf4-dq', 'fp4', 'fp4-dq', 'int8-training']], default: null)
quantize:
# How many devices/GPUs to use. (type: Union[int, str], default: 1)
devices: 1
# How many nodes to use. (type: int, default: 1)
num_nodes: 1
# The LoRA rank. (type: int, default: 8)
lora_r: 32
# The LoRA alpha. (type: int, default: 16)
lora_alpha: 16
# The LoRA dropout value. (type: float, default: 0.05)
lora_dropout: 0.05
# Whether to apply LoRA to the query weights in attention. (type: bool, default: True)
lora_query: true
# Whether to apply LoRA to the key weights in attention. (type: bool, default: False)
lora_key: false
# Whether to apply LoRA to the value weights in attention. (type: bool, default: True)
lora_value: true
# Whether to apply LoRA to the output projection in the attention block. (type: bool, default: False)
lora_projection: false
# Whether to apply LoRA to the weights of the MLP in the attention block. (type: bool, default: False)
lora_mlp: false
# Whether to apply LoRA to output head in GPT. (type: bool, default: False)
lora_head: false
# Data-related arguments. If not provided, the default is ``litgpt.data.Alpaca``.
data:
class_path: litgpt.data.Alpaca2k
init_args:
mask_prompt: false
prompt_style: alpaca
ignore_index: -100
seed: 42
num_workers: 4
# Training-related arguments. See ``litgpt.args.TrainArgs`` for details
train:
# Number of optimizer steps between saving checkpoints (type: Optional[int], default: 1000)
save_interval: 200
# Number of iterations between logging calls (type: int, default: 1)
log_interval: 1
# Number of samples between optimizer steps across data-parallel ranks (type: int, default: 128)
global_batch_size: 8
# Number of samples per data-parallel rank (type: int, default: 4)
micro_batch_size: 1
# Number of iterations with learning rate warmup active (type: int, default: 100)
lr_warmup_steps: 10
# Number of epochs to train on (type: Optional[int], default: 5)
epochs: 2
# Total number of tokens to train on (type: Optional[int], default: null)
max_tokens:
# Limits the number of optimizer steps to run. (type: Optional[int], default: null)
max_steps:
# Limits the length of samples. Off by default (type: Optional[int], default: null)
max_seq_length: 512
# Whether to tie the embedding weights with the language modeling head weights. (type: Optional[bool], default: null)
tie_embeddings:
# (type: Optional[float], default: null)
max_norm:
# (type: float, default: 6e-05)
min_lr: 6.0e-05
# Evaluation-related arguments. See ``litgpt.args.EvalArgs`` for details
eval:
# Number of optimizer steps between evaluation calls (type: int, default: 100)
interval: 100
# Number of tokens to generate (type: Optional[int], default: 100)
max_new_tokens: 100
# Number of iterations (type: int, default: 100)
max_iters: 100
# Whether to evaluate on the validation set at the beginning of the training
initial_validation: false
# Whether to evaluate on the validation set at the end the training
final_validation: true
# The name of the logger to send metrics to. (type: LoggerChoice, i.e. Literal['wandb', 'tensorboard', 'csv', 'mlflow', 'litlogger'], default: csv)
logger_name: csv
# The random seed to use for reproducibility. (type: int, default: 1337)
seed: 1337
# Optimizer-related arguments
optimizer:
class_path: torch.optim.AdamW
init_args:
# (type: float, default: 0.001)
lr: 0.0002
# (type: float, default: 0.01)
weight_decay: 0.0
# (type: tuple, default: (0.9,0.999))
betas:
- 0.9
- 0.95
config_hub/finetune/llama-3-8b/qlora.yaml
+133
View File
@@ -0,0 +1,133 @@
# The path to the base model's checkpoint directory to load for finetuning. (type: <class 'Path'>, default: checkpoints/stabilityai/stablelm-base-alpha-3b)
checkpoint_dir: checkpoints/meta-llama/Meta-Llama-3-8B
# Directory in which to save checkpoints and logs. (type: <class 'Path'>, default: out/lora)
out_dir: out/finetune/qlora-llama3-8b
# The precision to use for finetuning. Possible choices: "bf16-true", "bf16-mixed", "32-true". (type: Optional[str], default: null)
precision: bf16-true
# If set, quantize the model with this algorithm. See ``tutorials/quantize.md`` for more information. (type: Optional[Literal['nf4', 'nf4-dq', 'fp4', 'fp4-dq', 'int8-training']], default: null)
quantize: bnb.nf4
# How many devices/GPUs to use. (type: Union[int, str], default: 1)
devices: 1
# How many nodes to use. (type: int, default: 1)
num_nodes: 1
# The LoRA rank. (type: int, default: 8)
lora_r: 32
# The LoRA alpha. (type: int, default: 16)
lora_alpha: 16
# The LoRA dropout value. (type: float, default: 0.05)
lora_dropout: 0.05
# Whether to apply LoRA to the query weights in attention. (type: bool, default: True)
lora_query: true
# Whether to apply LoRA to the key weights in attention. (type: bool, default: False)
lora_key: false
# Whether to apply LoRA to the value weights in attention. (type: bool, default: True)
lora_value: true
# Whether to apply LoRA to the output projection in the attention block. (type: bool, default: False)
lora_projection: false
# Whether to apply LoRA to the weights of the MLP in the attention block. (type: bool, default: False)
lora_mlp: false
# Whether to apply LoRA to output head in GPT. (type: bool, default: False)
lora_head: false
# Data-related arguments. If not provided, the default is ``litgpt.data.Alpaca``.
data:
class_path: litgpt.data.Alpaca2k
init_args:
mask_prompt: false
val_split_fraction: 0.05
prompt_style: alpaca
ignore_index: -100
seed: 42
num_workers: 4
download_dir: data/alpaca2k
# Training-related arguments. See ``litgpt.args.TrainArgs`` for details
train:
# Number of optimizer steps between saving checkpoints (type: Optional[int], default: 1000)
save_interval: 200
# Number of iterations between logging calls (type: int, default: 1)
log_interval: 1
# Number of samples between optimizer steps across data-parallel ranks (type: int, default: 128)
global_batch_size: 8
# Number of samples per data-parallel rank (type: int, default: 4)
micro_batch_size: 2
# Number of iterations with learning rate warmup active (type: int, default: 100)
lr_warmup_steps: 10
# Number of epochs to train on (type: Optional[int], default: 5)
epochs: 2
# Total number of tokens to train on (type: Optional[int], default: null)
max_tokens:
# Limits the number of optimizer steps to run (type: Optional[int], default: null)
max_steps:
# Limits the length of samples (type: Optional[int], default: null)
max_seq_length: 512
# Whether to tie the embedding weights with the language modeling head weights (type: Optional[bool], default: null)
tie_embeddings:
# (type: Optional[float], default: null)
max_norm:
# (type: float, default: 6e-05)
min_lr: 6.0e-05
# Evaluation-related arguments. See ``litgpt.args.EvalArgs`` for details
eval:
# Number of optimizer steps between evaluation calls (type: int, default: 100)
interval: 100
# Number of tokens to generate (type: Optional[int], default: 100)
max_new_tokens: 100
# Number of iterations (type: int, default: 100)
max_iters: 100
# Whether to evaluate on the validation set at the beginning of the training
initial_validation: false
# Whether to evaluate on the validation set at the end the training
final_validation: true
# The name of the logger to send metrics to. (type: LoggerChoice, i.e. Literal['wandb', 'tensorboard', 'csv', 'mlflow', 'litlogger'], default: csv)
logger_name: csv
# The random seed to use for reproducibility. (type: int, default: 1337)
seed: 1337
# Optimizer-related arguments
optimizer:
class_path: torch.optim.AdamW
init_args:
# (type: float, default: 0.001)
lr: 0.0002
# (type: float, default: 0.01)
weight_decay: 0.0
# (type: tuple, default: (0.9,0.999))
betas:
- 0.9
- 0.95
config_hub/finetune/llama-3.1-8b/full.yaml
+107
View File
@@ -0,0 +1,107 @@
# The path to the base model's checkpoint directory to load for finetuning. (type: <class 'Path'>, default: checkpoints/stabilityai/stablelm-base-alpha-3b)
checkpoint_dir: checkpoints/meta-llama/Meta-Llama-3.1-8B
# Directory in which to save checkpoints and logs. (type: <class 'Path'>, default: out/finetune/full)
out_dir: out/finetune/full-llama-3.1-8b
# The precision to use for finetuning. Possible choices: "bf16-true", "bf16-mixed", "32-true". (type: Optional[str], default: null)
precision: bf16-true
# How many devices/GPUs to use (type: Union[int, str], default: 1)
devices: 4
# How many nodes to use. (type: int, default: 1)
num_nodes: 1
# Path to a checkpoint directory to resume from in case training was interrupted, or ``True`` to resume
# from the latest checkpoint in ``out_dir``. An error will be raised if no checkpoint is found. Passing
# ``'auto'`` will resume from the latest checkpoint but not error if no checkpoint exists.
# (type: Union[bool, Literal["auto"], Path], default: False)
resume: false
# Data-related arguments. If not provided, the default is ``litgpt.data.Alpaca``.
data:
class_path: litgpt.data.Alpaca2k
init_args:
mask_prompt: false
prompt_style: alpaca
ignore_index: -100
seed: 42
num_workers: 4
# Training-related arguments. See ``litgpt.args.TrainArgs`` for details
train:
# Number of optimizer steps between saving checkpoints (type: Optional[int], default: 1000)
save_interval: 200
# Number of iterations between logging calls (type: int, default: 1)
log_interval: 1
# Number of samples between optimizer steps across data-parallel ranks (type: int, default: 64)
global_batch_size: 64
# Number of samples per data-parallel rank (type: int, default: 1)
micro_batch_size: 4
# Number of iterations with learning rate warmup active (type: int, default: 100)
lr_warmup_steps: 25
# Number of epochs to train on (type: Optional[int], default: 5)
epochs: 1
# Total number of tokens to train on (type: Optional[int], default: null)
max_tokens:
# Limits the number of optimizer steps to run. (type: Optional[int], default: null)
max_steps:
# Limits the length of samples. Off by default (type: Optional[int], default: null)
max_seq_length: 512
# Whether to tie the embedding weights with the language modeling head weights. (type: Optional[bool], default: null)
tie_embeddings:
# (type: Optional[float], default: null)
max_norm:
# (type: float, default: 6e-05)
min_lr: 6.0e-05
# Evaluation-related arguments. See ``litgpt.args.EvalArgs`` for details
eval:
# Number of optimizer steps between evaluation calls (type: int, default: 600)
interval: 25
# Number of tokens to generate (type: Optional[int], default: 100)
max_new_tokens: 100
# Number of iterations (type: int, default: 100)
max_iters: 100
# Whether to evaluate on the validation set at the beginning of the training
initial_validation: false
# Whether to evaluate on the validation set at the end the training
final_validation: true
# The name of the logger to send metrics to. (type: LoggerChoice, i.e. Literal['wandb', 'tensorboard', 'csv', 'mlflow', 'litlogger'], default: csv)
logger_name: csv
# The random seed to use for reproducibility. (type: int, default: 1337)
seed: 1337
# Optimizer-related arguments
optimizer:
class_path: torch.optim.AdamW
init_args:
# (type: float, default: 0.001)
lr: 0.0002
# (type: float, default: 0.01)
weight_decay: 0.1
# (type: tuple, default: (0.9,0.999))
betas:
- 0.9
- 0.95
config_hub/finetune/llama-3.1-8b/lora.yaml
+131
View File
@@ -0,0 +1,131 @@
# The path to the base model's checkpoint directory to load for finetuning. (type: <class 'Path'>, default: checkpoints/stabilityai/stablelm-base-alpha-3b)
checkpoint_dir: checkpoints/meta-llama/Meta-Llama-3.1-8B
# Directory in which to save checkpoints and logs. (type: <class 'Path'>, default: out/lora)
out_dir: out/finetune/lora-llama-3.1-8b
# The precision to use for finetuning. Possible choices: "bf16-true", "bf16-mixed", "32-true". (type: Optional[str], default: null)
precision: bf16-true
# If set, quantize the model with this algorithm. See ``tutorials/quantize.md`` for more information. (type: Optional[Literal['nf4', 'nf4-dq', 'fp4', 'fp4-dq', 'int8-training']], default: null)
quantize:
# How many devices/GPUs to use. (type: Union[int, str], default: 1)
devices: 1
# How many nodes to use. (type: int, default: 1)
num_nodes: 1
# The LoRA rank. (type: int, default: 8)
lora_r: 32
# The LoRA alpha. (type: int, default: 16)
lora_alpha: 16
# The LoRA dropout value. (type: float, default: 0.05)
lora_dropout: 0.05
# Whether to apply LoRA to the query weights in attention. (type: bool, default: True)
lora_query: true
# Whether to apply LoRA to the key weights in attention. (type: bool, default: False)
lora_key: false
# Whether to apply LoRA to the value weights in attention. (type: bool, default: True)
lora_value: true
# Whether to apply LoRA to the output projection in the attention block. (type: bool, default: False)
lora_projection: false
# Whether to apply LoRA to the weights of the MLP in the attention block. (type: bool, default: False)
lora_mlp: false
# Whether to apply LoRA to output head in GPT. (type: bool, default: False)
lora_head: false
# Data-related arguments. If not provided, the default is ``litgpt.data.Alpaca``.
data:
class_path: litgpt.data.Alpaca2k
init_args:
mask_prompt: false
prompt_style: alpaca
ignore_index: -100
seed: 42
num_workers: 4
# Training-related arguments. See ``litgpt.args.TrainArgs`` for details
train:
# Number of optimizer steps between saving checkpoints (type: Optional[int], default: 1000)
save_interval: 200
# Number of iterations between logging calls (type: int, default: 1)
log_interval: 1
# Number of samples between optimizer steps across data-parallel ranks (type: int, default: 128)
global_batch_size: 8
# Number of samples per data-parallel rank (type: int, default: 4)
micro_batch_size: 1
# Number of iterations with learning rate warmup active (type: int, default: 100)
lr_warmup_steps: 10
# Number of epochs to train on (type: Optional[int], default: 5)
epochs: 2
# Total number of tokens to train on (type: Optional[int], default: null)
max_tokens:
# Limits the number of optimizer steps to run. (type: Optional[int], default: null)
max_steps:
# Limits the length of samples. Off by default (type: Optional[int], default: null)
max_seq_length: 512
# Whether to tie the embedding weights with the language modeling head weights. (type: Optional[bool], default: null)
tie_embeddings:
# (type: Optional[float], default: null)
max_norm:
# (type: float, default: 6e-05)
min_lr: 6.0e-05
# Evaluation-related arguments. See ``litgpt.args.EvalArgs`` for details
eval:
# Number of optimizer steps between evaluation calls (type: int, default: 100)
interval: 100
# Number of tokens to generate (type: Optional[int], default: 100)
max_new_tokens: 100
# Number of iterations (type: int, default: 100)
max_iters: 100
# Whether to evaluate on the validation set at the beginning of the training
initial_validation: false
# Whether to evaluate on the validation set at the end the training
final_validation: true
# The name of the logger to send metrics to. (type: LoggerChoice, i.e. Literal['wandb', 'tensorboard', 'csv', 'mlflow', 'litlogger'], default: csv)
logger_name: csv
# The random seed to use for reproducibility. (type: int, default: 1337)
seed: 1337
# Optimizer-related arguments
optimizer:
class_path: torch.optim.AdamW
init_args:
# (type: float, default: 0.001)
lr: 0.0002
# (type: float, default: 0.01)
weight_decay: 0.0
# (type: tuple, default: (0.9,0.999))
betas:
- 0.9
- 0.95
config_hub/finetune/llama-3.1-8b/qlora.yaml
+133
View File
@@ -0,0 +1,133 @@
# The path to the base model's checkpoint directory to load for finetuning. (type: <class 'Path'>, default: checkpoints/stabilityai/stablelm-base-alpha-3b)
checkpoint_dir: checkpoints/meta-llama/Meta-Llama-3.1-8B
# Directory in which to save checkpoints and logs. (type: <class 'Path'>, default: out/lora)
out_dir: out/finetune/qlora-llama3.1-8b
# The precision to use for finetuning. Possible choices: "bf16-true", "bf16-mixed", "32-true". (type: Optional[str], default: null)
precision: bf16-true
# If set, quantize the model with this algorithm. See ``tutorials/quantize.md`` for more information. (type: Optional[Literal['nf4', 'nf4-dq', 'fp4', 'fp4-dq', 'int8-training']], default: null)
quantize: bnb.nf4
# How many devices/GPUs to use. (type: Union[int, str], default: 1)
devices: 1
# How many nodes to use. (type: int, default: 1)
num_nodes: 1
# The LoRA rank. (type: int, default: 8)
lora_r: 32
# The LoRA alpha. (type: int, default: 16)
lora_alpha: 16
# The LoRA dropout value. (type: float, default: 0.05)
lora_dropout: 0.05
# Whether to apply LoRA to the query weights in attention. (type: bool, default: True)
lora_query: true
# Whether to apply LoRA to the key weights in attention. (type: bool, default: False)
lora_key: false
# Whether to apply LoRA to the value weights in attention. (type: bool, default: True)
lora_value: true
# Whether to apply LoRA to the output projection in the attention block. (type: bool, default: False)
lora_projection: false
# Whether to apply LoRA to the weights of the MLP in the attention block. (type: bool, default: False)
lora_mlp: false
# Whether to apply LoRA to output head in GPT. (type: bool, default: False)
lora_head: false
# Data-related arguments. If not provided, the default is ``litgpt.data.Alpaca``.
data:
class_path: litgpt.data.Alpaca2k
init_args:
mask_prompt: false
val_split_fraction: 0.05
prompt_style: alpaca
ignore_index: -100
seed: 42
num_workers: 4
download_dir: data/alpaca2k
# Training-related arguments. See ``litgpt.args.TrainArgs`` for details
train:
# Number of optimizer steps between saving checkpoints (type: Optional[int], default: 1000)
save_interval: 200
# Number of iterations between logging calls (type: int, default: 1)
log_interval: 1
# Number of samples between optimizer steps across data-parallel ranks (type: int, default: 128)
global_batch_size: 8
# Number of samples per data-parallel rank (type: int, default: 4)
micro_batch_size: 2
# Number of iterations with learning rate warmup active (type: int, default: 100)
lr_warmup_steps: 10
# Number of epochs to train on (type: Optional[int], default: 5)
epochs: 2
# Total number of tokens to train on (type: Optional[int], default: null)
max_tokens:
# Limits the number of optimizer steps to run (type: Optional[int], default: null)
max_steps:
# Limits the length of samples (type: Optional[int], default: null)
max_seq_length: 512
# Whether to tie the embedding weights with the language modeling head weights (type: Optional[bool], default: null)
tie_embeddings:
# (type: Optional[float], default: null)
max_norm:
# (type: float, default: 6e-05)
min_lr: 6.0e-05
# Evaluation-related arguments. See ``litgpt.args.EvalArgs`` for details
eval:
# Number of optimizer steps between evaluation calls (type: int, default: 100)
interval: 100
# Number of tokens to generate (type: Optional[int], default: 100)
max_new_tokens: 100
# Number of iterations (type: int, default: 100)
max_iters: 100
# Whether to evaluate on the validation set at the beginning of the training
initial_validation: false
# Whether to evaluate on the validation set at the end the training
final_validation: true
# The name of the logger to send metrics to. (type: LoggerChoice, i.e. Literal['wandb', 'tensorboard', 'csv', 'mlflow', 'litlogger'], default: csv)
logger_name: csv
# The random seed to use for reproducibility. (type: int, default: 1337)
seed: 1337
# Optimizer-related arguments
optimizer:
class_path: torch.optim.AdamW
init_args:
# (type: float, default: 0.001)
lr: 0.0002
# (type: float, default: 0.01)
weight_decay: 0.0
# (type: tuple, default: (0.9,0.999))
betas:
- 0.9
- 0.95
config_hub/finetune/llama-3.2-1B/full.yaml
+107
View File
@@ -0,0 +1,107 @@
# The path to the base model's checkpoint directory to load for finetuning. (type: <class 'Path'>, default: checkpoints/stabilityai/stablelm-base-alpha-3b)
checkpoint_dir: checkpoints/meta-llama/Llama-3.2-1B
# Directory in which to save checkpoints and logs. (type: <class 'Path'>, default: out/finetune/full)
out_dir: out/finetune/full-llama-3.2-1B
# The precision to use for finetuning. Possible choices: "bf16-true", "bf16-mixed", "32-true". (type: Optional[str], default: null)
precision: bf16-true
# How many devices/GPUs to use (type: Union[int, str], default: 1)
devices: 1
# How many nodes to use. (type: int, default: 1)
num_nodes: 1
# Path to a checkpoint directory to resume from in case training was interrupted, or ``True`` to resume
# from the latest checkpoint in ``out_dir``. An error will be raised if no checkpoint is found. Passing
# ``'auto'`` will resume from the latest checkpoint but not error if no checkpoint exists.
# (type: Union[bool, Literal["auto"], Path], default: False)
# resume: false
# Data-related arguments. If not provided, the default is ``litgpt.data.Alpaca``.
data:
class_path: litgpt.data.Alpaca2k
init_args:
mask_prompt: false
prompt_style: alpaca
ignore_index: -100
seed: 42
num_workers: 4
# Training-related arguments. See ``litgpt.args.TrainArgs`` for details
train:
# Number of optimizer steps between saving checkpoints (type: Optional[int], default: 1000)
save_interval: 200
# Number of iterations between logging calls (type: int, default: 1)
log_interval: 1
# Number of samples between optimizer steps across data-parallel ranks (type: int, default: 64)
global_batch_size: 64
# Number of samples per data-parallel rank (type: int, default: 1)
micro_batch_size: 4
# Number of iterations with learning rate warmup active (type: int, default: 100)
lr_warmup_steps: 25
# Number of epochs to train on (type: Optional[int], default: 5)
epochs: 1
# Total number of tokens to train on (type: Optional[int], default: null)
max_tokens:
# Limits the number of optimizer steps to run. (type: Optional[int], default: null)
max_steps:
# Limits the length of samples. Off by default (type: Optional[int], default: null)
max_seq_length: 512
# Whether to tie the embedding weights with the language modeling head weights. (type: Optional[bool], default: null)
tie_embeddings:
# (type: Optional[float], default: null)
max_norm:
# (type: float, default: 6e-05)
min_lr: 6.0e-05
# Evaluation-related arguments. See ``litgpt.args.EvalArgs`` for details
eval:
# Number of optimizer steps between evaluation calls (type: int, default: 600)
interval: 25
# Number of tokens to generate (type: Optional[int], default: 100)
max_new_tokens: 100
# Number of iterations (type: int, default: 100)
max_iters: 100
# Whether to evaluate on the validation set at the beginning of the training
initial_validation: false
# Whether to evaluate on the validation set at the end the training
final_validation: true
# The name of the logger to send metrics to. (type: LoggerChoice, i.e. Literal['wandb', 'tensorboard', 'csv', 'mlflow', 'litlogger'], default: csv)
logger_name: csv
# The random seed to use for reproducibility. (type: int, default: 1337)
seed: 1337
# Optimizer-related arguments
optimizer:
class_path: torch.optim.AdamW
init_args:
# (type: float, default: 0.001)
lr: 0.0002
# (type: float, default: 0.01)
weight_decay: 0.1
# (type: tuple, default: (0.9,0.999))
betas:
- 0.9
- 0.95
config_hub/finetune/llama-3.2-1B/lora.yaml
+131
View File
@@ -0,0 +1,131 @@
# The path to the base model's checkpoint directory to load for finetuning. (type: <class 'Path'>, default: checkpoints/stabilityai/stablelm-base-alpha-3b)
checkpoint_dir: checkpoints/meta-llama/Llama-3.2-1B
# Directory in which to save checkpoints and logs. (type: <class 'Path'>, default: out/lora)
out_dir: out/finetune/lora-llama-3.2-1B
# The precision to use for finetuning. Possible choices: "bf16-true", "bf16-mixed", "32-true". (type: Optional[str], default: null)
precision: bf16-true
# If set, quantize the model with this algorithm. See ``tutorials/quantize.md`` for more information. (type: Optional[Literal['nf4', 'nf4-dq', 'fp4', 'fp4-dq', 'int8-training']], default: null)
quantize:
# How many devices/GPUs to use. (type: Union[int, str], default: 1)
devices: 1
# How many nodes to use. (type: int, default: 1)
num_nodes: 1
# The LoRA rank. (type: int, default: 8)
lora_r: 32
# The LoRA alpha. (type: int, default: 16)
lora_alpha: 16
# The LoRA dropout value. (type: float, default: 0.05)
lora_dropout: 0.05
# Whether to apply LoRA to the query weights in attention. (type: bool, default: True)
lora_query: true
# Whether to apply LoRA to the key weights in attention. (type: bool, default: False)
lora_key: false
# Whether to apply LoRA to the value weights in attention. (type: bool, default: True)
lora_value: true
# Whether to apply LoRA to the output projection in the attention block. (type: bool, default: False)
lora_projection: false
# Whether to apply LoRA to the weights of the MLP in the attention block. (type: bool, default: False)
lora_mlp: false
# Whether to apply LoRA to output head in GPT. (type: bool, default: False)
lora_head: false
# Data-related arguments. If not provided, the default is ``litgpt.data.Alpaca``.
data:
class_path: litgpt.data.Alpaca2k
init_args:
mask_prompt: false
prompt_style: alpaca
ignore_index: -100
seed: 42
num_workers: 4
# Training-related arguments. See ``litgpt.args.TrainArgs`` for details
train:
# Number of optimizer steps between saving checkpoints (type: Optional[int], default: 1000)
save_interval: 200
# Number of iterations between logging calls (type: int, default: 1)
log_interval: 1
# Number of samples between optimizer steps across data-parallel ranks (type: int, default: 128)
global_batch_size: 8
# Number of samples per data-parallel rank (type: int, default: 4)
micro_batch_size: 1
# Number of iterations with learning rate warmup active (type: int, default: 100)
lr_warmup_steps: 10
# Number of epochs to train on (type: Optional[int], default: 5)
epochs: 2
# Total number of tokens to train on (type: Optional[int], default: null)
max_tokens:
# Limits the number of optimizer steps to run. (type: Optional[int], default: null)
max_steps:
# Limits the length of samples. Off by default (type: Optional[int], default: null)
max_seq_length: 512
# Whether to tie the embedding weights with the language modeling head weights. (type: Optional[bool], default: null)
tie_embeddings:
# (type: Optional[float], default: null)
max_norm:
# (type: float, default: 6e-05)
min_lr: 6.0e-05
# Evaluation-related arguments. See ``litgpt.args.EvalArgs`` for details
eval:
# Number of optimizer steps between evaluation calls (type: int, default: 100)
interval: 100
# Number of tokens to generate (type: Optional[int], default: 100)
max_new_tokens: 100
# Number of iterations (type: int, default: 100)
max_iters: 100
# Whether to evaluate on the validation set at the beginning of the training
initial_validation: false
# Whether to evaluate on the validation set at the end the training
final_validation: true
# The name of the logger to send metrics to. (type: LoggerChoice, i.e. Literal['wandb', 'tensorboard', 'csv', 'mlflow', 'litlogger'], default: csv)
logger_name: csv
# The random seed to use for reproducibility. (type: int, default: 1337)
seed: 1337
# Optimizer-related arguments
optimizer:
class_path: torch.optim.AdamW
init_args:
# (type: float, default: 0.001)
lr: 0.0002
# (type: float, default: 0.01)
weight_decay: 0.0
# (type: tuple, default: (0.9,0.999))
betas:
- 0.9
- 0.95
config_hub/finetune/llama-3.2-1B/qlora.yaml
+133
View File
@@ -0,0 +1,133 @@
# The path to the base model's checkpoint directory to load for finetuning. (type: <class 'Path'>, default: checkpoints/stabilityai/stablelm-base-alpha-3b)
checkpoint_dir: checkpoints/meta-llama/Llama-3.2-1B
# Directory in which to save checkpoints and logs. (type: <class 'Path'>, default: out/lora)
out_dir: out/finetune/qlora-llama3.2-1b
# The precision to use for finetuning. Possible choices: "bf16-true", "bf16-mixed", "32-true". (type: Optional[str], default: null)
precision: bf16-true
# If set, quantize the model with this algorithm. See ``tutorials/quantize.md`` for more information. (type: Optional[Literal['nf4', 'nf4-dq', 'fp4', 'fp4-dq', 'int8-training']], default: null)
quantize: bnb.nf4
# How many devices/GPUs to use. (type: Union[int, str], default: 1)
devices: 1
# How many nodes to use. (type: int, default: 1)
num_nodes: 1
# The LoRA rank. (type: int, default: 8)
lora_r: 32
# The LoRA alpha. (type: int, default: 16)
lora_alpha: 16
# The LoRA dropout value. (type: float, default: 0.05)
lora_dropout: 0.05
# Whether to apply LoRA to the query weights in attention. (type: bool, default: True)
lora_query: true
# Whether to apply LoRA to the key weights in attention. (type: bool, default: False)
lora_key: false
# Whether to apply LoRA to the value weights in attention. (type: bool, default: True)
lora_value: true
# Whether to apply LoRA to the output projection in the attention block. (type: bool, default: False)
lora_projection: false
# Whether to apply LoRA to the weights of the MLP in the attention block. (type: bool, default: False)
lora_mlp: false
# Whether to apply LoRA to output head in GPT. (type: bool, default: False)
lora_head: false
# Data-related arguments. If not provided, the default is ``litgpt.data.Alpaca``.
data:
class_path: litgpt.data.Alpaca2k
init_args:
mask_prompt: false
val_split_fraction: 0.05
prompt_style: alpaca
ignore_index: -100
seed: 42
num_workers: 4
download_dir: data/alpaca2k
# Training-related arguments. See ``litgpt.args.TrainArgs`` for details
train:
# Number of optimizer steps between saving checkpoints (type: Optional[int], default: 1000)
save_interval: 200
# Number of iterations between logging calls (type: int, default: 1)
log_interval: 1
# Number of samples between optimizer steps across data-parallel ranks (type: int, default: 128)
global_batch_size: 8
# Number of samples per data-parallel rank (type: int, default: 4)
micro_batch_size: 2
# Number of iterations with learning rate warmup active (type: int, default: 100)
lr_warmup_steps: 10
# Number of epochs to train on (type: Optional[int], default: 5)
epochs: 2
# Total number of tokens to train on (type: Optional[int], default: null)
max_tokens:
# Limits the number of optimizer steps to run (type: Optional[int], default: null)
max_steps:
# Limits the length of samples (type: Optional[int], default: null)
max_seq_length: 512
# Whether to tie the embedding weights with the language modeling head weights (type: Optional[bool], default: null)
tie_embeddings:
# (type: Optional[float], default: null)
max_norm:
# (type: float, default: 6e-05)
min_lr: 6.0e-05
# Evaluation-related arguments. See ``litgpt.args.EvalArgs`` for details
eval:
# Number of optimizer steps between evaluation calls (type: int, default: 100)
interval: 100
# Number of tokens to generate (type: Optional[int], default: 100)
max_new_tokens: 100
# Number of iterations (type: int, default: 100)
max_iters: 100
# Whether to evaluate on the validation set at the beginning of the training
initial_validation: false
# Whether to evaluate on the validation set at the end the training
final_validation: true
# The name of the logger to send metrics to. (type: LoggerChoice, i.e. Literal['wandb', 'tensorboard', 'csv', 'mlflow', 'litlogger'], default: csv)
logger_name: csv
# The random seed to use for reproducibility. (type: int, default: 1337)
seed: 1337
# Optimizer-related arguments
optimizer:
class_path: torch.optim.AdamW
init_args:
# (type: float, default: 0.001)
lr: 0.0002
# (type: float, default: 0.01)
weight_decay: 0.0
# (type: tuple, default: (0.9,0.999))
betas:
- 0.9
- 0.95
config_hub/finetune/llama-3.2-3B/full.yaml
+107
View File
@@ -0,0 +1,107 @@
# The path to the base model's checkpoint directory to load for finetuning. (type: <class 'Path'>, default: checkpoints/stabilityai/stablelm-base-alpha-3b)
checkpoint_dir: checkpoints/meta-llama/Llama-3.2-3B
# Directory in which to save checkpoints and logs. (type: <class 'Path'>, default: out/finetune/full)
out_dir: out/finetune/full-llama-3.2-3B
# The precision to use for finetuning. Possible choices: "bf16-true", "bf16-mixed", "32-true". (type: Optional[str], default: null)
precision: bf16-true
# How many devices/GPUs to use (type: Union[int, str], default: 1)
devices: 1
# How many nodes to use. (type: int, default: 1)
num_nodes: 1
# Path to a checkpoint directory to resume from in case training was interrupted, or ``True`` to resume
# from the latest checkpoint in ``out_dir``. An error will be raised if no checkpoint is found. Passing
# ``'auto'`` will resume from the latest checkpoint but not error if no checkpoint exists.
# (type: Union[bool, Literal["auto"], Path], default: False)
# resume: false
# Data-related arguments. If not provided, the default is ``litgpt.data.Alpaca``.
data:
class_path: litgpt.data.Alpaca2k
init_args:
mask_prompt: false
prompt_style: alpaca
ignore_index: -100
seed: 42
num_workers: 4
# Training-related arguments. See ``litgpt.args.TrainArgs`` for details
train:
# Number of optimizer steps between saving checkpoints (type: Optional[int], default: 1000)
save_interval: 200
# Number of iterations between logging calls (type: int, default: 1)
log_interval: 1
# Number of samples between optimizer steps across data-parallel ranks (type: int, default: 64)
global_batch_size: 64
# Number of samples per data-parallel rank (type: int, default: 1)
micro_batch_size: 4
# Number of iterations with learning rate warmup active (type: int, default: 100)
lr_warmup_steps: 25
# Number of epochs to train on (type: Optional[int], default: 5)
epochs: 1
# Total number of tokens to train on (type: Optional[int], default: null)
max_tokens:
# Limits the number of optimizer steps to run. (type: Optional[int], default: null)
max_steps:
# Limits the length of samples. Off by default (type: Optional[int], default: null)
max_seq_length: 512
# Whether to tie the embedding weights with the language modeling head weights. (type: Optional[bool], default: null)
tie_embeddings:
# (type: Optional[float], default: null)
max_norm:
# (type: float, default: 6e-05)
min_lr: 6.0e-05
# Evaluation-related arguments. See ``litgpt.args.EvalArgs`` for details
eval:
# Number of optimizer steps between evaluation calls (type: int, default: 600)
interval: 25
# Number of tokens to generate (type: Optional[int], default: 100)
max_new_tokens: 100
# Number of iterations (type: int, default: 100)
max_iters: 100
# Whether to evaluate on the validation set at the beginning of the training
initial_validation: false
# Whether to evaluate on the validation set at the end the training
final_validation: true
# The name of the logger to send metrics to. (type: LoggerChoice, i.e. Literal['wandb', 'tensorboard', 'csv', 'mlflow', 'litlogger'], default: csv)
logger_name: csv
# The random seed to use for reproducibility. (type: int, default: 1337)
seed: 1337
# Optimizer-related arguments
optimizer:
class_path: torch.optim.AdamW
init_args:
# (type: float, default: 0.001)
lr: 0.0002
# (type: float, default: 0.01)
weight_decay: 0.1
# (type: tuple, default: (0.9,0.999))
betas:
- 0.9
- 0.95
config_hub/finetune/llama-3.2-3B/lora.yaml
+131
View File
@@ -0,0 +1,131 @@
# The path to the base model's checkpoint directory to load for finetuning. (type: <class 'Path'>, default: checkpoints/stabilityai/stablelm-base-alpha-3b)
checkpoint_dir: checkpoints/meta-llama/Llama-3.2-3B
# Directory in which to save checkpoints and logs. (type: <class 'Path'>, default: out/lora)
out_dir: out/finetune/lora-llama-3.2-3B
# The precision to use for finetuning. Possible choices: "bf16-true", "bf16-mixed", "32-true". (type: Optional[str], default: null)
precision: bf16-true
# If set, quantize the model with this algorithm. See ``tutorials/quantize.md`` for more information. (type: Optional[Literal['nf4', 'nf4-dq', 'fp4', 'fp4-dq', 'int8-training']], default: null)
quantize:
# How many devices/GPUs to use. (type: Union[int, str], default: 1)
devices: 1
# How many nodes to use. (type: int, default: 1)
num_nodes: 1
# The LoRA rank. (type: int, default: 8)
lora_r: 32
# The LoRA alpha. (type: int, default: 16)
lora_alpha: 16
# The LoRA dropout value. (type: float, default: 0.05)
lora_dropout: 0.05
# Whether to apply LoRA to the query weights in attention. (type: bool, default: True)
lora_query: true
# Whether to apply LoRA to the key weights in attention. (type: bool, default: False)
lora_key: false
# Whether to apply LoRA to the value weights in attention. (type: bool, default: True)
lora_value: true
# Whether to apply LoRA to the output projection in the attention block. (type: bool, default: False)
lora_projection: false
# Whether to apply LoRA to the weights of the MLP in the attention block. (type: bool, default: False)
lora_mlp: false
# Whether to apply LoRA to output head in GPT. (type: bool, default: False)
lora_head: false
# Data-related arguments. If not provided, the default is ``litgpt.data.Alpaca``.
data:
class_path: litgpt.data.Alpaca2k
init_args:
mask_prompt: false
prompt_style: alpaca
ignore_index: -100
seed: 42
num_workers: 4
# Training-related arguments. See ``litgpt.args.TrainArgs`` for details
train:
# Number of optimizer steps between saving checkpoints (type: Optional[int], default: 1000)
save_interval: 200
# Number of iterations between logging calls (type: int, default: 1)
log_interval: 1
# Number of samples between optimizer steps across data-parallel ranks (type: int, default: 128)
global_batch_size: 8
# Number of samples per data-parallel rank (type: int, default: 4)
micro_batch_size: 1
# Number of iterations with learning rate warmup active (type: int, default: 100)
lr_warmup_steps: 10
# Number of epochs to train on (type: Optional[int], default: 5)
epochs: 2
# Total number of tokens to train on (type: Optional[int], default: null)
max_tokens:
# Limits the number of optimizer steps to run. (type: Optional[int], default: null)
max_steps:
# Limits the length of samples. Off by default (type: Optional[int], default: null)
max_seq_length: 512
# Whether to tie the embedding weights with the language modeling head weights. (type: Optional[bool], default: null)
tie_embeddings:
# (type: Optional[float], default: null)
max_norm:
# (type: float, default: 6e-05)
min_lr: 6.0e-05
# Evaluation-related arguments. See ``litgpt.args.EvalArgs`` for details
eval:
# Number of optimizer steps between evaluation calls (type: int, default: 100)
interval: 100
# Number of tokens to generate (type: Optional[int], default: 100)
max_new_tokens: 100
# Number of iterations (type: int, default: 100)
max_iters: 100
# Whether to evaluate on the validation set at the beginning of the training
initial_validation: false
# Whether to evaluate on the validation set at the end the training
final_validation: true
# The name of the logger to send metrics to. (type: LoggerChoice, i.e. Literal['wandb', 'tensorboard', 'csv', 'mlflow', 'litlogger'], default: csv)
logger_name: csv
# The random seed to use for reproducibility. (type: int, default: 1337)
seed: 1337
# Optimizer-related arguments
optimizer:
class_path: torch.optim.AdamW
init_args:
# (type: float, default: 0.001)
lr: 0.0002
# (type: float, default: 0.01)
weight_decay: 0.0
# (type: tuple, default: (0.9,0.999))
betas:
- 0.9
- 0.95
config_hub/finetune/llama-3.2-3B/qlora.yaml
+133
View File
@@ -0,0 +1,133 @@
# The path to the base model's checkpoint directory to load for finetuning. (type: <class 'Path'>, default: checkpoints/stabilityai/stablelm-base-alpha-3b)
checkpoint_dir: checkpoints/meta-llama/Llama-3.2-3B
# Directory in which to save checkpoints and logs. (type: <class 'Path'>, default: out/lora)
out_dir: out/finetune/qlora-llama3.2-3b
# The precision to use for finetuning. Possible choices: "bf16-true", "bf16-mixed", "32-true". (type: Optional[str], default: null)
precision: bf16-true
# If set, quantize the model with this algorithm. See ``tutorials/quantize.md`` for more information. (type: Optional[Literal['nf4', 'nf4-dq', 'fp4', 'fp4-dq', 'int8-training']], default: null)
quantize: bnb.nf4
# How many devices/GPUs to use. (type: Union[int, str], default: 1)
devices: 1
# How many nodes to use. (type: int, default: 1)
num_nodes: 1
# The LoRA rank. (type: int, default: 8)
lora_r: 32
# The LoRA alpha. (type: int, default: 16)
lora_alpha: 16
# The LoRA dropout value. (type: float, default: 0.05)
lora_dropout: 0.05
# Whether to apply LoRA to the query weights in attention. (type: bool, default: True)
lora_query: true
# Whether to apply LoRA to the key weights in attention. (type: bool, default: False)
lora_key: false
# Whether to apply LoRA to the value weights in attention. (type: bool, default: True)
lora_value: true
# Whether to apply LoRA to the output projection in the attention block. (type: bool, default: False)
lora_projection: false
# Whether to apply LoRA to the weights of the MLP in the attention block. (type: bool, default: False)
lora_mlp: false
# Whether to apply LoRA to output head in GPT. (type: bool, default: False)
lora_head: false
# Data-related arguments. If not provided, the default is ``litgpt.data.Alpaca``.
data:
class_path: litgpt.data.Alpaca2k
init_args:
mask_prompt: false
val_split_fraction: 0.05
prompt_style: alpaca
ignore_index: -100
seed: 42
num_workers: 4
download_dir: data/alpaca2k
# Training-related arguments. See ``litgpt.args.TrainArgs`` for details
train:
# Number of optimizer steps between saving checkpoints (type: Optional[int], default: 1000)
save_interval: 200
# Number of iterations between logging calls (type: int, default: 1)
log_interval: 1
# Number of samples between optimizer steps across data-parallel ranks (type: int, default: 128)
global_batch_size: 8
# Number of samples per data-parallel rank (type: int, default: 4)
micro_batch_size: 2
# Number of iterations with learning rate warmup active (type: int, default: 100)
lr_warmup_steps: 10
# Number of epochs to train on (type: Optional[int], default: 5)
epochs: 2
# Total number of tokens to train on (type: Optional[int], default: null)
max_tokens:
# Limits the number of optimizer steps to run (type: Optional[int], default: null)
max_steps:
# Limits the length of samples (type: Optional[int], default: null)
max_seq_length: 512
# Whether to tie the embedding weights with the language modeling head weights (type: Optional[bool], default: null)
tie_embeddings:
# (type: Optional[float], default: null)
max_norm:
# (type: float, default: 6e-05)
min_lr: 6.0e-05
# Evaluation-related arguments. See ``litgpt.args.EvalArgs`` for details
eval:
# Number of optimizer steps between evaluation calls (type: int, default: 100)
interval: 100
# Number of tokens to generate (type: Optional[int], default: 100)
max_new_tokens: 100
# Number of iterations (type: int, default: 100)
max_iters: 100
# Whether to evaluate on the validation set at the beginning of the training
initial_validation: false
# Whether to evaluate on the validation set at the end the training
final_validation: true
# The name of the logger to send metrics to. (type: LoggerChoice, i.e. Literal['wandb', 'tensorboard', 'csv', 'mlflow', 'litlogger'], default: csv)
logger_name: csv
# The random seed to use for reproducibility. (type: int, default: 1337)
seed: 1337
# Optimizer-related arguments
optimizer:
class_path: torch.optim.AdamW
init_args:
# (type: float, default: 0.001)
lr: 0.0002
# (type: float, default: 0.01)
weight_decay: 0.0
# (type: tuple, default: (0.9,0.999))
betas:
- 0.9
- 0.95
config_hub/finetune/mistral-7b-v0.2/lora.yaml
+131
View File
@@ -0,0 +1,131 @@
# The path to the base model's checkpoint directory to load for finetuning. (type: <class 'Path'>, default: checkpoints/stabilityai/stablelm-base-alpha-3b)
checkpoint_dir: checkpoints/unsloth/Mistral-7B-v0.2
# Directory in which to save checkpoints and logs. (type: <class 'Path'>, default: out/lora)
out_dir: out/finetune/lora-mistral-7b
# The precision to use for finetuning. Possible choices: "bf16-true", "bf16-mixed", "32-true". (type: Optional[str], default: null)
precision: bf16-true
# If set, quantize the model with this algorithm. See ``tutorials/quantize.md`` for more information. (type: Optional[Literal['nf4', 'nf4-dq', 'fp4', 'fp4-dq', 'int8-training']], default: null)
quantize:
# How many devices/GPUs to use. (type: Union[int, str], default: 1)
devices: 1
# How many nodes to use. (type: int, default: 1)
num_nodes: 1
# The LoRA rank. (type: int, default: 8)
lora_r: 32
# The LoRA alpha. (type: int, default: 16)
lora_alpha: 16
# The LoRA dropout value. (type: float, default: 0.05)
lora_dropout: 0.05
# Whether to apply LoRA to the query weights in attention. (type: bool, default: True)
lora_query: true
# Whether to apply LoRA to the key weights in attention. (type: bool, default: False)
lora_key: false
# Whether to apply LoRA to the value weights in attention. (type: bool, default: True)
lora_value: true
# Whether to apply LoRA to the output projection in the attention block. (type: bool, default: False)
lora_projection: false
# Whether to apply LoRA to the weights of the MLP in the attention block. (type: bool, default: False)
lora_mlp: false
# Whether to apply LoRA to output head in GPT. (type: bool, default: False)
lora_head: false
# Data-related arguments. If not provided, the default is ``litgpt.data.Alpaca``.
data:
class_path: litgpt.data.Alpaca2k
init_args:
mask_prompt: false
prompt_style: alpaca
ignore_index: -100
seed: 42
num_workers: 4
# Training-related arguments. See ``litgpt.args.TrainArgs`` for details
train:
# Number of optimizer steps between saving checkpoints (type: Optional[int], default: 1000)
save_interval: 200
# Number of iterations between logging calls (type: int, default: 1)
log_interval: 1
# Number of samples between optimizer steps across data-parallel ranks (type: int, default: 128)
global_batch_size: 8
# Number of samples per data-parallel rank (type: int, default: 4)
micro_batch_size: 2
# Number of iterations with learning rate warmup active (type: int, default: 100)
lr_warmup_steps: 10
# Number of epochs to train on (type: Optional[int], default: 5)
epochs: 4
# Total number of tokens to train on (type: Optional[int], default: null)
max_tokens:
# Limits the number of optimizer steps to run. (type: Optional[int], default: null)
max_steps:
# Limits the length of samples. Off by default (type: Optional[int], default: null)
max_seq_length: 512
# Whether to tie the embedding weights with the language modeling head weights. (type: Optional[bool], default: null)
tie_embeddings:
# (type: Optional[float], default: null)
max_norm:
# (type: float, default: 6e-05)
min_lr: 6.0e-05
# Evaluation-related arguments. See ``litgpt.args.EvalArgs`` for details
eval:
# Number of optimizer steps between evaluation calls (type: int, default: 100)
interval: 100
# Number of tokens to generate (type: Optional[int], default: 100)
max_new_tokens: 100
# Number of iterations (type: int, default: 100)
max_iters: 100
# Whether to evaluate on the validation set at the beginning of the training
initial_validation: false
# Whether to evaluate on the validation set at the end the training
final_validation: true
# The name of the logger to send metrics to. (type: LoggerChoice, i.e. Literal['wandb', 'tensorboard', 'csv', 'mlflow', 'litlogger'], default: csv)
logger_name: csv
# The random seed to use for reproducibility. (type: int, default: 1337)
seed: 1337
# Optimizer-related arguments
optimizer:
class_path: torch.optim.AdamW
init_args:
# (type: float, default: 0.001)
lr: 0.0002
# (type: float, default: 0.01)
weight_decay: 0.0
# (type: tuple, default: (0.9,0.999))
betas:
- 0.9
- 0.95
config_hub/finetune/mistral-7b-v0.2/qlora.yaml
+133
View File
@@ -0,0 +1,133 @@
# The path to the base model's checkpoint directory to load for finetuning. (type: <class 'Path'>, default: checkpoints/stabilityai/stablelm-base-alpha-3b)
checkpoint_dir: checkpoints/unsloth/Mistral-7B-v0.2
# Directory in which to save checkpoints and logs. (type: <class 'Path'>, default: out/lora)
out_dir: out/finetune/qlora-mistral-7b
# The precision to use for finetuning. Possible choices: "bf16-true", "bf16-mixed", "32-true". (type: Optional[str], default: null)
precision: bf16-true
# If set, quantize the model with this algorithm. See ``tutorials/quantize.md`` for more information. (type: Optional[Literal['nf4', 'nf4-dq', 'fp4', 'fp4-dq', 'int8-training']], default: null)
quantize: bnb.nf4
# How many devices/GPUs to use. (type: Union[int, str], default: 1)
devices: 1
# How many nodes to use. (type: int, default: 1)
num_nodes: 1
# The LoRA rank. (type: int, default: 8)
lora_r: 32
# The LoRA alpha. (type: int, default: 16)
lora_alpha: 16
# The LoRA dropout value. (type: float, default: 0.05)
lora_dropout: 0.05
# Whether to apply LoRA to the query weights in attention. (type: bool, default: True)
lora_query: true
# Whether to apply LoRA to the key weights in attention. (type: bool, default: False)
lora_key: false
# Whether to apply LoRA to the value weights in attention. (type: bool, default: True)
lora_value: true
# Whether to apply LoRA to the output projection in the attention block. (type: bool, default: False)
lora_projection: false
# Whether to apply LoRA to the weights of the MLP in the attention block. (type: bool, default: False)
lora_mlp: false
# Whether to apply LoRA to output head in GPT. (type: bool, default: False)
lora_head: false
# Data-related arguments. If not provided, the default is ``litgpt.data.Alpaca``.
data:
class_path: litgpt.data.Alpaca2k
init_args:
mask_prompt: false
val_split_fraction: 0.05
prompt_style: alpaca
ignore_index: -100
seed: 42
num_workers: 4
download_dir: data/alpaca2k
# Training-related arguments. See ``litgpt.args.TrainArgs`` for details
train:
# Number of optimizer steps between saving checkpoints (type: Optional[int], default: 1000)
save_interval: 200
# Number of iterations between logging calls (type: int, default: 1)
log_interval: 1
# Number of samples between optimizer steps across data-parallel ranks (type: int, default: 128)
global_batch_size: 8
# Number of samples per data-parallel rank (type: int, default: 4)
micro_batch_size: 2
# Number of iterations with learning rate warmup active (type: int, default: 100)
lr_warmup_steps: 10
# Number of epochs to train on (type: Optional[int], default: 5)
epochs: 4
# Total number of tokens to train on (type: Optional[int], default: null)
max_tokens:
# Limits the number of optimizer steps to run (type: Optional[int], default: null)
max_steps:
# Limits the length of samples (type: Optional[int], default: null)
max_seq_length: 512
# Whether to tie the embedding weights with the language modeling head weights (type: Optional[bool], default: null)
tie_embeddings:
# (type: Optional[float], default: null)
max_norm:
# (type: float, default: 6e-05)
min_lr: 6.0e-05
# Evaluation-related arguments. See ``litgpt.args.EvalArgs`` for details
eval:
# Number of optimizer steps between evaluation calls (type: int, default: 100)
interval: 100
# Number of tokens to generate (type: Optional[int], default: 100)
max_new_tokens: 100
# Number of iterations (type: int, default: 100)
max_iters: 100
# Whether to evaluate on the validation set at the beginning of the training
initial_validation: false
# Whether to evaluate on the validation set at the end the training
final_validation: true
# The name of the logger to send metrics to. (type: LoggerChoice, i.e. Literal['wandb', 'tensorboard', 'csv', 'mlflow', 'litlogger'], default: csv)
logger_name: csv
# The random seed to use for reproducibility. (type: int, default: 1337)
seed: 1337
# Optimizer-related arguments
optimizer:
class_path: torch.optim.AdamW
init_args:
# (type: float, default: 0.001)
lr: 0.0002
# (type: float, default: 0.01)
weight_decay: 0.0
# (type: tuple, default: (0.9,0.999))
betas:
- 0.9
- 0.95
config_hub/finetune/mistral-7b/lora.yaml
+131
View File
@@ -0,0 +1,131 @@
# The path to the base model's checkpoint directory to load for finetuning. (type: <class 'Path'>, default: checkpoints/stabilityai/stablelm-base-alpha-3b)
checkpoint_dir: checkpoints/mistralai/Mistral-7B-v0.1
# Directory in which to save checkpoints and logs. (type: <class 'Path'>, default: out/lora)
out_dir: out/finetune/lora-mistral-7b
# The precision to use for finetuning. Possible choices: "bf16-true", "bf16-mixed", "32-true". (type: Optional[str], default: null)
precision: bf16-true
# If set, quantize the model with this algorithm. See ``tutorials/quantize.md`` for more information. (type: Optional[Literal['nf4', 'nf4-dq', 'fp4', 'fp4-dq', 'int8-training']], default: null)
quantize:
# How many devices/GPUs to use. (type: Union[int, str], default: 1)
devices: 1
# How many nodes to use. (type: int, default: 1)
num_nodes: 1
# The LoRA rank. (type: int, default: 8)
lora_r: 32
# The LoRA alpha. (type: int, default: 16)
lora_alpha: 16
# The LoRA dropout value. (type: float, default: 0.05)
lora_dropout: 0.05
# Whether to apply LoRA to the query weights in attention. (type: bool, default: True)
lora_query: true
# Whether to apply LoRA to the key weights in attention. (type: bool, default: False)
lora_key: false
# Whether to apply LoRA to the value weights in attention. (type: bool, default: True)
lora_value: true
# Whether to apply LoRA to the output projection in the attention block. (type: bool, default: False)
lora_projection: false
# Whether to apply LoRA to the weights of the MLP in the attention block. (type: bool, default: False)
lora_mlp: false
# Whether to apply LoRA to output head in GPT. (type: bool, default: False)
lora_head: false
# Data-related arguments. If not provided, the default is ``litgpt.data.Alpaca``.
data:
class_path: litgpt.data.Alpaca2k
init_args:
mask_prompt: false
prompt_style: alpaca
ignore_index: -100
seed: 42
num_workers: 4
# Training-related arguments. See ``litgpt.args.TrainArgs`` for details
train:
# Number of optimizer steps between saving checkpoints (type: Optional[int], default: 1000)
save_interval: 200
# Number of iterations between logging calls (type: int, default: 1)
log_interval: 1
# Number of samples between optimizer steps across data-parallel ranks (type: int, default: 128)
global_batch_size: 8
# Number of samples per data-parallel rank (type: int, default: 4)
micro_batch_size: 2
# Number of iterations with learning rate warmup active (type: int, default: 100)
lr_warmup_steps: 10
# Number of epochs to train on (type: Optional[int], default: 5)
epochs: 4
# Total number of tokens to train on (type: Optional[int], default: null)
max_tokens:
# Limits the number of optimizer steps to run. (type: Optional[int], default: null)
max_steps:
# Limits the length of samples. Off by default (type: Optional[int], default: null)
max_seq_length: 512
# Whether to tie the embedding weights with the language modeling head weights. (type: Optional[bool], default: null)
tie_embeddings:
# (type: Optional[float], default: null)
max_norm:
# (type: float, default: 6e-05)
min_lr: 6.0e-05
# Evaluation-related arguments. See ``litgpt.args.EvalArgs`` for details
eval:
# Number of optimizer steps between evaluation calls (type: int, default: 100)
interval: 100
# Number of tokens to generate (type: Optional[int], default: 100)
max_new_tokens: 100
# Number of iterations (type: int, default: 100)
max_iters: 100
# Whether to evaluate on the validation set at the beginning of the training
initial_validation: false
# Whether to evaluate on the validation set at the end the training
final_validation: true
# The name of the logger to send metrics to. (type: LoggerChoice, i.e. Literal['wandb', 'tensorboard', 'csv', 'mlflow', 'litlogger'], default: csv)
logger_name: csv
# The random seed to use for reproducibility. (type: int, default: 1337)
seed: 1337
# Optimizer-related arguments
optimizer:
class_path: torch.optim.AdamW
init_args:
# (type: float, default: 0.001)
lr: 0.0002
# (type: float, default: 0.01)
weight_decay: 0.0
# (type: tuple, default: (0.9,0.999))
betas:
- 0.9
- 0.95
config_hub/finetune/mistral-7b/qlora.yaml
+133
View File
@@ -0,0 +1,133 @@
# The path to the base model's checkpoint directory to load for finetuning. (type: <class 'Path'>, default: checkpoints/stabilityai/stablelm-base-alpha-3b)
checkpoint_dir: checkpoints/mistralai/Mistral-7B-v0.1
# Directory in which to save checkpoints and logs. (type: <class 'Path'>, default: out/lora)
out_dir: out/finetune/qlora-mistral-7b
# The precision to use for finetuning. Possible choices: "bf16-true", "bf16-mixed", "32-true". (type: Optional[str], default: null)
precision: bf16-true
# If set, quantize the model with this algorithm. See ``tutorials/quantize.md`` for more information. (type: Optional[Literal['nf4', 'nf4-dq', 'fp4', 'fp4-dq', 'int8-training']], default: null)
quantize: bnb.nf4
# How many devices/GPUs to use. (type: Union[int, str], default: 1)
devices: 1
# How many nodes to use. (type: int, default: 1)
num_nodes: 1
# The LoRA rank. (type: int, default: 8)
lora_r: 32
# The LoRA alpha. (type: int, default: 16)
lora_alpha: 16
# The LoRA dropout value. (type: float, default: 0.05)
lora_dropout: 0.05
# Whether to apply LoRA to the query weights in attention. (type: bool, default: True)
lora_query: true
# Whether to apply LoRA to the key weights in attention. (type: bool, default: False)
lora_key: false
# Whether to apply LoRA to the value weights in attention. (type: bool, default: True)
lora_value: true
# Whether to apply LoRA to the output projection in the attention block. (type: bool, default: False)
lora_projection: false
# Whether to apply LoRA to the weights of the MLP in the attention block. (type: bool, default: False)
lora_mlp: false
# Whether to apply LoRA to output head in GPT. (type: bool, default: False)
lora_head: false
# Data-related arguments. If not provided, the default is ``litgpt.data.Alpaca``.
data:
class_path: litgpt.data.Alpaca2k
init_args:
mask_prompt: false
val_split_fraction: 0.05
prompt_style: alpaca
ignore_index: -100
seed: 42
num_workers: 4
download_dir: data/alpaca2k
# Training-related arguments. See ``litgpt.args.TrainArgs`` for details
train:
# Number of optimizer steps between saving checkpoints (type: Optional[int], default: 1000)
save_interval: 200
# Number of iterations between logging calls (type: int, default: 1)
log_interval: 1
# Number of samples between optimizer steps across data-parallel ranks (type: int, default: 128)
global_batch_size: 8
# Number of samples per data-parallel rank (type: int, default: 4)
micro_batch_size: 2
# Number of iterations with learning rate warmup active (type: int, default: 100)
lr_warmup_steps: 10
# Number of epochs to train on (type: Optional[int], default: 5)
epochs: 4
# Total number of tokens to train on (type: Optional[int], default: null)
max_tokens:
# Limits the number of optimizer steps to run (type: Optional[int], default: null)
max_steps:
# Limits the length of samples (type: Optional[int], default: null)
max_seq_length: 512
# Whether to tie the embedding weights with the language modeling head weights (type: Optional[bool], default: null)
tie_embeddings:
# (type: Optional[float], default: null)
max_norm:
# (type: float, default: 6e-05)
min_lr: 6.0e-05
# Evaluation-related arguments. See ``litgpt.args.EvalArgs`` for details
eval:
# Number of optimizer steps between evaluation calls (type: int, default: 100)
interval: 100
# Number of tokens to generate (type: Optional[int], default: 100)
max_new_tokens: 100
# Number of iterations (type: int, default: 100)
max_iters: 100
# Whether to evaluate on the validation set at the beginning of the training
initial_validation: false
# Whether to evaluate on the validation set at the end the training
final_validation: true
# The name of the logger to send metrics to. (type: LoggerChoice, i.e. Literal['wandb', 'tensorboard', 'csv', 'mlflow', 'litlogger'], default: csv)
logger_name: csv
# The random seed to use for reproducibility. (type: int, default: 1337)
seed: 1337
# Optimizer-related arguments
optimizer:
class_path: torch.optim.AdamW
init_args:
# (type: float, default: 0.001)
lr: 0.0002
# (type: float, default: 0.01)
weight_decay: 0.0
# (type: tuple, default: (0.9,0.999))
betas:
- 0.9
- 0.95
config_hub/finetune/phi-2/full.yaml
+101
View File
@@ -0,0 +1,101 @@
# The path to the base model's checkpoint directory to load for finetuning. (type: <class 'Path'>, default: checkpoints/stabilityai/stablelm-base-alpha-3b)
checkpoint_dir: checkpoints/microsoft/phi-2
# Directory in which to save checkpoints and logs. (type: <class 'Path'>, default: out/finetune/full)
out_dir: out/finetune/full-phi-2
# The precision to use for finetuning. Possible choices: "bf16-true", "bf16-mixed", "32-true". (type: Optional[str], default: null)
precision: bf16-true
# How many devices/GPUs to use (type: Union[int, str], default: 1)
devices: 2
# How many nodes to use. (type: int, default: 1)
num_nodes: 1
# Data-related arguments. If not provided, the default is ``litgpt.data.Alpaca``.
data:
class_path: litgpt.data.Alpaca2k
init_args:
mask_prompt: false
prompt_style: alpaca
ignore_index: -100
seed: 42
num_workers: 4
# Training-related arguments. See ``litgpt.args.TrainArgs`` for details
train:
# Number of optimizer steps between saving checkpoints (type: Optional[int], default: 1000)
save_interval: 200
# Number of iterations between logging calls (type: int, default: 1)
log_interval: 1
# Number of samples between optimizer steps across data-parallel ranks (type: int, default: 64)
global_batch_size: 8
# Number of samples per data-parallel rank (type: int, default: 1)
micro_batch_size: 4
# Number of iterations with learning rate warmup active (type: int, default: 100)
lr_warmup_steps: 200
# Number of epochs to train on (type: Optional[int], default: 5)
epochs: 1
# Total number of tokens to train on (type: Optional[int], default: null)
max_tokens:
# Limits the number of optimizer steps to run. (type: Optional[int], default: null)
max_steps: 100
# Limits the length of samples. Off by default (type: Optional[int], default: null)
max_seq_length: 512
# Whether to tie the embedding weights with the language modeling head weights. (type: Optional[bool], default: null)
tie_embeddings:
# (type: Optional[float], default: null)
max_norm:
# (type: float, default: 6e-05)
min_lr: 6.0e-05
# Evaluation-related arguments. See ``litgpt.args.EvalArgs`` for details
eval:
# Number of optimizer steps between evaluation calls (type: int, default: 600)
interval: 25
# Number of tokens to generate (type: Optional[int], default: 100)
max_new_tokens: 100
# Number of iterations (type: int, default: 100)
max_iters: 100
# Whether to evaluate on the validation set at the beginning of the training
initial_validation: false
# Whether to evaluate on the validation set at the end the training
final_validation: true
# The name of the logger to send metrics to. (type: LoggerChoice, i.e. Literal['wandb', 'tensorboard', 'csv', 'mlflow', 'litlogger'], default: csv)
logger_name: csv
# The random seed to use for reproducibility. (type: int, default: 1337)
seed: 1337
# Optimizer-related arguments
optimizer:
class_path: torch.optim.AdamW
init_args:
# (type: float, default: 0.001)
lr: 0.0002
# (type: float, default: 0.01)
weight_decay: 0.1
# (type: tuple, default: (0.9,0.999))
betas:
- 0.9
- 0.95
config_hub/finetune/phi-2/lora.yaml
+132
View File
@@ -0,0 +1,132 @@
# The path to the base model's checkpoint directory to load for finetuning. (type: <class 'Path'>, default: checkpoints/stabilityai/stablelm-base-alpha-3b)
checkpoint_dir: checkpoints/microsoft/phi-2
# Directory in which to save checkpoints and logs. (type: <class 'Path'>, default: out/lora)
out_dir: out/finetune/lora-phi-2
# The precision to use for finetuning. Possible choices: "bf16-true", "bf16-mixed", "32-true". (type: Optional[str], default: null)
precision: bf16-true
# If set, quantize the model with this algorithm. See ``tutorials/quantize.md`` for more information. (type: Optional[Literal['nf4', 'nf4-dq', 'fp4', 'fp4-dq', 'int8-training']], default: null)
quantize:
# How many devices/GPUs to use. (type: Union[int, str], default: 1)
devices: 1
# How many nodes to use. (type: int, default: 1)
num_nodes: 1
# The LoRA rank. (type: int, default: 8)
lora_r: 8
# The LoRA alpha. (type: int, default: 16)
lora_alpha: 16
# The LoRA dropout value. (type: float, default: 0.05)
lora_dropout: 0.05
# Whether to apply LoRA to the query weights in attention. (type: bool, default: True)
lora_query: true
# Whether to apply LoRA to the key weights in attention. (type: bool, default: False)
lora_key: true
# Whether to apply LoRA to the value weights in attention. (type: bool, default: True)
lora_value: true
# Whether to apply LoRA to the output projection in the attention block. (type: bool, default: False)
lora_projection: true
# Whether to apply LoRA to the weights of the MLP in the attention block. (type: bool, default: False)
lora_mlp: true
# Whether to apply LoRA to output head in GPT. (type: bool, default: False)
lora_head: true
# Data-related arguments. If not provided, the default is ``litgpt.data.Alpaca``.
data:
class_path: litgpt.data.Alpaca2k
init_args:
mask_prompt: false
val_split_fraction: 0.03847
prompt_style: alpaca
ignore_index: -100
seed: 42
num_workers: 4
# Training-related arguments. See ``litgpt.args.TrainArgs`` for details
train:
# Number of optimizer steps between saving checkpoints (type: Optional[int], default: 1000)
save_interval: 800
# Number of iterations between logging calls (type: int, default: 1)
log_interval: 1
# Number of samples between optimizer steps across data-parallel ranks (type: int, default: 128)
global_batch_size: 8
# Number of samples per data-parallel rank (type: int, default: 4)
micro_batch_size: 4
# Number of iterations with learning rate warmup active (type: int, default: 100)
lr_warmup_steps: 10
# Number of epochs to train on (type: Optional[int], default: 5)
epochs: 1
# Total number of tokens to train on (type: Optional[int], default: null)
max_tokens:
# Limits the number of optimizer steps to run. (type: Optional[int], default: null)
max_steps:
# Limits the length of samples. Off by default (type: Optional[int], default: null)
max_seq_length: 512
# Whether to tie the embedding weights with the language modeling head weights. (type: Optional[bool], default: null)
tie_embeddings:
# (type: Optional[float], default: null)
max_norm:
# (type: float, default: 6e-05)
min_lr: 6.0e-05
# Evaluation-related arguments. See ``litgpt.args.EvalArgs`` for details
eval:
# Number of optimizer steps between evaluation calls (type: int, default: 100)
interval: 100
# Number of tokens to generate (type: Optional[int], default: 100)
max_new_tokens: 100
# Number of iterations (type: int, default: 100)
max_iters: 100
# Whether to evaluate on the validation set at the beginning of the training
initial_validation: false
# Whether to evaluate on the validation set at the end the training
final_validation: true
# The name of the logger to send metrics to. (type: LoggerChoice, i.e. Literal['wandb', 'tensorboard', 'csv', 'mlflow', 'litlogger'], default: csv)
logger_name: csv
# The random seed to use for reproducibility. (type: int, default: 1337)
seed: 1337
# Optimizer-related arguments
optimizer:
class_path: torch.optim.AdamW
init_args:
# (type: float, default: 0.001)
lr: 0.0002
# (type: float, default: 0.01)
weight_decay: 0.0
# (type: tuple, default: (0.9,0.999))
betas:
- 0.9
- 0.95
config_hub/finetune/phi-2/qlora.yaml
+132
View File
@@ -0,0 +1,132 @@
# The path to the base model's checkpoint directory to load for finetuning. (type: <class 'Path'>, default: checkpoints/stabilityai/stablelm-base-alpha-3b)
checkpoint_dir: checkpoints/microsoft/phi-2
# Directory in which to save checkpoints and logs. (type: <class 'Path'>, default: out/lora)
out_dir: out/finetune/qlora-phi-2
# The precision to use for finetuning. Possible choices: "bf16-true", "bf16-mixed", "32-true". (type: Optional[str], default: null)
precision: bf16-true
# If set, quantize the model with this algorithm. See ``tutorials/quantize.md`` for more information. (type: Optional[Literal['nf4', 'nf4-dq', 'fp4', 'fp4-dq', 'int8-training']], default: null)
quantize: bnb.nf4
# How many devices/GPUs to use. (type: Union[int, str], default: 1)
devices: 1
# How many nodes to use. (type: int, default: 1)
num_nodes: 1
# The LoRA rank. (type: int, default: 8)
lora_r: 8
# The LoRA alpha. (type: int, default: 16)
lora_alpha: 16
# The LoRA dropout value. (type: float, default: 0.05)
lora_dropout: 0.05
# Whether to apply LoRA to the query weights in attention. (type: bool, default: True)
lora_query: true
# Whether to apply LoRA to the key weights in attention. (type: bool, default: False)
lora_key: true
# Whether to apply LoRA to the value weights in attention. (type: bool, default: True)
lora_value: true
# Whether to apply LoRA to the output projection in the attention block. (type: bool, default: False)
lora_projection: true
# Whether to apply LoRA to the weights of the MLP in the attention block. (type: bool, default: False)
lora_mlp: true
# Whether to apply LoRA to output head in GPT. (type: bool, default: False)
lora_head: true
# Data-related arguments. If not provided, the default is ``litgpt.data.Alpaca``.
data:
class_path: litgpt.data.Alpaca2k
init_args:
mask_prompt: false
val_split_fraction: 0.03847
prompt_style: alpaca
ignore_index: -100
seed: 42
num_workers: 4
# Training-related arguments. See ``litgpt.args.TrainArgs`` for details
train:
# Number of optimizer steps between saving checkpoints (type: Optional[int], default: 1000)
save_interval: 800
# Number of iterations between logging calls (type: int, default: 1)
log_interval: 1
# Number of samples between optimizer steps across data-parallel ranks (type: int, default: 128)
global_batch_size: 8
# Number of samples per data-parallel rank (type: int, default: 4)
micro_batch_size: 4
# Number of iterations with learning rate warmup active (type: int, default: 100)
lr_warmup_steps: 10
# Number of epochs to train on (type: Optional[int], default: 5)
epochs: 1
# Total number of tokens to train on (type: Optional[int], default: null)
max_tokens:
# Limits the number of optimizer steps to run. (type: Optional[int], default: null)
max_steps:
# Limits the length of samples. Off by default (type: Optional[int], default: null)
max_seq_length: 512
# Whether to tie the embedding weights with the language modeling head weights. (type: Optional[bool], default: null)
tie_embeddings:
# (type: Optional[float], default: null)
max_norm:
# (type: float, default: 6e-05)
min_lr: 6.0e-05
# Evaluation-related arguments. See ``litgpt.args.EvalArgs`` for details
eval:
# Number of optimizer steps between evaluation calls (type: int, default: 100)
interval: 100
# Number of tokens to generate (type: Optional[int], default: 100)
max_new_tokens: 100
# Number of iterations (type: int, default: 100)
max_iters: 100
# Whether to evaluate on the validation set at the beginning of the training
initial_validation: false
# Whether to evaluate on the validation set at the end the training
final_validation: true
# The name of the logger to send metrics to. (type: LoggerChoice, i.e. Literal['wandb', 'tensorboard', 'csv', 'mlflow', 'litlogger'], default: csv)
logger_name: csv
# The random seed to use for reproducibility. (type: int, default: 1337)
seed: 1337
# Optimizer-related arguments
optimizer:
class_path: torch.optim.AdamW
init_args:
# (type: float, default: 0.001)
lr: 0.0002
# (type: float, default: 0.01)
weight_decay: 0.0
# (type: tuple, default: (0.9,0.999))
betas:
- 0.9
- 0.95
config_hub/finetune/phi-3/full.yaml
+98
View File
@@ -0,0 +1,98 @@
# The path to the base model's checkpoint directory to load for finetuning. (type: <class 'Path'>, default: checkpoints/stabilityai/stablelm-base-alpha-3b)
checkpoint_dir: checkpoints/microsoft/Phi-3-mini-4k-instruct
# Directory in which to save checkpoints and logs. (type: <class 'Path'>, default: out/finetune/full)
out_dir: out/finetune/full-phi-3
# The precision to use for finetuning. Possible choices: "bf16-true", "bf16-mixed", "32-true". (type: Optional[str], default: null)
precision: bf16-true
# How many devices/GPUs to use (type: Union[int, str], default: 1)
devices: 1
# Data-related arguments. If not provided, the default is ``litgpt.data.Alpaca``.
data:
class_path: litgpt.data.Alpaca2k
init_args:
mask_prompt: false
prompt_style: alpaca
ignore_index: -100
seed: 42
num_workers: 4
# Training-related arguments. See ``litgpt.args.TrainArgs`` for details
train:
# Number of optimizer steps between saving checkpoints (type: Optional[int], default: 1000)
save_interval: 200
# Number of iterations between logging calls (type: int, default: 1)
log_interval: 1
# Number of samples between optimizer steps across data-parallel ranks (type: int, default: 64)
global_batch_size: 8
# Number of samples per data-parallel rank (type: int, default: 1)
micro_batch_size: 4
# Number of iterations with learning rate warmup active (type: int, default: 100)
lr_warmup_steps: 200
# Number of epochs to train on (type: Optional[int], default: 5)
epochs: 1
# Total number of tokens to train on (type: Optional[int], default: null)
max_tokens:
# Limits the number of optimizer steps to run. (type: Optional[int], default: null)
max_steps:
# Limits the length of samples. Off by default (type: Optional[int], default: null)
max_seq_length: 512
# Whether to tie the embedding weights with the language modeling head weights. (type: Optional[bool], default: null)
tie_embeddings:
# (type: Optional[float], default: null)
max_norm:
# (type: float, default: 6e-05)
min_lr: 6.0e-05
# Evaluation-related arguments. See ``litgpt.args.EvalArgs`` for details
eval:
# Number of optimizer steps between evaluation calls (type: int, default: 600)
interval: 25
# Number of tokens to generate (type: Optional[int], default: 100)
max_new_tokens: 100
# Number of iterations (type: int, default: 100)
max_iters: 100
# Whether to evaluate on the validation set at the beginning of the training
initial_validation: false
# Whether to evaluate on the validation set at the end the training
final_validation: true
# The name of the logger to send metrics to. (type: LoggerChoice, i.e. Literal['wandb', 'tensorboard', 'csv', 'mlflow', 'litlogger'], default: csv)
logger_name: csv
# The random seed to use for reproducibility. (type: int, default: 1337)
seed: 1337
# Optimizer-related arguments
optimizer:
class_path: torch.optim.AdamW
init_args:
# (type: float, default: 0.001)
lr: 0.0002
# (type: float, default: 0.01)
weight_decay: 0.1
# (type: tuple, default: (0.9,0.999))
betas:
- 0.9
- 0.95
config_hub/finetune/phi-3/lora.yaml
+129
View File
@@ -0,0 +1,129 @@
# The path to the base model's checkpoint directory to load for finetuning. (type: <class 'Path'>, default: checkpoints/stabilityai/stablelm-base-alpha-3b)
checkpoint_dir: checkpoints/microsoft/Phi-3-mini-4k-instruct
# Directory in which to save checkpoints and logs. (type: <class 'Path'>, default: out/lora)
out_dir: out/finetune/lora-phi-3
# The precision to use for finetuning. Possible choices: "bf16-true", "bf16-mixed", "32-true". (type: Optional[str], default: null)
precision: bf16-true
# If set, quantize the model with this algorithm. See ``tutorials/quantize.md`` for more information. (type: Optional[Literal['nf4', 'nf4-dq', 'fp4', 'fp4-dq', 'int8-training']], default: null)
quantize:
# How many devices/GPUs to use. (type: Union[int, str], default: 1)
devices: 1
# The LoRA rank. (type: int, default: 8)
lora_r: 8
# The LoRA alpha. (type: int, default: 16)
lora_alpha: 16
# The LoRA dropout value. (type: float, default: 0.05)
lora_dropout: 0.05
# Whether to apply LoRA to the query weights in attention. (type: bool, default: True)
lora_query: true
# Whether to apply LoRA to the key weights in attention. (type: bool, default: False)
lora_key: true
# Whether to apply LoRA to the value weights in attention. (type: bool, default: True)
lora_value: true
# Whether to apply LoRA to the output projection in the attention block. (type: bool, default: False)
lora_projection: true
# Whether to apply LoRA to the weights of the MLP in the attention block. (type: bool, default: False)
lora_mlp: true
# Whether to apply LoRA to output head in GPT. (type: bool, default: False)
lora_head: true
# Data-related arguments. If not provided, the default is ``litgpt.data.Alpaca``.
data:
class_path: litgpt.data.Alpaca2k
init_args:
mask_prompt: false
val_split_fraction: 0.03847
prompt_style: alpaca
ignore_index: -100
seed: 42
num_workers: 4
# Training-related arguments. See ``litgpt.args.TrainArgs`` for details
train:
# Number of optimizer steps between saving checkpoints (type: Optional[int], default: 1000)
save_interval: 800
# Number of iterations between logging calls (type: int, default: 1)
log_interval: 1
# Number of samples between optimizer steps across data-parallel ranks (type: int, default: 128)
global_batch_size: 8
# Number of samples per data-parallel rank (type: int, default: 4)
micro_batch_size: 4
# Number of iterations with learning rate warmup active (type: int, default: 100)
lr_warmup_steps: 10
# Number of epochs to train on (type: Optional[int], default: 5)
epochs: 1
# Total number of tokens to train on (type: Optional[int], default: null)
max_tokens:
# Limits the number of optimizer steps to run. (type: Optional[int], default: null)
max_steps:
# Limits the length of samples. Off by default (type: Optional[int], default: null)
max_seq_length: 512
# Whether to tie the embedding weights with the language modeling head weights. (type: Optional[bool], default: null)
tie_embeddings:
# (type: Optional[float], default: null)
max_norm:
# (type: float, default: 6e-05)
min_lr: 6.0e-05
# Evaluation-related arguments. See ``litgpt.args.EvalArgs`` for details
eval:
# Number of optimizer steps between evaluation calls (type: int, default: 100)
interval: 100
# Number of tokens to generate (type: Optional[int], default: 100)
max_new_tokens: 100
# Number of iterations (type: int, default: 100)
max_iters: 100
# Whether to evaluate on the validation set at the beginning of the training
initial_validation: false
# Whether to evaluate on the validation set at the end the training
final_validation: true
# The name of the logger to send metrics to. (type: LoggerChoice, i.e. Literal['wandb', 'tensorboard', 'csv', 'mlflow', 'litlogger'], default: csv)
logger_name: csv
# The random seed to use for reproducibility. (type: int, default: 1337)
seed: 1337
# Optimizer-related arguments
optimizer:
class_path: torch.optim.AdamW
init_args:
# (type: float, default: 0.001)
lr: 0.0002
# (type: float, default: 0.01)
weight_decay: 0.0
# (type: tuple, default: (0.9,0.999))
betas:
- 0.9
- 0.95
config_hub/finetune/phi-3/qlora.yaml
+129
View File
@@ -0,0 +1,129 @@
# The path to the base model's checkpoint directory to load for finetuning. (type: <class 'Path'>, default: checkpoints/stabilityai/stablelm-base-alpha-3b)
checkpoint_dir: checkpoints/microsoft/Phi-3-mini-4k-instruct
# Directory in which to save checkpoints and logs. (type: <class 'Path'>, default: out/lora)
out_dir: out/finetune/qlora-phi-3
# The precision to use for finetuning. Possible choices: "bf16-true", "bf16-mixed", "32-true". (type: Optional[str], default: null)
precision: bf16-true
# If set, quantize the model with this algorithm. See ``tutorials/quantize.md`` for more information. (type: Optional[Literal['nf4', 'nf4-dq', 'fp4', 'fp4-dq', 'int8-training']], default: null)
quantize: bnb.nf4
# How many devices/GPUs to use. (type: Union[int, str], default: 1)
devices: 1
# The LoRA rank. (type: int, default: 8)
lora_r: 8
# The LoRA alpha. (type: int, default: 16)
lora_alpha: 16
# The LoRA dropout value. (type: float, default: 0.05)
lora_dropout: 0.05
# Whether to apply LoRA to the query weights in attention. (type: bool, default: True)
lora_query: true
# Whether to apply LoRA to the key weights in attention. (type: bool, default: False)
lora_key: true
# Whether to apply LoRA to the value weights in attention. (type: bool, default: True)
lora_value: true
# Whether to apply LoRA to the output projection in the attention block. (type: bool, default: False)
lora_projection: true
# Whether to apply LoRA to the weights of the MLP in the attention block. (type: bool, default: False)
lora_mlp: true
# Whether to apply LoRA to output head in GPT. (type: bool, default: False)
lora_head: true
# Data-related arguments. If not provided, the default is ``litgpt.data.Alpaca``.
data:
class_path: litgpt.data.Alpaca2k
init_args:
mask_prompt: false
val_split_fraction: 0.03847
prompt_style: alpaca
ignore_index: -100
seed: 42
num_workers: 4
# Training-related arguments. See ``litgpt.args.TrainArgs`` for details
train:
# Number of optimizer steps between saving checkpoints (type: Optional[int], default: 1000)
save_interval: 800
# Number of iterations between logging calls (type: int, default: 1)
log_interval: 1
# Number of samples between optimizer steps across data-parallel ranks (type: int, default: 128)
global_batch_size: 8
# Number of samples per data-parallel rank (type: int, default: 4)
micro_batch_size: 4
# Number of iterations with learning rate warmup active (type: int, default: 100)
lr_warmup_steps: 10
# Number of epochs to train on (type: Optional[int], default: 5)
epochs: 1
# Total number of tokens to train on (type: Optional[int], default: null)
max_tokens:
# Limits the number of optimizer steps to run. (type: Optional[int], default: null)
max_steps:
# Limits the length of samples. Off by default (type: Optional[int], default: null)
max_seq_length: 512
# Whether to tie the embedding weights with the language modeling head weights. (type: Optional[bool], default: null)
tie_embeddings:
# (type: Optional[float], default: null)
max_norm:
# (type: float, default: 6e-05)
min_lr: 6.0e-05
# Evaluation-related arguments. See ``litgpt.args.EvalArgs`` for details
eval:
# Number of optimizer steps between evaluation calls (type: int, default: 100)
interval: 100
# Number of tokens to generate (type: Optional[int], default: 100)
max_new_tokens: 100
# Number of iterations (type: int, default: 100)
max_iters: 100
# Whether to evaluate on the validation set at the beginning of the training
initial_validation: false
# Whether to evaluate on the validation set at the end the training
final_validation: true
# The name of the logger to send metrics to. (type: LoggerChoice, i.e. Literal['wandb', 'tensorboard', 'csv', 'mlflow', 'litlogger'], default: csv)
logger_name: csv
# The random seed to use for reproducibility. (type: int, default: 1337)
seed: 1337
# Optimizer-related arguments
optimizer:
class_path: torch.optim.AdamW
init_args:
# (type: float, default: 0.001)
lr: 0.0002
# (type: float, default: 0.01)
weight_decay: 0.0
# (type: tuple, default: (0.9,0.999))
betas:
- 0.9
- 0.95
config_hub/finetune/stablelm-base-alpha-3b/full.yaml
+102
View File
@@ -0,0 +1,102 @@
# The path to the base model's checkpoint directory to load for finetuning. (type: <class 'Path'>, default: checkpoints/stabilityai/stablelm-base-alpha-3b)
checkpoint_dir: checkpoints/stabilityai/stablelm-base-alpha-3b
# Directory in which to save checkpoints and logs. (type: <class 'Path'>, default: out/lora)
out_dir: out/finetune/full-stablelm-base-alpha-3b
# The precision to use for finetuning. Possible choices: "bf16-true", "bf16-mixed", "32-true". (type: Optional[str], default: null)
precision: bf16-true
# How many devices/GPUs to use. (type: Union[int, str], default: 1)
devices: 2
# How many nodes to use. (type: int, default: 1)
num_nodes: 1
# Data-related arguments. If not provided, the default is ``litgpt.data.Alpaca``.
data:
class_path: litgpt.data.Alpaca2k
init_args:
mask_prompt: false
val_split_fraction: 0.03847
prompt_style: alpaca
ignore_index: -100
seed: 42
num_workers: 4
# Training-related arguments. See ``litgpt.args.TrainArgs`` for details
train:
# Number of optimizer steps between saving checkpoints (type: Optional[int], default: 1000)
save_interval: 800
# Number of iterations between logging calls (type: int, default: 1)
log_interval: 1
# Number of samples between optimizer steps across data-parallel ranks (type: int, default: 128)
global_batch_size: 8
# Number of samples per data-parallel rank (type: int, default: 4)
micro_batch_size: 1
# Number of iterations with learning rate warmup active (type: int, default: 100)
lr_warmup_steps: 1000
# Number of epochs to train on (type: Optional[int], default: 5)
epochs: 1
# Total number of tokens to train on (type: Optional[int], default: null)
max_tokens:
# Limits the number of optimizer steps to run. (type: Optional[int], default: null)
max_steps:
# Limits the length of samples. Off by default (type: Optional[int], default: null)
max_seq_length: 512
# Whether to tie the embedding weights with the language modeling head weights. (type: Optional[bool], default: null)
tie_embeddings:
# (type: Optional[float], default: null)
max_norm:
# (type: float, default: 6e-05)
min_lr: 6.0e-05
# Evaluation-related arguments. See ``litgpt.args.EvalArgs`` for details
eval:
# Number of optimizer steps between evaluation calls (type: int, default: 100)
interval: 25
# Number of tokens to generate (type: Optional[int], default: 100)
max_new_tokens: 100
# Number of iterations (type: int, default: 100)
max_iters: 100
# Whether to evaluate on the validation set at the beginning of the training
initial_validation: false
# Whether to evaluate on the validation set at the end the training
final_validation: true
# The name of the logger to send metrics to. (type: LoggerChoice, i.e. Literal['wandb', 'tensorboard', 'csv', 'mlflow', 'litlogger'], default: csv)
logger_name: csv
# The random seed to use for reproducibility. (type: int, default: 1337)
seed: 1337
# Optimizer-related arguments
optimizer:
class_path: torch.optim.AdamW
init_args:
# (type: float, default: 0.001)
lr: 0.0002
# (type: float, default: 0.01)
weight_decay: 0.1
# (type: tuple, default: (0.9,0.999))
betas:
- 0.9
- 0.95
config_hub/finetune/stablelm-base-alpha-3b/lora.yaml
+131
View File
@@ -0,0 +1,131 @@
# The path to the base model's checkpoint directory to load for finetuning. (type: <class 'Path'>, default: checkpoints/stabilityai/stablelm-base-alpha-3b)
checkpoint_dir: checkpoints/stabilityai/stablelm-base-alpha-3b
# Directory in which to save checkpoints and logs. (type: <class 'Path'>, default: out/lora)
out_dir: out/finetune/lora-stablelm-base-alpha-3b
# The precision to use for finetuning. Possible choices: "bf16-true", "bf16-mixed", "32-true". (type: Optional[str], default: null)
precision: bf16-true
# If set, quantize the model with this algorithm. See ``tutorials/quantize.md`` for more information. (type: Optional[Literal['nf4', 'nf4-dq', 'fp4', 'fp4-dq', 'int8-training']], default: null)
quantize:
# How many devices/GPUs to use. (type: Union[int, str], default: 1)
devices: 1
# How many nodes to use. (type: int, default: 1)
num_nodes: 1
# The LoRA rank. (type: int, default: 8)
lora_r: 32
# The LoRA alpha. (type: int, default: 16)
lora_alpha: 16
# The LoRA dropout value. (type: float, default: 0.05)
lora_dropout: 0.05
# Whether to apply LoRA to the query weights in attention. (type: bool, default: True)
lora_query: true
# Whether to apply LoRA to the key weights in attention. (type: bool, default: False)
lora_key: false
# Whether to apply LoRA to the value weights in attention. (type: bool, default: True)
lora_value: true
# Whether to apply LoRA to the output projection in the attention block. (type: bool, default: False)
lora_projection: false
# Whether to apply LoRA to the weights of the MLP in the attention block. (type: bool, default: False)
lora_mlp: false
# Whether to apply LoRA to output head in GPT. (type: bool, default: False)
lora_head: false
# Data-related arguments. If not provided, the default is ``litgpt.data.Alpaca``.
data:
class_path: litgpt.data.Alpaca2k
init_args:
mask_prompt: false
prompt_style: alpaca
ignore_index: -100
seed: 42
num_workers: 4
# Training-related arguments. See ``litgpt.args.TrainArgs`` for details
train:
# Number of optimizer steps between saving checkpoints (type: Optional[int], default: 1000)
save_interval: 200
# Number of iterations between logging calls (type: int, default: 1)
log_interval: 1
# Number of samples between optimizer steps across data-parallel ranks (type: int, default: 128)
global_batch_size: 8
# Number of samples per data-parallel rank (type: int, default: 4)
micro_batch_size: 1
# Number of iterations with learning rate warmup active (type: int, default: 100)
lr_warmup_steps: 10
# Number of epochs to train on (type: Optional[int], default: 5)
epochs: 4
# Total number of tokens to train on (type: Optional[int], default: null)
max_tokens:
# Limits the number of optimizer steps to run. (type: Optional[int], default: null)
max_steps:
# Limits the length of samples. Off by default (type: Optional[int], default: null)
max_seq_length: 512
# Whether to tie the embedding weights with the language modeling head weights. (type: Optional[bool], default: null)
tie_embeddings:
# (type: Optional[float], default: null)
max_norm:
# (type: float, default: 6e-05)
min_lr: 6.0e-05
# Evaluation-related arguments. See ``litgpt.args.EvalArgs`` for details
eval:
# Number of optimizer steps between evaluation calls (type: int, default: 100)
interval: 100
# Number of tokens to generate (type: Optional[int], default: 100)
max_new_tokens: 100
# Number of iterations (type: int, default: 100)
max_iters: 100
# Whether to evaluate on the validation set at the beginning of the training
initial_validation: false
# Whether to evaluate on the validation set at the end the training
final_validation: true
# The name of the logger to send metrics to. (type: LoggerChoice, i.e. Literal['wandb', 'tensorboard', 'csv', 'mlflow', 'litlogger'], default: csv)
logger_name: csv
# The random seed to use for reproducibility. (type: int, default: 1337)
seed: 1337
# Optimizer-related arguments
optimizer:
class_path: torch.optim.AdamW
init_args:
# (type: float, default: 0.001)
lr: 0.0002
# (type: float, default: 0.01)
weight_decay: 0.0
# (type: tuple, default: (0.9,0.999))
betas:
- 0.9
- 0.95
config_hub/finetune/stablelm-base-alpha-3b/qlora.yaml
+133
View File
@@ -0,0 +1,133 @@
# The path to the base model's checkpoint directory to load for finetuning. (type: <class 'Path'>, default: checkpoints/stabilityai/stablelm-base-alpha-3b)
checkpoint_dir: checkpoints/stabilityai/stablelm-base-alpha-3b
# Directory in which to save checkpoints and logs. (type: <class 'Path'>, default: out/lora)
out_dir: out/finetune/qlora-stablelm-base-alpha-3b
# The precision to use for finetuning. Possible choices: "bf16-true", "bf16-mixed", "32-true". (type: Optional[str], default: null)
precision: bf16-true
# If set, quantize the model with this algorithm. See ``tutorials/quantize.md`` for more information. (type: Optional[Literal['nf4', 'nf4-dq', 'fp4', 'fp4-dq', 'int8-training']], default: null)
quantize: bnb.nf4
# How many devices/GPUs to use. (type: Union[int, str], default: 1)
devices: 1
# How many nodes to use. (type: int, default: 1)
num_nodes: 1
# The LoRA rank. (type: int, default: 8)
lora_r: 32
# The LoRA alpha. (type: int, default: 16)
lora_alpha: 16
# The LoRA dropout value. (type: float, default: 0.05)
lora_dropout: 0.05
# Whether to apply LoRA to the query weights in attention. (type: bool, default: True)
lora_query: true
# Whether to apply LoRA to the key weights in attention. (type: bool, default: False)
lora_key: false
# Whether to apply LoRA to the value weights in attention. (type: bool, default: True)
lora_value: true
# Whether to apply LoRA to the output projection in the attention block. (type: bool, default: False)
lora_projection: false
# Whether to apply LoRA to the weights of the MLP in the attention block. (type: bool, default: False)
lora_mlp: false
# Whether to apply LoRA to output head in GPT. (type: bool, default: False)
lora_head: false
# Data-related arguments. If not provided, the default is ``litgpt.data.Alpaca``.
data:
class_path: litgpt.data.Alpaca2k
init_args:
mask_prompt: false
val_split_fraction: 0.05
prompt_style: alpaca
ignore_index: -100
seed: 42
num_workers: 4
download_dir: data/alpaca2k
# Training-related arguments. See ``litgpt.args.TrainArgs`` for details
train:
# Number of optimizer steps between saving checkpoints (type: Optional[int], default: 1000)
save_interval: 200
# Number of iterations between logging calls (type: int, default: 1)
log_interval: 1
# Number of samples between optimizer steps across data-parallel ranks (type: int, default: 128)
global_batch_size: 8
# Number of samples per data-parallel rank (type: int, default: 4)
micro_batch_size: 1
# Number of iterations with learning rate warmup active (type: int, default: 100)
lr_warmup_steps: 10
# Number of epochs to train on (type: Optional[int], default: 5)
epochs: 4
# Total number of tokens to train on (type: Optional[int], default: null)
max_tokens:
# Limits the number of optimizer steps to run (type: Optional[int], default: null)
max_steps:
# Limits the length of samples (type: Optional[int], default: null)
max_seq_length: 512
# Whether to tie the embedding weights with the language modeling head weights (type: Optional[bool], default: null)
tie_embeddings:
# (type: Optional[float], default: null)
max_norm:
# (type: float, default: 6e-05)
min_lr: 6.0e-05
# Evaluation-related arguments. See ``litgpt.args.EvalArgs`` for details
eval:
# Number of optimizer steps between evaluation calls (type: int, default: 100)
interval: 100
# Number of tokens to generate (type: Optional[int], default: 100)
max_new_tokens: 100
# Number of iterations (type: int, default: 100)
max_iters: 100
# Whether to evaluate on the validation set at the beginning of the training
initial_validation: false
# Whether to evaluate on the validation set at the end the training
final_validation: true
# The name of the logger to send metrics to. (type: LoggerChoice, i.e. Literal['wandb', 'tensorboard', 'csv', 'mlflow', 'litlogger'], default: csv)
logger_name: csv
# The random seed to use for reproducibility. (type: int, default: 1337)
seed: 1337
# Optimizer-related arguments
optimizer:
class_path: torch.optim.AdamW
init_args:
# (type: float, default: 0.001)
lr: 0.0002
# (type: float, default: 0.01)
weight_decay: 0.0
# (type: tuple, default: (0.9,0.999))
betas:
- 0.9
- 0.95
config_hub/finetune/tiny-llama/full.yaml
+102
View File
@@ -0,0 +1,102 @@
# The path to the base model's checkpoint directory to load for finetuning. (type: <class 'Path'>, default: checkpoints/stabilityai/stablelm-base-alpha-3b)
checkpoint_dir: checkpoints/TinyLlama/TinyLlama-1.1B-intermediate-step-1431k-3T
# Directory in which to save checkpoints and logs. (type: <class 'Path'>, default: out/lora)
out_dir: out/finetune/full-tiny-llama-1.1b
# The precision to use for finetuning. Possible choices: "bf16-true", "bf16-mixed", "32-true". (type: Optional[str], default: null)
precision: bf16-true
# How many devices/GPUs to use. (type: Union[int, str], default: 1)
devices: 1
# How many nodes to use. (type: int, default: 1)
num_nodes: 1
# Data-related arguments. If not provided, the default is ``litgpt.data.Alpaca``.
data:
class_path: litgpt.data.Alpaca2k
init_args:
mask_prompt: false
val_split_fraction: 0.03847
prompt_style: alpaca
ignore_index: -100
seed: 42
num_workers: 4
# Training-related arguments. See ``litgpt.args.TrainArgs`` for details
train:
# Number of optimizer steps between saving checkpoints (type: Optional[int], default: 1000)
save_interval: 800
# Number of iterations between logging calls (type: int, default: 1)
log_interval: 1
# Number of samples between optimizer steps across data-parallel ranks (type: int, default: 128)
global_batch_size: 32
# Number of samples per data-parallel rank (type: int, default: 4)
micro_batch_size: 4
# Number of iterations with learning rate warmup active (type: int, default: 100)
lr_warmup_steps: 1000
# Number of epochs to train on (type: Optional[int], default: 5)
epochs: 1
# Total number of tokens to train on (type: Optional[int], default: null)
max_tokens:
# Limits the number of optimizer steps to run. (type: Optional[int], default: null)
max_steps:
# Limits the length of samples. Off by default (type: Optional[int], default: null)
max_seq_length: 512
# Whether to tie the embedding weights with the language modeling head weights. (type: Optional[bool], default: null)
tie_embeddings:
# (type: Optional[float], default: null)
max_norm:
# (type: float, default: 6e-05)
min_lr: 6.0e-05
# Evaluation-related arguments. See ``litgpt.args.EvalArgs`` for details
eval:
# Number of optimizer steps between evaluation calls (type: int, default: 100)
interval: 25
# Number of tokens to generate (type: Optional[int], default: 100)
max_new_tokens: 100
# Number of iterations (type: int, default: 100)
max_iters: 100
# Whether to evaluate on the validation set at the beginning of the training
initial_validation: false
# Whether to evaluate on the validation set at the end the training
final_validation: true
# The name of the logger to send metrics to. (type: LoggerChoice, i.e. Literal['wandb', 'tensorboard', 'csv', 'mlflow', 'litlogger'], default: csv)
logger_name: csv
# The random seed to use for reproducibility. (type: int, default: 1337)
seed: 1337
# Optimizer-related arguments
optimizer:
class_path: torch.optim.AdamW
init_args:
# (type: float, default: 0.001)
lr: 0.0002
# (type: float, default: 0.01)
weight_decay: 0.0
# (type: tuple, default: (0.9,0.999))
betas:
- 0.9
- 0.95
config_hub/finetune/tiny-llama/lora.yaml
+132
View File
@@ -0,0 +1,132 @@
# The path to the base model's checkpoint directory to load for finetuning. (type: <class 'Path'>, default: checkpoints/stabilityai/stablelm-base-alpha-3b)
checkpoint_dir: checkpoints/TinyLlama/TinyLlama-1.1B-intermediate-step-1431k-3T
# Directory in which to save checkpoints and logs. (type: <class 'Path'>, default: out/lora)
out_dir: out/finetune/lora-tiny-llama-1.1b
# The precision to use for finetuning. Possible choices: "bf16-true", "bf16-mixed", "32-true". (type: Optional[str], default: null)
precision: bf16-true
# If set, quantize the model with this algorithm. See ``tutorials/quantize.md`` for more information. (type: Optional[Literal['nf4', 'nf4-dq', 'fp4', 'fp4-dq', 'int8-training']], default: null)
quantize:
# How many devices/GPUs to use. (type: Union[int, str], default: 1)
devices: 1
# How many nodes to use. (type: int, default: 1)
num_nodes: 1
# The LoRA rank. (type: int, default: 8)
lora_r: 32
# The LoRA alpha. (type: int, default: 16)
lora_alpha: 16
# The LoRA dropout value. (type: float, default: 0.05)
lora_dropout: 0.05
# Whether to apply LoRA to the query weights in attention. (type: bool, default: True)
lora_query: true
# Whether to apply LoRA to the key weights in attention. (type: bool, default: False)
lora_key: true
# Whether to apply LoRA to the value weights in attention. (type: bool, default: True)
lora_value: true
# Whether to apply LoRA to the output projection in the attention block. (type: bool, default: False)
lora_projection: true
# Whether to apply LoRA to the weights of the MLP in the attention block. (type: bool, default: False)
lora_mlp: true
# Whether to apply LoRA to output head in GPT. (type: bool, default: False)
lora_head: true
# Data-related arguments. If not provided, the default is ``litgpt.data.Alpaca``.
data:
class_path: litgpt.data.Alpaca2k
init_args:
mask_prompt: false
val_split_fraction: 0.03847
prompt_style: alpaca
ignore_index: -100
seed: 42
num_workers: 4
# Training-related arguments. See ``litgpt.args.TrainArgs`` for details
train:
# Number of optimizer steps between saving checkpoints (type: Optional[int], default: 1000)
save_interval: 800
# Number of iterations between logging calls (type: int, default: 1)
log_interval: 1
# Number of samples between optimizer steps across data-parallel ranks (type: int, default: 128)
global_batch_size: 8
# Number of samples per data-parallel rank (type: int, default: 4)
micro_batch_size: 8
# Number of iterations with learning rate warmup active (type: int, default: 100)
lr_warmup_steps: 10
# Number of epochs to train on (type: Optional[int], default: 5)
epochs: 3
# Total number of tokens to train on (type: Optional[int], default: null)
max_tokens:
# Limits the number of optimizer steps to run. (type: Optional[int], default: null)
max_steps:
# Limits the length of samples. Off by default (type: Optional[int], default: null)
max_seq_length: 512
# Whether to tie the embedding weights with the language modeling head weights. (type: Optional[bool], default: null)
tie_embeddings:
# (type: Optional[float], default: null)
max_norm:
# (type: float, default: 6e-05)
min_lr: 6.0e-05
# Evaluation-related arguments. See ``litgpt.args.EvalArgs`` for details
eval:
# Number of optimizer steps between evaluation calls (type: int, default: 100)
interval: 100
# Number of tokens to generate (type: Optional[int], default: 100)
max_new_tokens: 100
# Number of iterations (type: int, default: 100)
max_iters: 100
# Whether to evaluate on the validation set at the beginning of the training
initial_validation: false
# Whether to evaluate on the validation set at the end the training
final_validation: true
# The name of the logger to send metrics to. (type: LoggerChoice, i.e. Literal['wandb', 'tensorboard', 'csv', 'mlflow', 'litlogger'], default: csv)
logger_name: csv
# The random seed to use for reproducibility. (type: int, default: 1337)
seed: 1337
# Optimizer-related arguments
optimizer:
class_path: torch.optim.AdamW
init_args:
# (type: float, default: 0.001)
lr: 0.0002
# (type: float, default: 0.01)
weight_decay: 0.0
# (type: tuple, default: (0.9,0.999))
betas:
- 0.9
- 0.95
config_hub/finetune/tiny-llama/qlora.yaml
+132
View File
@@ -0,0 +1,132 @@
# The path to the base model's checkpoint directory to load for finetuning. (type: <class 'Path'>, default: checkpoints/stabilityai/stablelm-base-alpha-3b)
checkpoint_dir: checkpoints/TinyLlama/TinyLlama-1.1B-intermediate-step-1431k-3T
# Directory in which to save checkpoints and logs. (type: <class 'Path'>, default: out/lora)
out_dir: out/finetune/qlora-tiny-llama-1.1b
# The precision to use for finetuning. Possible choices: "bf16-true", "bf16-mixed", "32-true". (type: Optional[str], default: null)
precision: bf16-true
# If set, quantize the model with this algorithm. See ``tutorials/quantize.md`` for more information. (type: Optional[Literal['nf4', 'nf4-dq', 'fp4', 'fp4-dq', 'int8-training']], default: null)
quantize: bnb.nf4
# How many devices/GPUs to use. (type: Union[int, str], default: 1)
devices: 1
# How many nodes to use. (type: int, default: 1)
num_nodes: 1
# The LoRA rank. (type: int, default: 8)
lora_r: 32
# The LoRA alpha. (type: int, default: 16)
lora_alpha: 16
# The LoRA dropout value. (type: float, default: 0.05)
lora_dropout: 0.05
# Whether to apply LoRA to the query weights in attention. (type: bool, default: True)
lora_query: true
# Whether to apply LoRA to the key weights in attention. (type: bool, default: False)
lora_key: true
# Whether to apply LoRA to the value weights in attention. (type: bool, default: True)
lora_value: true
# Whether to apply LoRA to the output projection in the attention block. (type: bool, default: False)
lora_projection: true
# Whether to apply LoRA to the weights of the MLP in the attention block. (type: bool, default: False)
lora_mlp: true
# Whether to apply LoRA to output head in GPT. (type: bool, default: False)
lora_head: true
# Data-related arguments. If not provided, the default is ``litgpt.data.Alpaca``.
data:
class_path: litgpt.data.Alpaca2k
init_args:
mask_prompt: false
val_split_fraction: 0.03847
prompt_style: alpaca
ignore_index: -100
seed: 42
num_workers: 4
# Training-related arguments. See ``litgpt.args.TrainArgs`` for details
train:
# Number of optimizer steps between saving checkpoints (type: Optional[int], default: 1000)
save_interval: 800
# Number of iterations between logging calls (type: int, default: 1)
log_interval: 1
# Number of samples between optimizer steps across data-parallel ranks (type: int, default: 128)
global_batch_size: 8
# Number of samples per data-parallel rank (type: int, default: 4)
micro_batch_size: 8
# Number of iterations with learning rate warmup active (type: int, default: 100)
lr_warmup_steps: 10
# Number of epochs to train on (type: Optional[int], default: 5)
epochs: 3
# Total number of tokens to train on (type: Optional[int], default: null)
max_tokens:
# Limits the number of optimizer steps to run. (type: Optional[int], default: null)
max_steps:
# Limits the length of samples. Off by default (type: Optional[int], default: null)
max_seq_length: 512
# Whether to tie the embedding weights with the language modeling head weights. (type: Optional[bool], default: null)
tie_embeddings:
# (type: Optional[float], default: null)
max_norm:
# (type: float, default: 6e-05)
min_lr: 6.0e-05
# Evaluation-related arguments. See ``litgpt.args.EvalArgs`` for details
eval:
# Number of optimizer steps between evaluation calls (type: int, default: 100)
interval: 100
# Number of tokens to generate (type: Optional[int], default: 100)
max_new_tokens: 100
# Number of iterations (type: int, default: 100)
max_iters: 100
# Whether to evaluate on the validation set at the beginning of the training
initial_validation: false
# Whether to evaluate on the validation set at the end the training
final_validation: true
# The name of the logger to send metrics to. (type: LoggerChoice, i.e. Literal['wandb', 'tensorboard', 'csv', 'mlflow', 'litlogger'], default: csv)
logger_name: csv
# The random seed to use for reproducibility. (type: int, default: 1337)
seed: 1337
# Optimizer-related arguments
optimizer:
class_path: torch.optim.AdamW
init_args:
# (type: float, default: 0.001)
lr: 0.0002
# (type: float, default: 0.01)
weight_decay: 0.0
# (type: tuple, default: (0.9,0.999))
betas:
- 0.9
- 0.95
config_hub/pretrain/debug.yaml
+114
View File
@@ -0,0 +1,114 @@
# The name of the model to pretrain. Choose from names in ``litgpt.config``. Mutually exclusive with
# ``model_config``. (type: Optional[str], default: null)
model_name: pythia-14m
# A ``litgpt.Config`` object to define the model architecture. Mutually exclusive with
# ``model_config``. (type: Optional[Config], default: null)
model_config:
# Directory in which to save checkpoints and logs. If running in a Lightning Studio Job, look for it in
# /teamspace/jobs/<job-name>/share. (type: <class 'Path'>, default: out/pretrain)
out_dir: out/pretrain/debug
# The precision to use for pretraining. Possible choices: "bf16-true", "bf16-mixed", "32-true". (type: Optional[str], default: null)
precision: bf16-mixed
# Optional path to a checkpoint directory to initialize the model from.
# Useful for continued pretraining. Mutually exclusive with ``resume``. (type: Optional[Path], default: null)
initial_checkpoint_dir:
# Path to a checkpoint directory to resume from in case training was interrupted, or ``True`` to resume
# from the latest checkpoint in ``out_dir``. An error will be raised if no checkpoint is found. Passing
# ``'auto'`` will resume from the latest checkpoint but not error if no checkpoint exists.
# (type: Union[bool, Literal["auto"], Path], default: False)
resume: false
# Data-related arguments. If not provided, the default is ``litgpt.data.TinyLlama``.
data: TinyStories
# Training-related arguments. See ``litgpt.args.TrainArgs`` for details
train:
# Number of optimizer steps between saving checkpoints (type: Optional[int], default: 1000)
save_interval: 1000
# Number of iterations between logging calls (type: int, default: 1)
log_interval: 1
# Number of samples between optimizer steps across data-parallel ranks (type: int, default: 512)
global_batch_size: 125
# Number of samples per data-parallel rank (type: int, default: 4)
micro_batch_size: 5
# Number of iterations with learning rate warmup active (type: int, default: 2000)
lr_warmup_steps: 100
# Number of epochs to train on (type: Optional[int], default: null)
epochs:
# Total number of tokens to train on (type: Optional[int], default: 3000000000000)
max_tokens: 100000000
# Limits the number of optimizer steps to run. (type: Optional[int], default: null)
max_steps:
# Limits the length of samples. Off by default (type: Optional[int], default: null)
max_seq_length:
# Whether to tie the embedding weights with the language modeling head weights. (type: Optional[bool], default: False)
tie_embeddings:
# (type: Optional[float], default: 1.0)
max_norm: 1.0
# (type: float, default: 4e-05)
min_lr: 6e-5
# Evaluation-related arguments. See ``litgpt.args.EvalArgs`` for details
eval:
# Number of optimizer steps between evaluation calls (type: int, default: 1000)
interval: 1000
# Number of tokens to generate (type: Optional[int], default: null)
max_new_tokens:
# Number of iterations (type: int, default: 100)
max_iters: 100
# Whether to evaluate on the validation set at the beginning of the training
initial_validation: false
# Whether to evaluate on the validation set at the end the training
final_validation: false
# Optimizer-related arguments
optimizer:
class_path: torch.optim.AdamW
init_args:
# (type: float, default: 0.001)
lr: 6e-4
# (type: float, default: 0.01)
weight_decay: 0.1
# (type: tuple, default: (0.9,0.999))
betas:
- 0.9
- 0.95
# How many devices/GPUs to use. Uses all GPUs by default. (type: Union[int, str], default: auto)
devices: auto
# How many nodes to use. (type: int, default: 1)
num_nodes: 1
# Optional path to the tokenizer dir that was used for preprocessing the dataset. Only some data
# module require this. (type: Optional[Path], default: null)
tokenizer_dir: checkpoints/EleutherAI/pythia-14m
# The name of the logger to send metrics to. (type: LoggerChoice, i.e. Literal['wandb', 'tensorboard', 'csv', 'mlflow', 'litlogger'], default: tensorboard)
logger_name: tensorboard
# The random seed to use for reproducibility. (type: int, default: 42)
seed: 42
config_hub/pretrain/microllama.yaml
+115
View File
@@ -0,0 +1,115 @@
# The name of the model to pretrain. Choose from names in ``litgpt.config``. Mutually exclusive with
# ``model_config``. (type: Optional[str], default: null)
model_name: micro-llama-300M
# A ``litgpt.Config`` object to define the model architecture. Mutually exclusive with
# ``model_config``. (type: Optional[Config], default: null)
model_config:
# Directory in which to save checkpoints and logs. If running in a Lightning Studio Job, look for it in
# /teamspace/jobs/<job-name>/share. (type: <class 'Path'>, default: out/pretrain)
out_dir: out/pretrain/micro-llama
# The precision to use for pretraining. Possible choices: "bf16-true", "bf16-mixed", "32-true". (type: Optional[str], default: null)
precision: bf16-mixed
# Optional path to a checkpoint directory to initialize the model from.
# Useful for continued pretraining. Mutually exclusive with ``resume``. (type: Optional[Path], default: null)
initial_checkpoint_dir:
# Path to a checkpoint directory to resume from in case training was interrupted, or ``True`` to resume
# from the latest checkpoint in ``out_dir``. An error will be raised if no checkpoint is found. Passing
# ``'auto'`` will resume from the latest checkpoint but not error if no checkpoint exists.
# (type: Union[bool, Literal["auto"], Path], default: False)
resume: false
# Data-related arguments. If not provided, the default is ``litgpt.data.TinyLlama``.
data: MicroLlama
# Training-related arguments. See ``litgpt.args.TrainArgs`` for details
train:
# Number of optimizer steps between saving checkpoints (type: Optional[int], default: 1000)
save_interval: 1000
# Number of iterations between logging calls (type: int, default: 1)
log_interval: 1
# Number of samples between optimizer steps across data-parallel ranks (type: int, default: 48)
# Scale this number according to the number of GPU and memory size per GPU
# For example, we used 48 for 4 x 24G 4090
global_batch_size: 48
# Number of samples per data-parallel rank (type: int, default: 12)
# Scale this number according to the memory size per GPU
# For example, we used 12 for 24G 4090
micro_batch_size: 12
# Number of iterations with learning rate warmup active (type: int, default: 2000)
lr_warmup_steps: 2000
# Number of epochs to train on (type: Optional[int], default: null)
epochs:
# Total number of tokens to train on (type: Optional[int], default: 3000000000000)
max_tokens: 3000000000000
# Limits the number of optimizer steps to run. (type: Optional[int], default: null)
max_steps:
# Limits the length of samples. Off by default (type: Optional[int], default: null)
max_seq_length: 2048
# Whether to tie the embedding weights with the language modeling head weights. (type: Optional[bool], default: False)
tie_embeddings:
# (type: Optional[float], default: 1.0)
max_norm: 1.0
# (type: float, default: 4e-05)
min_lr: 4.0e-05
# Evaluation-related arguments. See ``litgpt.args.EvalArgs`` for details
eval:
# Number of optimizer steps between evaluation calls (type: int, default: 1000)
interval: 1000
# Number of tokens to generate (type: Optional[int], default: null)
max_new_tokens:
# Number of iterations (type: int, default: 100)
max_iters: 100
# Whether to evaluate on the validation set at the beginning of the training
initial_validation: false
# Optimizer-related arguments
optimizer:
class_path: torch.optim.AdamW
init_args:
# (type: float, default: 0.001)
lr: 4e-4
# (type: float, default: 0.01)
weight_decay: 0.1
# (type: tuple, default: (0.9,0.999))
betas:
- 0.9
- 0.95
# How many devices/GPUs to use. Uses all GPUs by default. (type: Union[int, str], default: auto)
devices: auto
# How many nodes to use. (type: int, default: 1)
num_nodes: 1
# Optional path to the tokenizer dir that was used for preprocessing the dataset. Only some data
# module require this. (type: Optional[Path], default: null)
tokenizer_dir: checkpoints/meta-llama/Llama-2-7b-hf
# The name of the logger to send metrics to. (type: LoggerChoice, i.e. Literal['wandb', 'tensorboard', 'csv', 'mlflow', 'litlogger'], default: tensorboard)
logger_name: tensorboard
# The random seed to use for reproducibility. (type: int, default: 42)
seed: 42
config_hub/pretrain/tinyllama.yaml
+114
View File
@@ -0,0 +1,114 @@
# The name of the model to pretrain. Choose from names in ``litgpt.config``. Mutually exclusive with
# ``model_config``. (type: Optional[str], default: null)
model_name: tiny-llama-1.1b
# A ``litgpt.Config`` object to define the model architecture. Mutually exclusive with
# ``model_config``. (type: Optional[Config], default: null)
model_config:
# Directory in which to save checkpoints and logs. If running in a Lightning Studio Job, look for it in
# /teamspace/jobs/<job-name>/share. (type: <class 'Path'>, default: out/pretrain)
out_dir: out/pretrain/tiny-llama
# The precision to use for pretraining. Possible choices: "bf16-true", "bf16-mixed", "32-true". (type: Optional[str], default: null)
precision: bf16-mixed
# Optional path to a checkpoint directory to initialize the model from.
# Useful for continued pretraining. Mutually exclusive with ``resume``. (type: Optional[Path], default: null)
initial_checkpoint_dir:
# Path to a checkpoint directory to resume from in case training was interrupted, or ``True`` to resume
# from the latest checkpoint in ``out_dir``. An error will be raised if no checkpoint is found. Passing
# ``'auto'`` will resume from the latest checkpoint but not error if no checkpoint exists.
# (type: Union[bool, Literal["auto"], Path], default: False)
resume: false
# Data-related arguments. If not provided, the default is ``litgpt.data.TinyLlama``.
data: TinyLlama
# Training-related arguments. See ``litgpt.args.TrainArgs`` for details
train:
# Number of optimizer steps between saving checkpoints (type: Optional[int], default: 1000)
save_interval: 1000
# Number of iterations between logging calls (type: int, default: 1)
log_interval: 1
# Number of samples between optimizer steps across data-parallel ranks (type: int, default: 512)
global_batch_size: 512
# Number of samples per data-parallel rank (type: int, default: 4)
micro_batch_size: 4
# Number of iterations with learning rate warmup active (type: int, default: 2000)
lr_warmup_steps: 2000
# Number of epochs to train on (type: Optional[int], default: null)
epochs:
# Total number of tokens to train on (type: Optional[int], default: 3000000000000)
max_tokens: 3000000000000
# Limits the number of optimizer steps to run. (type: Optional[int], default: null)
max_steps:
# Limits the length of samples. Off by default (type: Optional[int], default: null)
max_seq_length: 2048
# Whether to tie the embedding weights with the language modeling head weights. (type: Optional[bool], default: False)
tie_embeddings:
# (type: Optional[float], default: 1.0)
max_norm: 1.0
# (type: float, default: 4e-05)
min_lr: 4.0e-05
# Evaluation-related arguments. See ``litgpt.args.EvalArgs`` for details
eval:
# Number of optimizer steps between evaluation calls (type: int, default: 1000)
interval: 1000
# Number of tokens to generate (type: Optional[int], default: null)
max_new_tokens:
# Number of iterations (type: int, default: 100)
max_iters: 100
# Whether to evaluate on the validation set at the beginning of the training
initial_validation: false
# Whether to evaluate on the validation set at the end the training
final_validation: false
# Optimizer-related arguments
optimizer:
class_path: torch.optim.AdamW
init_args:
# (type: float, default: 0.001)
lr: 4e-4
# (type: float, default: 0.01)
weight_decay: 0.1
# (type: tuple, default: (0.9,0.999))
betas:
- 0.9
- 0.95
# How many devices/GPUs to use. Uses all GPUs by default. (type: Union[int, str], default: auto)
devices: auto
# How many nodes to use. (type: int, default: 1)
num_nodes: 1
# Optional path to the tokenizer dir that was used for preprocessing the dataset. Only some data
# module require this. (type: Optional[Path], default: null)
tokenizer_dir: checkpoints/meta-llama/Llama-2-7b-hf
# The name of the logger to send metrics to. (type: LoggerChoice, i.e. Literal['wandb', 'tensorboard', 'csv', 'mlflow', 'litlogger'], default: tensorboard)
logger_name: tensorboard
# The random seed to use for reproducibility. (type: int, default: 42)
seed: 42
config_hub/pretrain/tinystories.yaml
+130
View File
@@ -0,0 +1,130 @@
# The name of the model to pretrain. Choose from names in ``litgpt.config``. Mutually exclusive with
# ``model_config``. (type: Optional[str], default: null)
model_name: stories15M
# A ``litgpt.Config`` object to define the model architecture. Mutually exclusive with
# ``model_config``. (type: Optional[Config], default: null)
model_config:
name: stories15M
hf_config: {}
scale_embeddings: false
block_size: 256
padded_vocab_size: 32000
n_layer: 6
n_head: 6
n_query_groups: 6
n_embd: 288
head_size: 48
rotary_percentage: 1.0
parallel_residual: false
bias: false
norm_class_name: RMSNorm
mlp_class_name: LLaMAMLP
intermediate_size: 768
# Directory in which to save checkpoints and logs. If running in a Lightning Studio Job, look for it in
# /teamspace/jobs/<job-name>/share. (type: <class 'Path'>, default: out/pretrain)
out_dir: out/pretrain/stories15M
# The precision to use for pretraining. Possible choices: "bf16-true", "bf16-mixed", "32-true". (type: Optional[str], default: null)
precision: bf16-mixed
# Optional path to a checkpoint directory to initialize the model from.
# Useful for continued pretraining. Mutually exclusive with ``resume``. (type: Optional[Path], default: null)
initial_checkpoint_dir:
# Path to a checkpoint directory to resume from in case training was interrupted, or ``True`` to resume
# from the latest checkpoint in ``out_dir``. An error will be raised if no checkpoint is found. Passing
# ``'auto'`` will resume from the latest checkpoint but not error if no checkpoint exists.
# (type: Union[bool, Literal["auto"], Path], default: False)
resume: false
# Data-related arguments. If not provided, the default is ``litgpt.data.TinyLlama``.
data: TinyStories
# Training-related arguments. See ``litgpt.args.TrainArgs`` for details
train:
# Number of optimizer steps between saving checkpoints (type: Optional[int], default: 1000)
save_interval: 1000
# Number of iterations between logging calls (type: int, default: 1)
log_interval: 1
# Number of samples between optimizer steps across data-parallel ranks (type: int, default: 512)
global_batch_size: 512
# Number of samples per data-parallel rank (type: int, default: 4)
micro_batch_size: 128
# Number of iterations with learning rate warmup active (type: int, default: 2000)
lr_warmup_steps: 1000
# Number of epochs to train on (type: Optional[int], default: null)
epochs:
# Total number of tokens to train on (type: Optional[int], default: 3000000000000)
max_tokens: 9700000000 # original did 298,000 iters
# Limits the number of optimizer steps to run. (type: Optional[int], default: null)
max_steps:
# Limits the length of samples. Off by default (type: Optional[int], default: null)
max_seq_length: 256
# Whether to tie the embedding weights with the language modeling head weights. (type: Optional[bool], default: False)
tie_embeddings: true
# (type: Optional[float], default: 1.0)
max_norm: 1.0
# (type: float, default: 4e-05)
min_lr: 0.0
# Evaluation-related arguments. See ``litgpt.args.EvalArgs`` for details
eval:
# Number of optimizer steps between evaluation calls (type: int, default: 1000)
interval: 2000
# Number of tokens to generate (type: Optional[int], default: null)
max_new_tokens:
# Number of iterations (type: int, default: 100)
max_iters: 100
# Whether to evaluate on the validation set at the beginning of the training
initial_validation: false
# Whether to evaluate on the validation set at the end the training
final_validation: false
# Optimizer-related arguments
optimizer:
class_path: torch.optim.AdamW
init_args:
# (type: float, default: 0.001)
lr: 0.0005
# (type: float, default: 0.01)
weight_decay: 0.1
# (type: tuple, default: (0.9,0.999))
betas:
- 0.9
- 0.95
# How many devices/GPUs to use. Uses all GPUs by default. (type: Union[int, str], default: auto)
devices: auto
# How many nodes to use. (type: int, default: 1)
num_nodes: 1
# Optional path to the tokenizer dir that was used for preprocessing the dataset. Only some data
# module require this. (type: Optional[Path], default: null)
tokenizer_dir: checkpoints/meta-llama/Llama-2-7b-hf
# The name of the logger to send metrics to. (type: LoggerChoice, i.e. Literal['wandb', 'tensorboard', 'csv', 'mlflow', 'litlogger'], default: tensorboard)
logger_name: csv
# The random seed to use for reproducibility. (type: int, default: 42)
seed: 42
extensions/thunder/README.md
+617
View File
@@ -0,0 +1,617 @@
# Lightning Thunder: a source-to-source compiler for PyTorch
[Lightning Thunder](https://github.com/Lightning-AI/lightning-thunder) makes PyTorch programs faster both on single accelerators or in distributed settings.
Thunder aims to be usable, understandable, and extensible and can achieve significant speedups over standard PyTorch eager code, through the compounding effects of optimizations and the use of best in class executors.
This extension directory shows how Thunder can be used with LitGPT.
> [!WARNING]
> This document is an early-access development version that is currently only for internal use. We recommend users checking out the [Lightning Thunder](https://github.com/Lightning-AI/lightning-thunder) project directly, which provides more up-to-date usage information.
&nbsp;
## Thunder 👉👈 LitGPT: a short showcase
To try Lightning Thunder with your model simply `thunder.jit()` it.
```python
from litgpt import GPT
import thunder
import torch
# Use only two layers to keep the traces shorter for the demonstration
model = GPT.from_name("Llama-2-7b-hf", n_layer=2).cuda()
model = thunder.jit(model)
x = torch.randint(model.max_seq_length, (2, 5), device="cuda")
y = model(x) # forward, this may take a bit
```
This will require some compilation time on the first forward call.
### Traces
The JIT is will acquire a Python program (what we call a "trace") from the Python program (`GPT`, a `torch.nn.Module` in this example) that was given.
This process targets PyTorch operators (like `Tensor.view()`, `+`, `torch.nn.functional.scaled_dot_product_atttention()`) and optionally custom operators (more about that later).
We can visualize the thunder trace generated under the hood:
```python
forward_trace = thunder.last_traces(model)[-1].python()
print(forward_trace)
```
```python
@torch.no_grad()
@no_autocast()
def augmented_forward_fn(*args):
# args: "Collection"
t0, t1, t2, t3, t4, t5, t6, t7, t8, t9, t10, t11, t12, t13, t14, t15, t16, t17, \
t18, t19, = args
del args
t24 = torch.nn.functional.embedding(t0, t19, None, None, 2.0, False, False) # t24: "cuda:0 f32[2, 5, 4096]"
t20 = torch_slice_prim_impl(t1, [0, 0], [5, 128], [1, 1]) # t20: "cuda:0 f32[5, 128]"
t21 = torch_slice_prim_impl(t2, [0, 0], [5, 128], [1, 1]) # t21: "cuda:0 f32[5, 128]"
t200 = torch.unsqueeze(t11, 0) # t200: "cuda:0 f32[1, 4096]"
t201 = torch.unsqueeze(t200, 1) # t201: "cuda:0 f32[1, 1, 4096]"
del t200
t33 = Tensor.expand(t201, (2, 5, 4096)) # t33: "cuda:0 f32[2, 5, 4096]"
del t201
t229 = torch.unsqueeze(t13, 0) # t229: "cuda:0 f32[1, 4096]"
t230 = torch.unsqueeze(t229, 1) # t230: "cuda:0 f32[1, 1, 4096]"
del t229
t84 = Tensor.expand(t230, (2, 5, 4096)) # t84: "cuda:0 f32[2, 5, 4096]"
del t230
t232 = torch.unsqueeze(t12, 0) # t232: "cuda:0 f32[1, 4096]"
t233 = torch.unsqueeze(t232, 1) # t233: "cuda:0 f32[1, 1, 4096]"
del t232
t104 = Tensor.expand(t233, (2, 5, 4096)) # t104: "cuda:0 f32[2, 5, 4096]"
del t233
t253 = torch.unsqueeze(t14, 0) # t253: "cuda:0 f32[1, 4096]"
t254 = torch.unsqueeze(t253, 1) # t254: "cuda:0 f32[1, 1, 4096]"
del t253
t155 = Tensor.expand(t254, (2, 5, 4096)) # t155: "cuda:0 f32[2, 5, 4096]"
del t254
t256 = torch.unsqueeze(t10, 0) # t256: "cuda:0 f32[1, 4096]"
t257 = torch.unsqueeze(t256, 1) # t257: "cuda:0 f32[1, 1, 4096]"
del t256
t175 = Tensor.expand(t257, (2, 5, 4096)) # t175: "cuda:0 f32[2, 5, 4096]"
del t257
t221 = torch.unsqueeze(t20, 0) # t221: "cuda:0 f32[1, 5, 128]"
del t20
t222 = torch.unsqueeze(t221, 1) # t222: "cuda:0 f32[1, 1, 5, 128]"
del t221
t49 = Tensor.expand(t222, (2, 32, 5, 128)) # t49: "cuda:0 f32[2, 32, 5, 128]"
del t222
t224 = torch.unsqueeze(t21, 0) # t224: "cuda:0 f32[1, 5, 128]"
del t21
t225 = torch.unsqueeze(t224, 1) # t225: "cuda:0 f32[1, 1, 5, 128]"
del t224
t51 = Tensor.expand(t225, (2, 32, 5, 128)) # t51: "cuda:0 f32[2, 32, 5, 128]"
del t225
[t30, t34] = nvFusion0(t24, t33)
t35 = torch.nn.functional.linear(t34, t3, None) # t35: "cuda:0 f32[2, 5, 12288]"
t36 = torch.reshape(t35, (2, 5, 32, 3, 128)) # t36: "cuda:0 f32[2, 5, 32, 3, 128]"
del t35
t37 = torch.permute(t36, (0, 2, 3, 1, 4)) # t37: "cuda:0 f32[2, 32, 3, 5, 128]"
del t36
(t38, t39, t40) = torch.split(t37, (1, 1, 1), 2)
del t37
t41 = torch.reshape(t38, (2, 32, 5, 128)) # t41: "cuda:0 f32[2, 32, 5, 128]"
del t38
t42 = torch.reshape(t39, (2, 32, 5, 128)) # t42: "cuda:0 f32[2, 32, 5, 128]"
del t39
t43 = torch.reshape(t40, (2, 32, 5, 128)) # t43: "cuda:0 f32[2, 32, 5, 128]"
del t40
t44 = torch_slice_prim_impl(t41, [0, 0, 0, 0], [2, 32, 5, 128], [1, 1, 1, 1]) # t44: "cuda:0 f32[2, 32, 5, 128]"
t54 = torch_slice_prim_impl(t42, [0, 0, 0, 0], [2, 32, 5, 128], [1, 1, 1, 1]) # t54: "cuda:0 f32[2, 32, 5, 128]"
t64 = torch_slice_prim_impl(t41, [0, 0, 0, 0], [2, 32, 5, 0], [1, 1, 1, 1]) # t64: "cuda:0 f32[2, 32, 5, 0]"
del t41
t66 = torch_slice_prim_impl(t42, [0, 0, 0, 0], [2, 32, 5, 0], [1, 1, 1, 1]) # t66: "cuda:0 f32[2, 32, 5, 0]"
del t42
t46 = torch_slice_prim_impl(t44, [0, 0, 0, 64], [2, 32, 5, 128], [1, 1, 1, 1]) # t46: "cuda:0 f32[2, 32, 5, 64]"
t45 = torch_slice_prim_impl(t44, [0, 0, 0, 0], [2, 32, 5, 64], [1, 1, 1, 1]) # t45: "cuda:0 f32[2, 32, 5, 64]"
t55 = torch_slice_prim_impl(t54, [0, 0, 0, 0], [2, 32, 5, 64], [1, 1, 1, 1]) # t55: "cuda:0 f32[2, 32, 5, 64]"
t56 = torch_slice_prim_impl(t54, [0, 0, 0, 64], [2, 32, 5, 128], [1, 1, 1, 1]) # t56: "cuda:0 f32[2, 32, 5, 64]"
[t47, t57] = nvFusion1(t46, t56)
del t46, t56
t48 = torch.cat((t47, t45), -1) # t48: "cuda:0 f32[2, 32, 5, 128]"
del t47, t45
t58 = torch.cat((t57, t55), -1) # t58: "cuda:0 f32[2, 32, 5, 128]"
del t57, t55
[t53, t63] = nvFusion2(t44, t48, t49, t51, t54, t58)
del t44, t48, t54, t58
t65 = torch.cat((t53, t64), -1) # t65: "cuda:0 f32[2, 32, 5, 128]"
del t53, t64
t67 = torch.cat((t63, t66), -1) # t67: "cuda:0 f32[2, 32, 5, 128]"
del t63, t66
(t68, t69, t70, t71) = sdpaex_grad_forward_scaled_dot_product_efficient_attention(t65, t67, t43, None, 0.0, True, 0.08838834764831843)
t72 = torch.permute(t68, (0, 2, 1, 3)) # t72: "cuda:0 f32[2, 5, 32, 128]"
t73 = torch.reshape(t72, (2, 5, 4096)) # t73: "cuda:0 f32[2, 5, 4096]"
del t72
t74 = torch.nn.functional.linear(t73, t15, None) # t74: "cuda:0 f32[2, 5, 4096]"
[t75, t81, t85] = nvFusion3(t24, t74, t84)
del t74
t86 = torch.nn.functional.linear(t85, t5, None) # t86: "cuda:0 f32[2, 5, 11008]"
t87 = torch.nn.functional.linear(t85, t7, None) # t87: "cuda:0 f32[2, 5, 11008]"
[t93] = nvFusion4(t86, t87)
t94 = torch.nn.functional.linear(t93, t16, None) # t94: "cuda:0 f32[2, 5, 4096]"
[t101, t105, t95] = nvFusion5(t104, t75, t94)
del t94
t106 = torch.nn.functional.linear(t105, t4, None) # t106: "cuda:0 f32[2, 5, 12288]"
t107 = torch.reshape(t106, (2, 5, 32, 3, 128)) # t107: "cuda:0 f32[2, 5, 32, 3, 128]"
del t106
t108 = torch.permute(t107, (0, 2, 3, 1, 4)) # t108: "cuda:0 f32[2, 32, 3, 5, 128]"
del t107
(t109, t110, t111) = torch.split(t108, (1, 1, 1), 2)
del t108
t112 = torch.reshape(t109, (2, 32, 5, 128)) # t112: "cuda:0 f32[2, 32, 5, 128]"
del t109
t113 = torch.reshape(t110, (2, 32, 5, 128)) # t113: "cuda:0 f32[2, 32, 5, 128]"
del t110
t114 = torch.reshape(t111, (2, 32, 5, 128)) # t114: "cuda:0 f32[2, 32, 5, 128]"
del t111
t135 = torch_slice_prim_impl(t112, [0, 0, 0, 0], [2, 32, 5, 0], [1, 1, 1, 1]) # t135: "cuda:0 f32[2, 32, 5, 0]"
t137 = torch_slice_prim_impl(t113, [0, 0, 0, 0], [2, 32, 5, 0], [1, 1, 1, 1]) # t137: "cuda:0 f32[2, 32, 5, 0]"
t115 = torch_slice_prim_impl(t112, [0, 0, 0, 0], [2, 32, 5, 128], [1, 1, 1, 1]) # t115: "cuda:0 f32[2, 32, 5, 128]"
del t112
t125 = torch_slice_prim_impl(t113, [0, 0, 0, 0], [2, 32, 5, 128], [1, 1, 1, 1]) # t125: "cuda:0 f32[2, 32, 5, 128]"
del t113
t116 = torch_slice_prim_impl(t115, [0, 0, 0, 0], [2, 32, 5, 64], [1, 1, 1, 1]) # t116: "cuda:0 f32[2, 32, 5, 64]"
t117 = torch_slice_prim_impl(t115, [0, 0, 0, 64], [2, 32, 5, 128], [1, 1, 1, 1]) # t117: "cuda:0 f32[2, 32, 5, 64]"
t127 = torch_slice_prim_impl(t125, [0, 0, 0, 64], [2, 32, 5, 128], [1, 1, 1, 1]) # t127: "cuda:0 f32[2, 32, 5, 64]"
t126 = torch_slice_prim_impl(t125, [0, 0, 0, 0], [2, 32, 5, 64], [1, 1, 1, 1]) # t126: "cuda:0 f32[2, 32, 5, 64]"
[t118, t128] = nvFusion6(t117, t127)
del t117, t127
t129 = torch.cat((t128, t126), -1) # t129: "cuda:0 f32[2, 32, 5, 128]"
del t128, t126
t119 = torch.cat((t118, t116), -1) # t119: "cuda:0 f32[2, 32, 5, 128]"
del t118, t116
[t124, t134] = nvFusion7(t115, t119, t125, t129, t49, t51)
del t115, t119, t125, t129
t136 = torch.cat((t124, t135), -1) # t136: "cuda:0 f32[2, 32, 5, 128]"
del t124, t135
t138 = torch.cat((t134, t137), -1) # t138: "cuda:0 f32[2, 32, 5, 128]"
del t134, t137
(t139, t140, t141, t142) = sdpaex_grad_forward_scaled_dot_product_efficient_attention(t136, t138, t114, None, 0.0, True, 0.08838834764831843)
t143 = torch.permute(t139, (0, 2, 1, 3)) # t143: "cuda:0 f32[2, 5, 32, 128]"
t144 = torch.reshape(t143, (2, 5, 4096)) # t144: "cuda:0 f32[2, 5, 4096]"
del t143
t145 = torch.nn.functional.linear(t144, t17, None) # t145: "cuda:0 f32[2, 5, 4096]"
[t146, t152, t156] = nvFusion8(t145, t155, t95)
del t145
t158 = torch.nn.functional.linear(t156, t8, None) # t158: "cuda:0 f32[2, 5, 11008]"
t157 = torch.nn.functional.linear(t156, t6, None) # t157: "cuda:0 f32[2, 5, 11008]"
[t164] = nvFusion9(t157, t158)
t165 = torch.nn.functional.linear(t164, t18, None) # t165: "cuda:0 f32[2, 5, 4096]"
[t166, t172, t176] = nvFusion10(t146, t165, t175)
del t165
t177 = torch.nn.functional.linear(t176, t9, None) # t177: "cuda:0 f32[2, 5, 32000]"
return {'output': t177, 'flat_args': [t0, t1, t2, t3, t4, t5, t6, t7, t8, t9, t10, t11, t12, t13, t14, t15, t16, t17, t18, t19], 'flat_output': (t177,)}, ((t0, t101, t104, t105, t114, t136, t138, t139, t140, t141, t142, t144, t146, t15, t152, t155, t156, t157, t158, t16, t164, t166, t17, t172, t175, t176, t18, t24, t3, t30, t33, t34, t4, t43, t49, t5, t51, t6, t65, t67, t68, t69, t7, t70, t71, t73, t75, t8, t81, t84, t85, t86, t87, t9, t93, t95), (False, False, True, True, 4096.0, 4096.0, 0.0, 0.08838834764831843, 4096.0, 4096.0, 4096.0, 0.0, 0.08838834764831843, 32000, 2, 2))
```
This is a straight-lined version of `GPT.forward` that has been optimized. Since it's running on CUDA, the [NvFuser](https://github.com/NVIDIA/Fuser) executor has created regions (look for "nvFusion") that fuse multiple operators together.
Operator fusion is very desirable with modern hardware and helps out in overhead-bound or device-bound settings by:
- Launching less kernels, thus reducing the kernel launch overhead.
- Reducing the number of memory accesses performed by reusing them in a fused operation
- Minimizing host-device communications
Thunder also uses a multi-level intermediate representation. If we let it print all levels
```python
forward_trace = thunder.last_traces(model)[-1]
print(forward_trace)
```
We can see as comments the primitives that compose the fusion regions. For instance, this is the region associated to [the `RMSNorm` implementation](https://github.com/Lightning-AI/litgpt/blob/9b6475dabf90c7acee506a026bd9fa86251835bf/litgpt/model.py#L409-L420)
```python
[t146, t152, t156] = nvFusion8(t145, t155, t95)
# t146 = prims.add(t145, t95) # t146: "cuda:0 f32[2, 5, 4096]"
# t147 = prims.mul(t146, t146) # t147: "cuda:0 f32[2, 5, 4096]"
# t148 = prims.sum(t147, (2,)) # t148: "cuda:0 f32[2, 5]"
# t149 = prims.broadcast_in_dim(t148, [2, 5, 1], [0, 1]) # t149: "cuda:0 f32[2, 5, 1]"
# t150 = prims.div(t149, 4096.0) # t150: "cuda:0 f32[2, 5, 1]"
# t151 = prims.add(t150, 1e-05) # t151: "cuda:0 f32[2, 5, 1]"
# t152 = prims.rsqrt(t151) # t152: "cuda:0 f32[2, 5, 1]"
# t153 = prims.broadcast_in_dim(t152, (2, 5, 4096), (0, 1, 2)) # t153: "cuda:0 f32[2, 5, 4096]"
# t154 = prims.mul(t146, t153) # t154: "cuda:0 f32[2, 5, 4096]"
# t156 = prims.mul(t154, t155) # t156: "cuda:0 f32[2, 5, 4096]"
```
Similarly, we can visualize the backward trace:
```python
backward_trace = thunder.last_backward_traces(model)[-1].python()
print(backward_trace)
```
```python
@torch.no_grad()
@no_autocast()
def backward_fn(saved_for_backward, cotangents):
# saved_for_backward: "Collection"
# cotangents: "Collection"
C0, C1, = saved_for_backward
clear_collection(saved_for_backward)
del saved_for_backward
t178, = cotangents
clear_collection(cotangents)
del cotangents
t0, t101, t104, t105, t114, t136, t138, t139, t140, t141, t142, t144, t146, \
t15, t152, t155, t156, t157, t158, t16, t164, t166, t17, t172, t175, t176, t18, \
t24, t3, t30, t33, t34, t4, t43, t49, t5, t51, t6, t65, t67, t68, t69, t7, t70, \
t71, t73, t75, t8, t81, t84, t85, t86, t87, t9, t93, t95, = C0
clear_collection(C0)
del C0
b1, b2, b41, b91, f101, f106, f40, f42, f51, f56, f6, f90, f92, i0, i23, i73, \
= C1
clear_collection(C1)
del C1
t639 = torch.reshape(t178, (-1, 32000)) # t639: "cuda:0 f32[10, 32000]"
del t178
t643 = torch.permute(t639, (1, 0)) # t643: "cuda:0 f32[32000, 10]"
t644 = torch.reshape(t176, (-1, 4096)) # t644: "cuda:0 f32[10, 4096]"
del t176
t669 = torch.reshape(t164, (-1, 11008)) # t669: "cuda:0 f32[10, 11008]"
del t164
t686 = torch.reshape(t156, (-1, 4096)) # t686: "cuda:0 f32[10, 4096]"
del t156
t720 = torch.reshape(t144, (-1, 4096)) # t720: "cuda:0 f32[10, 4096]"
del t144
t776 = torch.reshape(t105, (-1, 4096)) # t776: "cuda:0 f32[10, 4096]"
del t105
t802 = torch.reshape(t93, (-1, 11008)) # t802: "cuda:0 f32[10, 11008]"
del t93
t819 = torch.reshape(t85, (-1, 4096)) # t819: "cuda:0 f32[10, 4096]"
del t85
t853 = torch.reshape(t73, (-1, 4096)) # t853: "cuda:0 f32[10, 4096]"
del t73
t911 = torch.reshape(t34, (-1, 4096)) # t911: "cuda:0 f32[10, 4096]"
del t34
t640 = torch.matmul(t639, t9) # t640: "cuda:0 f32[10, 4096]"
del t639, t9
t645 = torch.matmul(t643, t644) # t645: "cuda:0 f32[32000, 4096]"
del t643, t644
t641 = torch.reshape(t640, (2, 5, 4096)) # t641: "cuda:0 f32[2, 5, 4096]"
del t640
[t648, t663] = nvFusion0(f106, t166, t172, t175, t641)
del f106, t166, t172, t175, t641
t664 = torch.reshape(t663, (-1, 4096)) # t664: "cuda:0 f32[10, 4096]"
t668 = torch.permute(t664, (1, 0)) # t668: "cuda:0 f32[4096, 10]"
t665 = torch.matmul(t664, t18) # t665: "cuda:0 f32[10, 11008]"
del t664, t18
t670 = torch.matmul(t668, t669) # t670: "cuda:0 f32[4096, 11008]"
del t668, t669
t666 = torch.reshape(t665, (2, 5, 11008)) # t666: "cuda:0 f32[2, 5, 11008]"
del t665
[t672, t680] = nvFusion1(t157, t158, t666)
del t157, t158, t666
t681 = torch.reshape(t672, (-1, 11008)) # t681: "cuda:0 f32[10, 11008]"
del t672
t685 = torch.permute(t681, (1, 0)) # t685: "cuda:0 f32[11008, 10]"
t688 = torch.reshape(t680, (-1, 11008)) # t688: "cuda:0 f32[10, 11008]"
del t680
t692 = torch.permute(t688, (1, 0)) # t692: "cuda:0 f32[11008, 10]"
t689 = torch.matmul(t688, t6) # t689: "cuda:0 f32[10, 4096]"
del t688, t6
t682 = torch.matmul(t681, t8) # t682: "cuda:0 f32[10, 4096]"
del t681, t8
t694 = torch.matmul(t692, t686) # t694: "cuda:0 f32[11008, 4096]"
del t692
t687 = torch.matmul(t685, t686) # t687: "cuda:0 f32[11008, 4096]"
del t685, t686
t683 = torch.reshape(t682, (2, 5, 4096)) # t683: "cuda:0 f32[2, 5, 4096]"
del t682
t690 = torch.reshape(t689, (2, 5, 4096)) # t690: "cuda:0 f32[2, 5, 4096]"
del t689
[t698, t714] = nvFusion2(f101, t146, t152, t155, t663, t683, t690)
del f101, t146, t152, t155, t663, t683, t690
t715 = torch.reshape(t714, (-1, 4096)) # t715: "cuda:0 f32[10, 4096]"
t719 = torch.permute(t715, (1, 0)) # t719: "cuda:0 f32[4096, 10]"
t716 = torch.matmul(t715, t17) # t716: "cuda:0 f32[10, 4096]"
del t715, t17
t721 = torch.matmul(t719, t720) # t721: "cuda:0 f32[4096, 4096]"
del t719, t720
t717 = torch.reshape(t716, (2, 5, 4096)) # t717: "cuda:0 f32[2, 5, 4096]"
del t716
t722 = torch.reshape(t717, (2, 5, 32, 128)) # t722: "cuda:0 f32[2, 5, 32, 128]"
del t717
t723 = torch.permute(t722, (0, 2, 1, 3)) # t723: "cuda:0 f32[2, 32, 5, 128]"
del t722
(t724, t725, t726, _) = sdpaex_scaled_dot_product_efficient_attention_backward(t723, t136, t138, t114, None, t139, t140, t141, t142, f90, b91, scale=f92)
del t723, t136, t138, t114, t139, t140, t141, t142, f90, b91, f92
t765 = torch.reshape(t726, (2, 32, 1, 5, 128)) # t765: "cuda:0 f32[2, 32, 1, 5, 128]"
del t726
t727 = torch_slice_prim_impl(t725, [0, 0, 0, 0], [2, 32, 5, 128], [1, 1, 1, 1]) # t727: "cuda:0 f32[2, 32, 5, 128]"
del t725
t730 = torch_slice_prim_impl(t724, [0, 0, 0, 0], [2, 32, 5, 128], [1, 1, 1, 1]) # t730: "cuda:0 f32[2, 32, 5, 128]"
del t724
[t747, t764] = nvFusion3(t49, t51, t727, t730)
del t727, t730
t766 = torch.reshape(t747, (2, 32, 1, 5, 128)) # t766: "cuda:0 f32[2, 32, 1, 5, 128]"
del t747
t767 = torch.reshape(t764, (2, 32, 1, 5, 128)) # t767: "cuda:0 f32[2, 32, 1, 5, 128]"
del t764
t768 = torch.cat((t767, t766, t765), i73) # t768: "cuda:0 f32[2, 32, 3, 5, 128]"
del t767, t766, t765, i73
t769 = torch.permute(t768, (0, 3, 1, 2, 4)) # t769: "cuda:0 f32[2, 5, 32, 3, 128]"
del t768
t770 = torch.reshape(t769, (2, 5, 12288)) # t770: "cuda:0 f32[2, 5, 12288]"
del t769
t771 = torch.reshape(t770, (-1, 12288)) # t771: "cuda:0 f32[10, 12288]"
del t770
t775 = torch.permute(t771, (1, 0)) # t775: "cuda:0 f32[12288, 10]"
t777 = torch.matmul(t775, t776) # t777: "cuda:0 f32[12288, 4096]"
del t775, t776
t772 = torch.matmul(t771, t4) # t772: "cuda:0 f32[10, 4096]"
del t771, t4
t773 = torch.reshape(t772, (2, 5, 4096)) # t773: "cuda:0 f32[2, 5, 4096]"
del t772
[t780, t796] = nvFusion4(f56, t101, t104, t714, t773, t95)
del f56, t101, t104, t714, t773, t95
t797 = torch.reshape(t796, (-1, 4096)) # t797: "cuda:0 f32[10, 4096]"
t801 = torch.permute(t797, (1, 0)) # t801: "cuda:0 f32[4096, 10]"
t798 = torch.matmul(t797, t16) # t798: "cuda:0 f32[10, 11008]"
del t797, t16
t803 = torch.matmul(t801, t802) # t803: "cuda:0 f32[4096, 11008]"
del t801, t802
t799 = torch.reshape(t798, (2, 5, 11008)) # t799: "cuda:0 f32[2, 5, 11008]"
del t798
[t805, t813] = nvFusion5(t799, t86, t87)
del t799, t86, t87
t814 = torch.reshape(t805, (-1, 11008)) # t814: "cuda:0 f32[10, 11008]"
del t805
t818 = torch.permute(t814, (1, 0)) # t818: "cuda:0 f32[11008, 10]"
t821 = torch.reshape(t813, (-1, 11008)) # t821: "cuda:0 f32[10, 11008]"
del t813
t825 = torch.permute(t821, (1, 0)) # t825: "cuda:0 f32[11008, 10]"
t822 = torch.matmul(t821, t5) # t822: "cuda:0 f32[10, 4096]"
del t821, t5
t815 = torch.matmul(t814, t7) # t815: "cuda:0 f32[10, 4096]"
del t814, t7
t827 = torch.matmul(t825, t819) # t827: "cuda:0 f32[11008, 4096]"
del t825
t820 = torch.matmul(t818, t819) # t820: "cuda:0 f32[11008, 4096]"
del t818, t819
t816 = torch.reshape(t815, (2, 5, 4096)) # t816: "cuda:0 f32[2, 5, 4096]"
del t815
t823 = torch.reshape(t822, (2, 5, 4096)) # t823: "cuda:0 f32[2, 5, 4096]"
del t822
[t831, t847] = nvFusion6(f51, t75, t796, t81, t816, t823, t84)
del f51, t75, t796, t81, t816, t823, t84
t848 = torch.reshape(t847, (-1, 4096)) # t848: "cuda:0 f32[10, 4096]"
t852 = torch.permute(t848, (1, 0)) # t852: "cuda:0 f32[4096, 10]"
t849 = torch.matmul(t848, t15) # t849: "cuda:0 f32[10, 4096]"
del t848, t15
t854 = torch.matmul(t852, t853) # t854: "cuda:0 f32[4096, 4096]"
del t852, t853
t850 = torch.reshape(t849, (2, 5, 4096)) # t850: "cuda:0 f32[2, 5, 4096]"
del t849
t855 = torch.reshape(t850, (2, 5, 32, 128)) # t855: "cuda:0 f32[2, 5, 32, 128]"
del t850
t856 = torch.permute(t855, (0, 2, 1, 3)) # t856: "cuda:0 f32[2, 32, 5, 128]"
del t855
(t857, t858, t859, _) = sdpaex_scaled_dot_product_efficient_attention_backward(t856, t65, t67, t43, None, t68, t69, t70, t71, f40, b41, scale=f42)
del t856, t65, t67, t43, t68, t69, t70, t71, f40, b41, f42
t900 = torch.reshape(t859, (2, 32, 1, 5, 128)) # t900: "cuda:0 f32[2, 32, 1, 5, 128]"
del t859
t863 = torch_slice_prim_impl(t857, [0, 0, 0, 0], [2, 32, 5, 128], [1, 1, 1, 1]) # t863: "cuda:0 f32[2, 32, 5, 128]"
del t857
t860 = torch_slice_prim_impl(t858, [0, 0, 0, 0], [2, 32, 5, 128], [1, 1, 1, 1]) # t860: "cuda:0 f32[2, 32, 5, 128]"
del t858
[t882, t899] = nvFusion7(t49, t51, t860, t863)
del t49, t51, t860, t863
t902 = torch.reshape(t899, (2, 32, 1, 5, 128)) # t902: "cuda:0 f32[2, 32, 1, 5, 128]"
del t899
t901 = torch.reshape(t882, (2, 32, 1, 5, 128)) # t901: "cuda:0 f32[2, 32, 1, 5, 128]"
del t882
t903 = torch.cat((t902, t901, t900), i23) # t903: "cuda:0 f32[2, 32, 3, 5, 128]"
del t902, t901, t900, i23
t904 = torch.permute(t903, (0, 3, 1, 2, 4)) # t904: "cuda:0 f32[2, 5, 32, 3, 128]"
del t903
t905 = torch.reshape(t904, (2, 5, 12288)) # t905: "cuda:0 f32[2, 5, 12288]"
del t904
t906 = torch.reshape(t905, (-1, 12288)) # t906: "cuda:0 f32[10, 12288]"
del t905
t910 = torch.permute(t906, (1, 0)) # t910: "cuda:0 f32[12288, 10]"
t907 = torch.matmul(t906, t3) # t907: "cuda:0 f32[10, 4096]"
del t906, t3
t912 = torch.matmul(t910, t911) # t912: "cuda:0 f32[12288, 4096]"
del t910, t911
t908 = torch.reshape(t907, (2, 5, 4096)) # t908: "cuda:0 f32[2, 5, 4096]"
del t907
[t915, t931] = nvFusion8(f6, t24, t30, t33, t847, t908)
del f6, t24, t30, t33, t847, t908
t932 = torch.torch.ops.aten.embedding_backward(t931, t0, i0, -1, b1, b2) # t932: "cuda:0 f32[32000, 4096]"
del t931, t0, i0, b1, b2
return (None, None, None, t912, t777, t827, t694, t820, t687, t645, t648, t915, t780, t831, t698, t854, t803, t721, t670, t932)
```
These traces are long, and require some familiarity with the model implementation to follow them, but they allow you to:
- Inspect exactly what operations are run including their decompositions.
- Inspect the sizes of tensors, their device, data type and conversions.
- Apply transformations to the traces since the computations are completely decoupled from the data.
- Inspect the backward operations generated for each forward operation to understand what autograd is doing.
### Transforms
Transforms are one of the core features of Thunder. For example, they enable easy data parallel distribution. That is replicated data parallelism (DDP) and fully-sharded data parallelism (FSDP).
We provide ready-to-use Fabric strategies that integrate Thunder DDP|FSDP. Under the hood, the code is quite straightforward:
```python
model = thunder.distributed.ddp(model)
# or
# model = thunder.distributed.fsdp(model)
model = thunder.jit(model)
```
After applying the DDP transformation, the backward trace will include the expected all-reduce collectives:
```python
p1022 = torch_all_reduce_prim_impl(t1021, _DistributedReduceOps_0, _torch_distributed_distributed_c10d_ProcessGroup_1, True, False) # p1022: "FUTURE cuda:0 f32[16797696]"
...
t1059 = torch_wait_prim_impl(p1025) # t1059: "cuda:0 f32[131072000]"
```
With `L.Fabric`, this is how to use them:
```python
from extensions.extensions.thunder.strategies import ThunderFSDPStrategy, ThunderDDPStrategy
# fully-sharded data parallel
strategy = ThunderFSDPStrategy(
sharding_strategy="ZERO3",
bucketing_strategy="BLOCK",
executors=("sdpa", "torchcompile_cat", "nvfuser", "torch"),
state_dict_type="full",
)
# replicated data parallel
strategy = ThunderDDPStrategy(executors=("sdpa", "torchcompile_cat", "nvfuser", "torch"))
fabric = L.Fabric(devices=devices, strategy=strategy)
fabric.launch()
model = fabric.setup(model) # JIT is called here
```
And in the case of FSDP all-gathers in forward and reduce-scatters in backward.
Meaning that Thunder automatically introduced the necessary collective operations to support data parallelism.
### Executors
Thunder allows you to define a priority list of executors that can map operators:
```python
import thunder
model = thunder.jit(
model,
executors=["sdpa", "torchcompile_cat", "nvfuser", "torch"]
)
```
Notice how `torch.compile` is a valid executor. This executor registers a few operators with improved performance so that you can utilize the fastest set of operator implementations possible.
### Custom executors
Lightning Thunder provides extension points to integrate fast kernels for operators in your model without having to modify your implementation.
For instance, the [Unsloth project](https://github.com/unslothai/unsloth/) provides several Triton kernels that can be used with LitGPT:
- Cross entropy loss
- SwiGLU (part of `LLaMAMLP`)
- RoPE
The [`unsloth` directory](unsloth) contains a [custom executor](unsloth/executor.py) that registers these operators for LitGPT.
We can enable this executor by passing it to the list of executors available. The order matters because we want to run its custom operators before
`NvFuser` creates its fusion regions.
```python
import thunder
model = thunder.jit(
model,
executors=["sdpa", "unsloth", "torchcompile_cat", "nvfuser", "torch"]
)
```
Doing this, the model trace now includes the Unsloth kernel calls:
```python
def augmented_forward_fn(*args):
...
(t121, _, _, _, _, _) = unsloth_apply_rope(t120, t21, t22)
...
(t189, t190) = unsloth_cross_entropy(t187, t188)
...
def backward_fn(saved_for_backward, cotangents):
...
t652 = unsloth_cross_entropy_backward(t651, t187, t188, t190) # t652: "cuda:0 f32[6, 320]"
...
t763 = unsloth_apply_rope_backward(t757, t21, t22, 1, 8, 4) # t763: "cuda:0 f32[2, 4, 3, 16]"
```
We provide a specific [pre-training script copy](pretrain.py) that uses this executor.
Given the Unsloth results below, these hand-written kernels do not seem to be worth it, showcasing the power of automated fusion compilers like [NvFuser](https://github.com/NVIDIA/Fuser).
## Examples and benchmarks
> [!WARNING]
> Lightning Thunder is alpha and not ready for production runs. Feel free to try it out, expect a few bumps along the way.
> We expect speed and memory usage to improve as we continue to develop it.
We provide a version of the main pre-training script [that integrates Thunder](pretrain.py) that uses TinyLlama, a 1.1B parameter LLM.
| Setting | Compiler | Executors | Devices | ms/iter @ step 10 | Memory (GB) |
|----------------------|----------|----------------------------------------|---------|-------------------|---------------|
| Fully-sharded ZeRO 3 | Eager | - | 8 | 456.57 | 22.13 |
| Fully-sharded ZeRO 3 | torch | - | 8 | Not supported | Not supported |
| Fully-sharded ZeRO 3 | Thunder | sdpa, torchcompile | 8 | Not supported | Not supported |
| Fully-sharded ZeRO 3 | Thunder | sdpa, torchcompile_cat, nvfuser, torch | 8 | 333.56 | 21.40 |
| | | | | | |
| Replicated | Eager | - | 8 | 569.46 | 32.04 |
| Replicated | torch | - | 8 | Not supported | Not supported |
| Replicated | Thunder | sdpa, torchcompile | 8 | 426.44 | 22.19 |
| Replicated | Thunder | sdpa, torchcompile_cat, nvfuser, torch | 8 | 356.01 | 27.42 |
| | | | | | |
| - | Eager | - | 1 | 447.65 | 29.84 |
| - | torch | - | 1 | Not supported | Not supported |
| - | Thunder | sdpa, torchcompile | 1 | 373.37 | 22.19 |
| - | Thunder | sdpa, torchcompile_cat, nvfuser, torch | 1 | 322.25 | 27.42 |
| | | | | | |
| Unsloth | Thunder | sdpa, torchcompile_cat, nvfuser, torch | 1 | 331.92 | 25.19 |
<details>
<summary>Reproduction details</summary>
Config:
```yaml
out_dir: out/pretrain-thunder
data: TinyStories
tokenizer_dir: checkpoints/TinyLlama/TinyLlama-1.1B-Chat-v1.0
logger_name: csv
```
Commands:
```bash
litgpt download --repo_id TinyLlama/TinyLlama-1.1B-Chat-v1.0 --tokenizer_only true
python extensions/thunder/pretrain.py --config config.yaml --compiler null --train.global_batch_size 32
python extensions/thunder/pretrain.py --config config.yaml --executors '[sdpa, torchcompile]' --train.global_batch_size 32
python extensions/thunder/pretrain.py --config config.yaml --executors '[sdpa, torchcompile_cat, nvfuser, torch]' --train.global_batch_size 32
python extensions/thunder/pretrain.py --config config.yaml --compiler null --strategy ddp
python extensions/thunder/pretrain.py --config config.yaml --executors '[sdpa, torchcompile]' --strategy ddp
python extensions/thunder/pretrain.py --config config.yaml --executors '[sdpa, torchcompile_cat, nvfuser, torch]' --strategy ddp
python extensions/thunder/pretrain.py --config config.yaml --compiler null --devices 1
python extensions/thunder/pretrain.py --config config.yaml --executors '[sdpa, torchcompile]' --devices 1
python extensions/thunder/pretrain.py --config config.yaml --executors '[sdpa, torchcompile_cat, nvfuser, torch]' --devices 1
python extensions/thunder/pretrain.py --config config.yaml --executors '[sdpa, unsloth, torchcompile_cat, nvfuser, torch]' --devices 1
```
`--compiler torch` (`torch.compile` without `thunder`) is not include because it does not support compiling the `_FabricModule` due to this issue: https://github.com/pytorch/pytorch/issues/112787#issuecomment-1986827601
The CUDA devices are all NVIDIA A100-SXM4-40GB.
```text
Python version: 3.10.12 [GCC 11.4.0] (64-bit runtime)
Is debug build: False
CUDA used to build PyTorch: 12.1
CUDA runtime version: 12.3.107
Nvidia driver version: 545.23.08
pytorch-triton==3.0.0+45fff310c8
torch==2.4.0.dev20240427+cu121
lightning==2.3.0.dev20240328
lightning-thunder==0.2.0.dev20240505
nvfuser_cu121==0.2.3.dev20240428
```
</details>
extensions/thunder/__init__.py
+6
View File
@@ -0,0 +1,6 @@
import sys
from pathlib import Path
# support running without installing as a package, adding extensions to the Python path
wd = Path(__file__).parent.parent.resolve()
sys.path.append(str(wd))
extensions/thunder/pretrain.py
+520
View File
@@ -0,0 +1,520 @@
# Copyright Lightning AI. Licensed under the Apache License 2.0, see LICENSE file.
import math
import os
import pprint
import sys
import time
from collections.abc import Callable
from dataclasses import asdict
from datetime import timedelta
from functools import partial
from pathlib import Path
from typing import Any, Literal
import lightning as L
import torch
import torch.nn as nn
from lightning.fabric.strategies import FSDPStrategy
from lightning.fabric.utilities.throughput import ThroughputMonitor, measure_flops
from torch.utils.data import DataLoader
from torchmetrics.aggregation import RunningMean
from litgpt import Tokenizer
from litgpt.args import EvalArgs, LogArgs, TrainArgs
from litgpt.data import DataModule, TinyLlama
from litgpt.model import GPT, Block, CausalSelfAttention, Config, LLaMAMLP, MultiheadLatentAttention
from litgpt.parser_config import save_hyperparameters
from litgpt.types import LoggerChoice
from litgpt.utils import (
CLI,
CycleIterator,
capture_hparams,
choose_logger,
chunked_cross_entropy,
copy_config_files,
find_resume_path,
instantiate_torch_optimizer,
num_parameters,
parse_devices,
reset_parameters,
save_config,
)
# support running without installing as a package
wd = Path(__file__).parent.resolve()
sys.path.append(str(wd))
def forward_and_loss(model: nn.Module, input_ids: torch.Tensor, targets: torch.Tensor) -> torch.Tensor:
logits = model(input_ids)
# disable chunk_size to enable the unsloth cross entropy kernel
loss = chunked_cross_entropy(logits, targets, chunk_size=0)
return loss
def setup(
model_name: str | None = None,
model_config: Config | None = None,
out_dir: Path = Path("out/pretrain"),
initial_checkpoint_dir: Path | None = None,
resume: bool | Literal["auto"] | Path = False,
data: DataModule | None = None,
train: TrainArgs = TrainArgs(
save_interval=1000,
log_interval=1,
global_batch_size=512,
micro_batch_size=4,
max_tokens=int(3e12), # 3 trillion
max_norm=1.0,
min_lr=4e-5,
lr_warmup_steps=2000,
tie_embeddings=False,
),
eval: EvalArgs = EvalArgs(interval=1000, max_iters=100),
log: LogArgs = LogArgs(),
optimizer: str | dict = "AdamW",
devices: int | str = "auto",
num_nodes: int = 1,
tokenizer_dir: Path | None = None,
logger_name: LoggerChoice = "tensorboard",
seed: int = 42,
compiler: Literal["thunder", "torch"] | None = "thunder",
executors: list[str] | None = ("sdpa", "torchcompile", "nvfuser", "torch"),
strategy: Literal["auto", "ddp", "fsdp"] = "fsdp",
):
"""Pretrain a model.
Arguments:
model_name: The name of the model to pretrain. Choose from names in ``litgpt.config``. Mutually exclusive with
``model_config``.
model_config: A ``litgpt.Config`` object to define the model architecture. Mutually exclusive with
``model_config``.
out_dir: Directory in which to save checkpoints and logs. If running in a Lightning Studio Job, look for it in
/teamspace/jobs/<job-name>/share.
initial_checkpoint_dir: Optional path to a checkpoint directory to initialize the model from.
Useful for continued pretraining. Mutually exclusive with ``resume``.
resume: Path to a checkpoint directory to resume from in case training was interrupted, or ``True`` to resume
from the latest checkpoint in ``out_dir``. An error will be raised if no checkpoint is found. Passing
``'auto'`` will resume from the latest checkpoint but not error if no checkpoint exists.
data: Data-related arguments. If not provided, the default is ``litgpt.data.TinyLlama``.
train: Training-related arguments. See ``litgpt.args.TrainArgs`` for details.
eval: Evaluation-related arguments. See ``litgpt.args.EvalArgs`` for details.
optimizer: An optimizer name (such as "AdamW") or config.
devices: How many devices/GPUs to use. Uses all GPUs by default.
num_nodes: How many nodes the code is being run on.
tokenizer_dir: Optional path to the tokenizer dir that was used for preprocessing the dataset. Only some data
module require this.
logger_name: The name of the logger to send metrics to.
seed: The random seed to use for reproducibility.
compiler: If desired, the compiler/JIT to use.
executors: If using Thunder, the executors to enable.
strategy: If desired, the strategy to use.
"""
hparams = capture_hparams()
data = TinyLlama() if data is None else data
if model_config is not None and model_name is not None:
raise ValueError("Only one of `model_name` or `model_config` can be set.")
elif model_config is None and model_name is None:
model_name = "tiny-llama-1.1b"
config = Config.from_name(model_name) if model_config is None else model_config
devices = parse_devices(devices)
out_dir = init_out_dir(out_dir)
# in case the dataset requires the Tokenizer
tokenizer = Tokenizer(tokenizer_dir) if tokenizer_dir is not None else None
logger = choose_logger(
logger_name,
out_dir,
name=f"pretrain-{config.name}",
resume=bool(resume),
log_interval=train.log_interval,
log_args=asdict(log),
)
if devices * num_nodes > 1:
if compiler == "thunder":
if strategy == "fsdp":
from extensions.thunder.strategies import ThunderFSDPStrategy
strategy = ThunderFSDPStrategy(
sharding_strategy="ZERO3",
bucketing_strategy="BLOCK",
state_dict_type="full",
jit=False,
)
elif strategy == "ddp":
from extensions.thunder.strategies import ThunderDDPStrategy
strategy = ThunderDDPStrategy(jit=False)
else:
if strategy == "fsdp":
strategy = FSDPStrategy(
auto_wrap_policy={Block}, state_dict_type="full", sharding_strategy="FULL_SHARD"
)
else:
strategy = "auto"
fabric = L.Fabric(devices=devices, num_nodes=num_nodes, strategy=strategy, precision="bf16-true", loggers=[logger])
fabric.launch()
if compiler is not None:
global forward_and_loss
forward_and_loss = (
jit(forward_and_loss, executors) if compiler == "thunder" else torch.compile(forward_and_loss)
)
fabric.print(pprint.pformat(hparams))
if logger_name in ("tensorboard", "wandb", "mlflow"):
fabric.logger.log_hyperparams(hparams)
main(
fabric=fabric,
devices=devices,
num_nodes=num_nodes,
seed=seed,
initial_checkpoint_dir=initial_checkpoint_dir,
resume=resume,
config=config,
data=data,
out_dir=out_dir,
tokenizer_dir=tokenizer_dir,
tokenizer=tokenizer,
train=train,
eval=eval,
optimizer=optimizer,
compiler=compiler,
)
def main(
fabric: L.Fabric,
devices: int,
seed: int,
initial_checkpoint_dir: Path | None,
resume: bool | Literal["auto"] | Path,
config: Config,
data: DataModule,
out_dir: Path,
tokenizer_dir: Path | None,
tokenizer: Tokenizer | None,
train: TrainArgs,
eval: EvalArgs,
optimizer: str | dict,
compiler: Literal["thunder", "torch"] | None,
num_nodes: int = 1,
) -> None:
validate_args(train, eval, initial_checkpoint_dir, resume)
if fabric.global_rank == 0:
out_dir.mkdir(parents=True, exist_ok=True)
fabric.seed_everything(seed) # same seed for every process to init model (FSDP)
t0 = time.perf_counter()
with fabric.init_module(empty_init=True):
model = GPT(config)
initialize_weights(fabric, model, n_layer=config.n_layer, n_embd=config.n_embd)
if train.tie_embeddings:
model.transformer.wte.weight = model.lm_head.weight
if train.max_seq_length:
model.max_seq_length = train.max_seq_length
fabric.print(f"Time to instantiate model: {time.perf_counter() - t0:.02f} seconds.")
fabric.print(f"Total parameters: {num_parameters(model):,}")
model = fabric.setup(model)
if compiler == "thunder":
# avoid `Tensor.register_hook` which is unsupported
model._register_backward_hook = lambda *_: None
optimizer = instantiate_torch_optimizer(optimizer, model.parameters())
optimizer = fabric.setup_optimizers(optimizer)
train_dataloader, val_dataloader = get_dataloaders(fabric, data, tokenizer, train, model.max_seq_length)
train_dataloader, val_dataloader = fabric.setup_dataloaders(train_dataloader, val_dataloader)
if initial_checkpoint_dir:
fabric.load_raw(initial_checkpoint_dir / "lit_model.pth", model)
state = {
"model": model,
"optimizer": optimizer,
"train_dataloader": train_dataloader,
"iter_num": 0,
"step_count": 0,
}
resume = find_resume_path(resume, out_dir)
if resume:
fabric.print(f"Resuming training from {resume}")
fabric.load(resume, state)
train_time = time.perf_counter()
fit(
fabric=fabric,
devices=devices,
num_nodes=num_nodes,
state=state,
train_dataloader=train_dataloader,
val_dataloader=val_dataloader,
out_dir=out_dir,
tokenizer_dir=tokenizer_dir,
train=train,
eval=eval,
optimizer=optimizer,
)
fabric.print(f"Training time: {(time.perf_counter() - train_time):.2f}s")
# Save final checkpoint
save_checkpoint(fabric, state, tokenizer_dir, out_dir / "final" / "lit_model.pth")
if fabric.device.type == "cuda":
fabric.print(f"Memory used: {torch.cuda.max_memory_allocated() / 1e9:.02f} GB")
def fit(
fabric: L.Fabric,
devices: int,
state: dict,
train_dataloader: DataLoader,
val_dataloader: DataLoader,
out_dir: Path,
tokenizer_dir: Path | None,
train: TrainArgs,
eval: EvalArgs,
optimizer: str | dict,
num_nodes: int = 1,
) -> None:
model = state["model"]
optimizer = state["optimizer"]
validate(fabric, model, val_dataloader, max_iters=2) # sanity check
throughput = ThroughputMonitor(fabric, window_size=5)
with torch.device("meta"):
meta_model = GPT(model.config)
x = torch.randint(0, 1, (train.micro_batch_size, meta_model.max_seq_length))
model_fwd = lambda: meta_model(x) # noqa: F821
model_loss = lambda y: chunked_cross_entropy(y, x, chunk_size=0) # noqa: F821
measured_flops = measure_flops(meta_model, model_fwd, model_loss)
fabric.print(f"Measured TFLOPs: {measured_flops * fabric.world_size / 1e12:.2f}")
del meta_model, x
max_tokens_per_device = train.max_tokens // fabric.world_size
tokens_per_iter = train.micro_batch_size * model.max_seq_length
max_iters = max_tokens_per_device // tokens_per_iter
log_iter_interval = train.log_interval * train.gradient_accumulation_iters(devices, num_nodes)
initial_iter = state["iter_num"]
train_iterator = CycleIterator(train_dataloader)
running_loss = RunningMean(window=train.gradient_accumulation_iters(devices, num_nodes), sync_on_compute=False).to(
fabric.device
)
fabric.barrier()
total_t0 = time.perf_counter()
val_loss = "n/a"
warmup_iters = train.warmup_iters(devices, num_nodes, max_iters, train_dataloader)
for train_data in train_iterator:
if state["iter_num"] >= max_iters:
break
# determine and set the learning rate for this iteration
lr = get_lr(optimizer.defaults["lr"], state["iter_num"], warmup_iters, max_iters, train.min_lr)
for param_group in optimizer.param_groups:
param_group["lr"] = lr
state["iter_num"] += 1
iter_t0 = time.perf_counter()
input_ids = train_data[:, 0 : model.max_seq_length].contiguous().long()
targets = train_data[:, 1 : (model.max_seq_length + 1)].contiguous().long()
is_accumulating = state["iter_num"] % train.gradient_accumulation_iters(devices, num_nodes) != 0
with fabric.no_backward_sync(model, enabled=is_accumulating):
loss = forward_and_loss(model, input_ids, targets)
fabric.backward(loss / train.gradient_accumulation_iters(devices, num_nodes))
running_loss.update(loss.detach())
if not is_accumulating:
# THUNDER unsupported: https://github.com/Lightning-AI/lightning-thunder/issues/2357
# fabric.clip_gradients(model, optimizer, max_norm=train.max_norm)
optimizer.step()
optimizer.zero_grad()
state["step_count"] += 1
if state["iter_num"] % log_iter_interval == 0:
loss = running_loss.compute().item() # expensive device-to-host synchronization
t1 = time.perf_counter()
throughput.update(
time=(t1 - total_t0),
flops=(measured_flops * log_iter_interval),
batches=state["iter_num"],
samples=(state["iter_num"] * train.micro_batch_size),
lengths=(state["iter_num"] * train.micro_batch_size * model.max_seq_length),
)
metrics = {
"loss": loss,
"iter": state["iter_num"],
"step": state["step_count"],
"epoch": train_iterator.epoch,
"iter_time": t1 - iter_t0,
"remaining_time": (
(t1 - total_t0) / (state["iter_num"] - initial_iter) * (max_iters - state["iter_num"])
),
"tokens": state["iter_num"] * train.micro_batch_size * model.max_seq_length,
"total_tokens": (state["iter_num"] * train.micro_batch_size * model.max_seq_length * fabric.world_size),
"learning_rate": lr,
}
if isinstance(val_loss, float):
val_loss = f"{val_loss:.3f}"
fabric.print(
f"Epoch {metrics['epoch'] + 1} | iter {metrics['iter']} step {metrics['step']} |"
f" loss train: {metrics['loss']:.3f},"
f" val: {val_loss} |"
f" iter time: {metrics['iter_time'] * 1000:.2f} ms"
f"{' (step)' if not is_accumulating else ''}"
f" remaining time: {timedelta(seconds=int(metrics['remaining_time']))!s}"
)
throughput_metrics = throughput.compute()
metrics.update(throughput_metrics)
fabric.log_dict(metrics, step=state["iter_num"] - 1)
if val_dataloader is not None and not is_accumulating and state["step_count"] % eval.interval == 0:
t0 = time.perf_counter()
val_loss = validate(fabric, model, val_dataloader, max_iters=eval.max_iters)
val_loss = val_loss.item()
td = time.perf_counter() - t0
fabric.print(f"iter {state['iter_num']}: val loss {val_loss:.4f}, val time: {td * 1000:.2f} ms")
metrics = {"val_loss": val_loss, "val_ppl": math.exp(val_loss)}
fabric.log_dict(metrics, step=state["iter_num"] - 1)
fabric.barrier()
if train.save_interval is not None and not is_accumulating and state["step_count"] % train.save_interval == 0:
save_checkpoint(fabric, state, tokenizer_dir, out_dir / f"step-{state['step_count']:08d}" / "lit_model.pth")
@torch.no_grad()
def validate(fabric: L.Fabric, model: nn.Module, val_dataloader: DataLoader, max_iters: int) -> torch.Tensor:
fabric.barrier()
fabric.print("Validating ...")
model.eval()
losses = []
for k, batch in enumerate(val_dataloader):
if k >= max_iters:
break
input_ids = batch[:, 0 : model.max_seq_length].contiguous().long()
targets = batch[:, 1 : (model.max_seq_length + 1)].contiguous().long()
loss = forward_and_loss(model, input_ids, targets)
losses.append(loss)
val_loss = torch.stack(losses).mean()
model.train()
fabric.barrier()
return val_loss
def get_dataloaders(
fabric: L.Fabric, data: DataModule, tokenizer: Tokenizer, train: TrainArgs, block_size: int
) -> tuple[DataLoader, DataLoader]:
data.connect(tokenizer=tokenizer, batch_size=train.micro_batch_size, max_seq_length=block_size)
with fabric.rank_zero_first():
data.prepare_data()
data.setup()
train_dataloader = data.train_dataloader()
val_dataloader = data.val_dataloader()
return train_dataloader, val_dataloader
# learning rate decay scheduler (cosine with linear warmup)
def get_lr(learning_rate: float, it: int, warmup_iters: int, max_iters: int, min_lr: float) -> float:
# 1) linear warmup for warmup_iters steps
if it < warmup_iters:
return learning_rate * it / warmup_iters
# 2) if it > max_iters, return min learning rate
if it > max_iters:
return min_lr
# 3) in between, use cosine decay down to min learning rate
decay_ratio = (it - warmup_iters) / (max_iters - warmup_iters)
assert 0 <= decay_ratio <= 1
coeff = 0.5 * (1.0 + math.cos(math.pi * decay_ratio)) # coeff ranges 0..1
return min_lr + coeff * (learning_rate - min_lr)
def initialize_weights(fabric: L.Fabric, model: GPT, n_layer: int, n_embd: int) -> None:
"""GPT-NeoX weight initialization (https://arxiv.org/abs/2204.06745)."""
# Adapted from https://github.com/jzhang38/TinyLlama
def init_weights(module, std):
nn.init.normal_(module.weight, mean=0.0, std=std)
if getattr(module, "bias", None) is not None:
nn.init.zeros_(module.bias)
for mod in model.modules():
if isinstance(mod, (nn.Embedding, nn.Linear)):
mod.reset_parameters = partial(init_weights, mod, std=math.sqrt(2.0 / 5 / n_embd))
# need a separate loop because `mod.proj` below is a `nn.Linear` too
for mod in model.modules():
if isinstance(mod, (LLaMAMLP, CausalSelfAttention, MultiheadLatentAttention)):
mod.proj.reset_parameters = partial(init_weights, mod.proj, std=(1 / math.sqrt(n_embd) / n_layer))
if not isinstance(fabric.strategy, FSDPStrategy):
reset_parameters(model)
def init_out_dir(out_dir: Path) -> Path:
if not out_dir.is_absolute() and "LIGHTNING_ARTIFACTS_DIR" in os.environ:
return Path(os.getenv("LIGHTNING_ARTIFACTS_DIR")) / out_dir
return out_dir
def save_checkpoint(fabric, state, tokenizer_dir, checkpoint_file):
model = state["model"]
checkpoint_file.parent.mkdir(parents=True, exist_ok=True)
fabric.print(f"Saving checkpoint to {str(checkpoint_file)!r}")
fabric.save(checkpoint_file, state)
if fabric.global_rank == 0:
save_hyperparameters(setup, checkpoint_file.parent)
if tokenizer_dir is not None:
copy_config_files(tokenizer_dir, checkpoint_file.parent)
save_config(model.config, checkpoint_file.parent)
def validate_args(train: TrainArgs, eval: EvalArgs, initial_checkpoint_dir, resume) -> None:
issues = []
unsupported = [(train, ["max_steps", "epochs"]), (eval, ["max_new_tokens"])]
for args, names in unsupported:
for name in names:
if getattr(args, name) is not None:
issues.append(f"{__file__} doesn't support the {name!r} argument. This is set in {args}")
required = [(train, ["max_tokens", "max_norm"])]
for args, names in required:
for name in names:
if getattr(args, name) is None:
issues.append(f"{__file__} requires the {name!r} argument. This is set in {args}")
if initial_checkpoint_dir and resume:
issues.append("Can't provide both `--resume` and `--initial_checkpoint_dir`. Choose one.")
if issues:
raise ValueError("\n".join(issues))
def jit(fn: Callable, executors: list[str]) -> Any:
assert executors is not None
from unsloth.executor import unsloth_ex # import for registration # noqa: F401
import thunder
return thunder.jit(fn, executors=executors)
if __name__ == "__main__":
torch.set_float32_matmul_precision("high")
CLI(setup)
extensions/thunder/strategies/__init__.py
+2
View File
@@ -0,0 +1,2 @@
from .thunder_ddp import ThunderDDPStrategy # noqa: F401
from .thunder_fsdp import ThunderFSDPStrategy # noqa: F401
extensions/thunder/strategies/thunder_ddp.py
+256
View File
@@ -0,0 +1,256 @@
"""Fabric Strategy to support Thunder DDP: To be upstreamed into Fabric eventually."""
from contextlib import AbstractContextManager, nullcontext
from datetime import timedelta
from typing import TYPE_CHECKING, Any, Union
import torch
import torch.distributed
from lightning.fabric.accelerators.accelerator import Accelerator
from lightning.fabric.plugins.collectives.torch_collective import default_pg_timeout
from lightning.fabric.plugins.environments.cluster_environment import ClusterEnvironment
from lightning.fabric.plugins.io.checkpoint_io import CheckpointIO
from lightning.fabric.plugins.precision import Precision
from lightning.fabric.strategies.launchers.subprocess_script import _SubprocessScriptLauncher
from lightning.fabric.strategies.parallel import ParallelStrategy
from lightning.fabric.strategies.strategy import TBroadcast, _BackwardSyncControl
from lightning.fabric.utilities.distributed import (
ReduceOp,
_distributed_is_initialized,
_get_default_process_group_backend_for_device,
_init_dist_connection,
_sync_ddp_if_available,
)
from lightning.fabric.utilities.rank_zero import rank_zero_only
from lightning_utilities.core.rank_zero import rank_zero_only as utils_rank_zero_only
from torch import Tensor
from torch.nn import Module
from typing_extensions import override
from litgpt.constants import _THUNDER_AVAILABLE
if TYPE_CHECKING:
from thunder import Executor
class ThunderDDPStrategy(ParallelStrategy):
def __init__(
self,
accelerator: Accelerator | None = None,
parallel_devices: list[torch.device] | None = None,
cluster_environment: ClusterEnvironment | None = None,
checkpoint_io: CheckpointIO | None = None,
precision: Precision | None = None,
jit: bool = True,
executors: tuple[Union["Executor", str], ...] | None = None,
process_group_backend: str | None = None,
timeout: timedelta | None = default_pg_timeout,
**kwargs: Any,
):
r"""Strategy for Replicated Data Parallel provided by Lightning Thunder.
.. warning:: This is an :ref:`experimental <versioning:Experimental API>` feature.
Arguments:
jit: Whether to automatically call ``thunder.jit(model)`` if necessary. Disable this if you are manually
jitting a function that includes the model.
executors: The list of Thunder executors to enable. They can be either string aliases for the executors
or the actual executor instances.
\**kwargs: See available parameters in :func:`thunder.distributed.ddp`.
"""
if not _THUNDER_AVAILABLE:
raise ModuleNotFoundError(str(_THUNDER_AVAILABLE))
super().__init__(accelerator=accelerator, checkpoint_io=checkpoint_io, precision=precision)
self.parallel_devices = parallel_devices
self.cluster_environment: ClusterEnvironment | None = cluster_environment
if not jit and executors is not None:
raise ValueError(f"Passing executors={executors} doesn't have an effect with `jit={jit}`")
self.jit = jit
self.executors = executors
self._num_nodes = 1
self._process_group_backend: str | None = process_group_backend
self._timeout: timedelta | None = timeout
self._backward_sync_control = _ThunderDataParalellBackwardSyncControl()
self._ddp_kwargs = kwargs
@property
@override
def root_device(self) -> torch.device:
assert self.parallel_devices is not None
return self.parallel_devices[self.local_rank]
@property
def num_nodes(self) -> int:
return self._num_nodes
@num_nodes.setter
def num_nodes(self, num_nodes: int) -> None:
# note that world ranks is related to num_nodes, when resetting it, need to reset world ranks
self._num_nodes = num_nodes
@property
def num_processes(self) -> int:
return len(self.parallel_devices) if self.parallel_devices is not None else 0
@property
@override
def distributed_sampler_kwargs(self) -> dict[str, Any]:
return {"num_replicas": self.num_nodes * self.num_processes, "rank": self.global_rank}
@override
def _configure_launcher(self) -> None:
assert self.cluster_environment is not None
if not self.cluster_environment.creates_processes_externally:
self._launcher = _SubprocessScriptLauncher(self.cluster_environment, self.num_processes, self.num_nodes)
@property
def process_group_backend(self) -> str | None:
return self._process_group_backend
@override
def _configure_launcher(self) -> None:
assert self.cluster_environment is not None
self._launcher = _SubprocessScriptLauncher(self.cluster_environment, self.num_processes, self.num_nodes)
@override
def setup_environment(self) -> None:
super().setup_environment()
self._setup_distributed()
@override
def setup_module(self, module: Module) -> Module:
import thunder
if (cd := thunder.compile_data(module)) is not None:
# the module was already jitted
if thunder.compile_stats(module).last_traces is not None:
raise RuntimeError(
"You already called `thunder.jit()` and generated an execution trace. It's too late to apply the"
" DDP transform. Remove the `forward` call before `fabric.setup()`"
)
assert cd.is_module # sanity check
ddp_module = thunder.distributed.ddp(cd.fn, **self._ddp_kwargs)
# update the compile data state
cd.fn = ddp_module
cd.process_group_for_ddp = ddp_module.process_group_for_ddp
return module
else:
module = thunder.distributed.ddp(module, **self._ddp_kwargs)
if not self.jit:
return module
return thunder.jit(module, executors=self.executors)
@override
def module_to_device(self, module: Module) -> None:
module.to(self.root_device)
@override
def all_reduce(self, tensor: Tensor, group: Any | None = None, reduce_op: ReduceOp | str | None = "mean") -> Tensor:
if isinstance(tensor, Tensor):
return _sync_ddp_if_available(tensor, group, reduce_op=reduce_op)
return tensor
@override
def barrier(self, *args: Any, **kwargs: Any) -> None:
if not _distributed_is_initialized():
return
if torch.distributed.get_backend() == "nccl":
torch.distributed.barrier(device_ids=[self.root_device.index])
else:
torch.distributed.barrier()
@override
def broadcast(self, obj: TBroadcast, src: int = 0) -> TBroadcast:
if not _distributed_is_initialized():
return obj
obj = [obj]
torch.distributed.broadcast_object_list(obj, src)
return obj[0]
def _setup_distributed(self) -> None:
self._set_world_ranks()
self._process_group_backend = self._get_process_group_backend()
assert self.cluster_environment is not None
_init_dist_connection(self.cluster_environment, self._process_group_backend, timeout=self._timeout)
def _get_process_group_backend(self) -> str:
return self._process_group_backend or _get_default_process_group_backend_for_device(self.root_device)
def _set_world_ranks(self) -> None:
if self.cluster_environment is not None:
self.cluster_environment.set_global_rank(self.node_rank * self.num_processes + self.local_rank)
self.cluster_environment.set_world_size(self.num_nodes * self.num_processes)
# `LightningEnvironment.set_global_rank` will do this too, but we cannot rely on that implementation detail
# additionally, for some implementations, the setter is a no-op, so it's safer to access the getter
rank_zero_only.rank = utils_rank_zero_only.rank = self.global_rank
class _ThunderDataParalellBackwardSyncControl(_BackwardSyncControl):
def __init__(self):
self._enabled = False
@override
def no_backward_sync(self, module: Module, enabled: bool) -> AbstractContextManager:
"""
In Thunder, we cannot use ``module.no_sync()`` because reduction happens at the end of the context manager.
It assumes that the user will reuse it across all gradient accumulation iterations:
.. code-block:: python
with model.no_sync():
for _ in range(len(gradient_accumulation_iters)):
fwd()
bwd() # uses no-sync-backward trace
fwd()
bwd() # uses regular-backward trace
However, Fabric is designed to the context manager every iteration:
.. code-block:: python
for i in range(iters):
is_accumulating = (i + 1) % gradient_accumulation_iters != 0
ctx = model.no_sync() if is_accumulating else nullcontext()
with ctx:
fwd()
bwd()
So we need to be smart about when to sync grads based on the ``enabled`` value.
More info in https://github.com/Lightning-AI/lit-thunder-LEGACY/issues/2085
"""
if not getattr(module, "use_ddp", False) and not getattr(module, "use_fsdp", False):
raise TypeError(
"Blocking backward sync is only possible if the module passed to"
f" `{self.__class__.__name__}.no_backward_sync` is applied DDP or FSDP."
f" Got: {module.__class__.__name__}."
)
from thunder.distributed import skip_data_parallel_grad_sync
previous, self._enabled = self._enabled, enabled
if enabled:
return skip_data_parallel_grad_sync()
if not enabled and previous:
return _SyncGradsContextManager(module)
return nullcontext()
class _SyncGradsContextManager:
def __init__(self, module: Module) -> None:
self._module = module
@override
def __enter__(self) -> None:
from thunder.distributed import _sync_grads
_sync_grads(self._module)
@override
def __exit__(self, exc_type: Any, exc_value: Any, traceback: Any) -> None:
pass
extensions/thunder/strategies/thunder_fsdp.py
+458
View File
@@ -0,0 +1,458 @@
"""Fabric Strategy to support Thunder FSDP: To be upstreamed into Fabric eventually."""
import shutil
from collections.abc import Callable
from contextlib import AbstractContextManager, ExitStack, nullcontext
from pathlib import Path
from typing import TYPE_CHECKING, Any, Literal
import torch
from lightning.fabric.accelerators.accelerator import Accelerator
from lightning.fabric.plugins.environments.cluster_environment import ClusterEnvironment
from lightning.fabric.plugins.io.checkpoint_io import CheckpointIO
from lightning.fabric.plugins.precision import Precision
from lightning.fabric.strategies.launchers.subprocess_script import _SubprocessScriptLauncher
from lightning.fabric.strategies.parallel import ParallelStrategy
from lightning.fabric.strategies.strategy import TBroadcast, _apply_filter, _Sharded, _validate_keys_for_strict_loading
from lightning.fabric.utilities.distributed import (
ReduceOp,
_distributed_is_initialized,
_get_default_process_group_backend_for_device,
_init_dist_connection,
_sync_ddp_if_available,
)
from lightning.fabric.utilities.imports import _TORCH_GREATER_EQUAL_2_2
from lightning.fabric.utilities.load import _METADATA_FILENAME, _move_state_into
from lightning.fabric.utilities.rank_zero import rank_zero_only
from lightning.fabric.utilities.seed import reset_seed
from lightning.fabric.utilities.types import _PATH, _Stateful
from lightning_utilities.core.rank_zero import rank_zero_only as utils_rank_zero_only
from torch import Tensor
from torch.nn import Module
from torch.optim import Optimizer
from typing_extensions import override
from extensions.thunder.strategies.thunder_ddp import _ThunderDataParalellBackwardSyncControl
from litgpt.constants import _THUNDER_AVAILABLE
if TYPE_CHECKING:
from thunder import Executor
from thunder.distributed import FSDPBucketingStrategy, FSDPType
from thunder.distributed.checkpoint import StateDictOptions
_FSDP_TYPE = FSDPType | Literal["ZERO2", "ZERO3"]
_BUCKETING_STRATEGY = FSDPBucketingStrategy | Literal["NONE", "LAYER", "BLOCK"]
class ThunderFSDPStrategy(ParallelStrategy, _Sharded):
def __init__(
self,
accelerator: Accelerator | None = None,
parallel_devices: list[torch.device] | None = None,
cluster_environment: ClusterEnvironment | None = None,
checkpoint_io: CheckpointIO | None = None,
precision: Precision | None = None,
jit: bool = True,
executors: tuple["Executor | str", ...] | None = None,
sharding_strategy: "_FSDP_TYPE" = "ZERO3",
bucketing_strategy: "_BUCKETING_STRATEGY" = "NONE",
state_dict_type: Literal["full", "sharded"] = "sharded",
**kwargs: Any,
):
r"""Strategy for Fully Sharded Data Parallel provided by Lightning Thunder.
.. warning:: This is an :ref:`experimental <versioning:Experimental API>` feature.
Fully Sharded Training shards the entire model across all available GPUs, allowing you to scale model
size, whilst using efficient communication to reduce overhead. In practice, this means we can remain
at parity with PyTorch DDP, whilst scaling our model sizes dramatically.
Arguments:
jit: Whether to automatically call ``thunder.jit(model)`` if necessary. Disable this if you are manually
jitting a function that includes the model.
executors: The list of Thunder executors to enable. They can be either string aliases for the executors
or the actual executor instances.
sharding_strategy: Select whether to shard model parameters, gradients, optimizer states, or a combination
of them:
- ``"ZERO3"``: Shards model parameters, gradients, and optimizer states (default).
- ``"ZERO2"``: Shards gradients and optimizer states only. Model parameters get replicated.
Also accepts a :class:`thunder.distributed.FSDPType` enum value.
bucketing_strategy: Enables combining the collective operations for sets of layers.
- ``"NONE"``: No bucketing (default).
- ``"LAYER"``: Create buckets per layer class.
- ``"BLOCK"``: Create buckets per layer block.
Also accepts a :class:`thunder.distributed.FSDPBucketingStrategy` enum value.
state_dict_type: The format in which the state of the model and optimizers gets saved into the checkpoint.
- ``"full"``: The full weights and optimizer states get assembled on rank 0 and saved to a single file
(default).
- ``"sharded"``: Each rank saves its shard of weights and optimizer states to a file. The checkpoint is
a folder with as many files as the world size.
\**kwargs: See available parameters in :func:`thunder.distributed.fsdp`.
"""
if not _TORCH_GREATER_EQUAL_2_2:
raise ImportError("Thunder's FSDP strategy requires PyTorch 2.2 or higher.")
if not _THUNDER_AVAILABLE:
raise ModuleNotFoundError(str(_THUNDER_AVAILABLE))
super().__init__(accelerator=accelerator, checkpoint_io=checkpoint_io, precision=precision)
self.parallel_devices = parallel_devices
self.cluster_environment: ClusterEnvironment | None = cluster_environment
from thunder.distributed import FSDPBucketingStrategy, FSDPType
self.sharding_strategy = (
FSDPType[sharding_strategy.upper()] if isinstance(sharding_strategy, str) else sharding_strategy
)
self.bucketing_strategy = (
FSDPBucketingStrategy[bucketing_strategy.upper()]
if isinstance(bucketing_strategy, str)
else bucketing_strategy
)
if not jit and executors is not None:
raise ValueError(f"Passing executors={executors} doesn't have an effect with `jit={jit}`")
self.jit = jit
self.executors = executors
self._state_dict_type = state_dict_type
self._backward_sync_control = _ThunderDataParalellBackwardSyncControl()
self._fsdp_kwargs = kwargs
@property
@override
def root_device(self) -> torch.device:
assert self.parallel_devices is not None
return self.parallel_devices[self.local_rank]
@property
def num_nodes(self) -> int:
return 1
@property
def num_processes(self) -> int:
return len(self.parallel_devices) if self.parallel_devices is not None else 0
@property
@override
def distributed_sampler_kwargs(self) -> dict[str, Any]:
return {"num_replicas": self.num_nodes * self.num_processes, "rank": self.global_rank}
@override
def _configure_launcher(self) -> None:
assert self.cluster_environment is not None
if not self.cluster_environment.creates_processes_externally:
self._launcher = _SubprocessScriptLauncher(self.cluster_environment, self.num_processes, self.num_nodes)
@override
def setup_environment(self) -> None:
super().setup_environment()
self._setup_distributed()
@override
def setup_module(self, module: Module) -> Module:
import thunder
if (cd := thunder.compile_data(module)) is not None:
# the module was already jitted
if thunder.compile_stats(module).last_traces is not None:
raise RuntimeError(
"You already called `thunder.jit()` and generated an execution trace. It's too late to apply the"
" FSDP transform. Remove the `forward` call before `fabric.setup()`"
)
assert cd.is_module # sanity check
fsdp_module = thunder.distributed.fsdp(
cd.fn,
device=self.root_device,
sharding_strategy=self.sharding_strategy,
bucketing_strategy=self.bucketing_strategy,
**self._fsdp_kwargs,
)
# update the compile data state
cd.fn = fsdp_module
cd.process_group_for_ddp = fsdp_module.process_group_for_ddp
return module
else:
module = thunder.distributed.fsdp(
module,
device=self.root_device,
sharding_strategy=self.sharding_strategy,
bucketing_strategy=self.bucketing_strategy,
**self._fsdp_kwargs,
)
if not self.jit:
return module
return thunder.jit(module, executors=self.executors)
@override
def module_to_device(self, module: Module) -> None:
pass
@override
def module_init_context(self, empty_init: bool | None = None) -> AbstractContextManager:
precision_init_ctx = self.precision.module_init_context()
module_sharded_ctx = self.module_sharded_context()
stack = ExitStack()
if empty_init:
# Materialization happens in `setup`. When modules get wrapped by FSDP
stack.enter_context(torch.device("meta"))
stack.enter_context(precision_init_ctx)
stack.enter_context(module_sharded_ctx)
return stack
@override
def module_sharded_context(self) -> AbstractContextManager:
return nullcontext()
@override
def all_reduce(self, tensor: Tensor, group: Any | None = None, reduce_op: ReduceOp | str | None = "mean") -> Tensor:
if isinstance(tensor, Tensor):
return _sync_ddp_if_available(tensor, group, reduce_op=reduce_op)
return tensor
@override
def barrier(self, *args: Any, **kwargs: Any) -> None:
if not _distributed_is_initialized():
return
if torch.distributed.get_backend() == "nccl":
torch.distributed.barrier(device_ids=[self.root_device.index])
else:
torch.distributed.barrier()
@override
def broadcast(self, obj: TBroadcast, src: int = 0) -> TBroadcast:
if not _distributed_is_initialized():
return obj
obj = [obj]
torch.distributed.broadcast_object_list(obj, src)
return obj[0]
@override
def clip_gradients_norm(
self,
module: Module,
optimizer: Optimizer,
max_norm: float | int,
norm_type: float | int = 2.0,
error_if_nonfinite: bool = True,
) -> Tensor:
raise NotImplementedError
@override
def save_checkpoint(
self,
path: _PATH,
state: dict[str, Module | Optimizer | Any],
storage_options: Any | None = None,
filter: dict[str, Callable[[str, Any], bool]] | None = None,
) -> None:
if storage_options is not None:
raise TypeError(
"`FSDPStrategy.save_checkpoint(..., storage_options=...)` is not supported because"
" `FSDPStrategy` does not use the `CheckpointIO`."
)
if filter is not None:
raise NotImplementedError("Filtering checkpoint paths is not implemented")
# broadcast the path from rank 0 to ensure all the states are saved in a common path
path = Path(self.broadcast(path))
if path.is_dir() and self._state_dict_type == "full" and not _is_sharded_checkpoint(path):
raise IsADirectoryError(f"The checkpoint path exists and is a directory: {path}")
from thunder.distributed.checkpoint import StateDictOptions, has_fsdp_modules, save
modules = [module for module in state.values() if has_fsdp_modules(module)]
if len(modules) == 0:
raise ValueError(
"Could not find a FSDP model in the provided checkpoint state. Please provide the model as"
" part of the state like so: `save_checkpoint(..., state={'model': model, ...})`. Make sure"
" you set up the model (and optimizers if any) through the strategy before saving the checkpoint."
)
if len(modules) > 1:
raise ValueError(
"Found multiple FSDP models in the given state. Saving checkpoints with FSDP is"
" currently limited to a single model per checkpoint. To save multiple models, call the"
" save method for each model separately with a different path."
)
if self._state_dict_type == "sharded":
if _is_full_checkpoint(path):
path.unlink()
path.mkdir(parents=True, exist_ok=True)
options = StateDictOptions(full_state_dict=False, cpu_offload=True, rank0_only=False)
converted_state, metadata = _get_state_dict(state, filter, options, self.local_rank)
save(converted_state, path)
if self.global_rank == 0:
torch.save(metadata, path / _METADATA_FILENAME)
elif self._state_dict_type == "full":
if _is_sharded_checkpoint(path):
shutil.rmtree(path)
options = StateDictOptions(full_state_dict=True, cpu_offload=True, rank0_only=True)
converted_state, metadata = _get_state_dict(state, filter, options, self.local_rank)
converted_state.update(metadata)
if self.global_rank == 0:
torch.save(converted_state, path)
else:
raise ValueError(f"Unknown state_dict_type: {self._state_dict_type}")
@override
def load_checkpoint(
self,
path: _PATH,
state: Module | Optimizer | dict[str, Module | Optimizer | Any] | None = None,
strict: bool = True,
) -> dict[str, Any]:
if not state:
raise ValueError(
f"Got `FSDPStrategy.load_checkpoint(..., state={state!r})` but a state with at least"
" a model instance to reload is required. Pass it in like so:"
" `FSDPStrategy.load_checkpoint(..., state={'model': model, ...})`"
)
# broadcast the path from rank 0 to ensure all the states are loaded from a common path
path = Path(self.broadcast(path))
from thunder.distributed.checkpoint import StateDictOptions, has_fsdp_modules, load, load_model_state_dict
if isinstance(state, Module):
if not _is_full_checkpoint(path):
raise ValueError(
"Failed to load checkpoint directly into the model. The given path must be a single file"
f" containing the full state dict: {path}"
)
state_dict = torch.load(str(path), mmap=True, map_location="cpu")
options = StateDictOptions(full_state_dict=True, cpu_offload=True, strict=strict, rank0_only=False)
load_model_state_dict(state_dict, _unwrap_tom(state), options, self.local_rank)
return {}
if isinstance(state, Optimizer):
raise NotImplementedError(
"Loading a single optimizer object from a checkpoint is not supported yet with the FSDP strategy."
)
modules = {key: module for key, module in state.items() if has_fsdp_modules(module)}
if len(modules) == 0:
raise ValueError(
"Could not find a FSDP model in the provided checkpoint state. Please provide the model as"
" part of the state like so: `load_checkpoint(..., state={'model': model, ...})`. Make sure"
" you set up the model (and optimizers if any) through the strategy before loading the checkpoint."
)
if len(modules) > 1:
raise ValueError(
"Found multiple FSDP models in the given state. Loading checkpoints with FSDP is"
" currently limited to a single model per checkpoint. To load multiple models, call the"
" load method for each model separately with a different path."
)
optimizers = {key: optim for key, optim in state.items() if isinstance(optim, Optimizer)}
module_key, module = list(modules.items())[0]
module = _unwrap_tom(module)
if _is_sharded_checkpoint(path):
options = StateDictOptions(full_state_dict=False, cpu_offload=True, strict=strict, rank0_only=False)
# Load the DCP state dict, which requires a holder state dict
converted_state, _ = _get_state_dict(state, None, options, self.local_rank)
load(converted_state, path)
load_model_state_dict(converted_state[module_key], module, options, self.local_rank)
# Load metadata (anything not a module or optimizer)
metadata = torch.load(path / _METADATA_FILENAME)
requested_metadata_keys = state.keys() - modules.keys() - optimizers.keys()
_validate_keys_for_strict_loading(requested_metadata_keys, metadata.keys(), strict=strict)
for key in requested_metadata_keys:
if key not in metadata:
continue
state[key] = metadata.pop(key)
# return the remaining metadata that wasn't requested as part of `state`
return metadata
if _is_full_checkpoint(path):
options = StateDictOptions(full_state_dict=True, cpu_offload=True, strict=strict, rank0_only=False)
if not options.rank0_only or self.local_rank == 0:
map_location = "cpu" if options.cpu_offload else None
checkpoint = torch.load(str(path), mmap=True, map_location=map_location)
load_model_state_dict(checkpoint[module_key], module, options, self.local_rank)
else:
checkpoint = {}
requested_metadata_keys = state.keys() - modules.keys() - optimizers.keys()
_validate_keys_for_strict_loading(requested_metadata_keys, checkpoint.keys(), strict=strict)
# Load metadata (anything not a module or optimizer)
_move_state_into(source=checkpoint, destination=state, keys=requested_metadata_keys)
# return the remaining metadata that wasn't requested as part of `state`
return checkpoint
raise ValueError(
f"The path {str(path)!r} does not point to a valid checkpoint. Make sure the path points to either a"
" directory with FSDP checkpoint shards, or a single file with a full checkpoint."
)
def _setup_distributed(self) -> None:
reset_seed()
self._set_world_ranks()
process_group_backend = _get_default_process_group_backend_for_device(self.root_device)
assert self.cluster_environment is not None
_init_dist_connection(self.cluster_environment, process_group_backend)
def _set_world_ranks(self) -> None:
if self.cluster_environment is not None:
self.cluster_environment.set_global_rank(self.node_rank * self.num_processes + self.local_rank)
self.cluster_environment.set_world_size(self.num_nodes * self.num_processes)
# `LightningEnvironment.set_global_rank` will do this too, but we cannot rely on that implementation detail
# additionally, for some implementations, the setter is a no-op, so it's safer to access the getter
rank_zero_only.rank = utils_rank_zero_only.rank = self.global_rank
def _is_sharded_checkpoint(path: Path) -> bool:
"""A heuristic check to determine whether the path points to a directory with checkpoint shards."""
return path.is_dir() and (path / _METADATA_FILENAME).is_file()
def _is_full_checkpoint(path: Path) -> bool:
return path.is_file()
def _get_state_dict(
state: dict[str, Any],
filter: dict[str, Callable[[str, Any], bool]] | None,
options: "StateDictOptions",
rank: int,
) -> tuple[dict[str, Any], dict[str, Any]]:
from thunder.distributed.checkpoint import get_model_state_dict
# replace the modules and optimizer objects in the state with their local state dict
# and separate the user's metadata
converted_state: dict[str, Any] = {}
metadata: dict[str, Any] = {}
for key, obj in state.items():
converted: Any
if isinstance(obj, Module):
converted = get_model_state_dict(_unwrap_tom(obj), options, rank)
target_dict = converted_state
elif isinstance(obj, Optimizer):
# TODO: optimizer support
converted = obj.state_dict()
target_dict = converted_state
else: # everything not a module or optimizer is considered metadata
converted = obj.state_dict() if isinstance(obj, _Stateful) else obj
target_dict = metadata
_apply_filter(key, filter or {}, converted, target_dict)
return converted_state, metadata
def _unwrap_tom(obj: object) -> object:
# TODO: this unwrap won't be required when Fabric's `_unwrap_objects` supports Thunder
from thunder import ThunderModule
if isinstance(obj, ThunderModule):
return obj._model
return obj
extensions/thunder/unsloth/__init__.py
View File
extensions/thunder/unsloth/executor.py
+283
View File
@@ -0,0 +1,283 @@
# Copyright Lightning AI. Licensed under the Apache License 2.0, see LICENSE file.
import sys
from pathlib import Path
import torch
from torch import Tensor
import litgpt.model
from litgpt.constants import _THUNDER_AVAILABLE
from litgpt.model import LLaMAMLP as OriginalLLaMAMLP
from thunder.core.proxies import TensorProxy
from thunder.core.transforms import get_grad, mean_backward, put_grads
from thunder.extend import OperatorExecutor, register_executor
from thunder.torch import ne, sum, true_divide
if _THUNDER_AVAILABLE:
import thunder
import thunder.torch as ltorch
sys.path.append(str(Path(__file__).parent))
import kernels
unsloth_ex = OperatorExecutor("unsloth", version="0.1")
register_executor(unsloth_ex)
"""
====================
Cross Entropy Loss
====================
"""
def unsloth_cross_entropy_meta(logits: TensorProxy, labels: TensorProxy) -> tuple[TensorProxy, TensorProxy]:
return (
TensorProxy(
shape=(logits.shape[0],),
# the cross entropy kernel only supports float32
dtype=thunder.dtypes.float32,
device=logits.device,
requires_grad=logits.requires_grad,
),
TensorProxy(shape=(logits.shape[0],), dtype=thunder.dtypes.float32, device=logits.device, requires_grad=False),
)
unsloth_cross_entropy = unsloth_ex.register_operator(
"unsloth_cross_entropy", meta=unsloth_cross_entropy_meta, fn=kernels.cross_entropy_loss._cross_entropy_forward_impl
)
def unsloth_cross_entropy_backward_impl(dlosses: Tensor, logits: Tensor, labels: Tensor, logsumexp: Tensor) -> Tensor:
# clone() because the kernel writes the grads in the logits
return kernels.cross_entropy_loss._cross_entropy_backward_impl(dlosses, logits.clone(), logsumexp, labels)
def unsloth_cross_entropy_backward_meta(
dlosses: TensorProxy, logits: TensorProxy, logsumexp: TensorProxy, labels: TensorProxy
) -> TensorProxy:
return thunder.TensorProxy(like=logits)
unsloth_cross_entropy_backward = unsloth_ex.register_operator(
"unsloth_cross_entropy_backward", meta=unsloth_cross_entropy_backward_meta, fn=unsloth_cross_entropy_backward_impl
)
def unsloth_cross_entropy_checker(
logits: TensorProxy,
labels: TensorProxy,
weight: TensorProxy | None = None,
size_average: bool | None = None,
ignore_index: int = -100,
reduce: bool | None = None,
reduction: str = "mean",
label_smoothing: float = 0.0,
) -> bool:
return (
weight is None
and size_average is None
and reduce is None
and reduction in ("none", "mean")
and ignore_index == -100
and label_smoothing == 0.0
and logits.device.type == "cuda"
and labels.device.type == "cuda"
)
def cross_entropy_to_unsloth(
logits: TensorProxy,
labels: TensorProxy,
weight: TensorProxy | None = None,
size_average: bool | None = None,
ignore_index: int = -100,
reduce: bool | None = None,
reduction: str = "mean",
label_smoothing: float = 0.0,
) -> tuple[TensorProxy, TensorProxy]:
loss, logsumexp = unsloth_cross_entropy(logits, labels)
if reduction == "mean":
# "mean" reduction is not part of the kernel
# TODO: this doesn't consider that all elements could be masked, causing a division by 0
n_items = sum(ne(labels, -100))
loss = true_divide(sum(loss), n_items)
elif reduction != "none":
raise NotImplementedError(reduction)
return loss, logsumexp
def unsloth_cross_entropy_grad(
logits: TensorProxy,
labels: TensorProxy,
weight: TensorProxy | None = None,
size_average: bool | None = None,
ignore_index: int = -100,
reduce: bool | None = None,
reduction: str = "mean",
label_smoothing: float = 0.0,
) -> TensorProxy:
loss, logsumexp = cross_entropy_to_unsloth(**locals())
grad = get_grad(loss)
if reduction == "mean":
grad = mean_backward(logsumexp.ndim, logsumexp.shape, (0,), grad)
logits_grad = unsloth_cross_entropy_backward(grad, logits, labels, logsumexp)
put_grads((logits,), (logits_grad,))
return loss
# registers as cross entropy implementation, including the execution transform and now a grad transform
unsloth_ex.register_implementation(
ltorch.cross_entropy,
checker=unsloth_cross_entropy_checker,
execution_transform=lambda *args: cross_entropy_to_unsloth(*args)[0],
grad_transform=unsloth_cross_entropy_grad,
)
"""
=========
RMSNorm
=========
The RMSNorm kernel is not integrated because it's not numerically equal and it doesn't compute the gradient for the
weight, just for the input.
"""
"""
========
SwiGLU
========
"""
def swiglu(e: torch.Tensor, g: torch.Tensor) -> torch.Tensor:
return torch.nn.functional.silu(e) * g
class ThunderLLaMAMLP(OriginalLLaMAMLP):
def forward(self, x: torch.Tensor) -> torch.Tensor:
x_fc_1 = self.fc_1(x)
x_fc_2 = self.fc_2(x)
x = swiglu(x_fc_1, x_fc_2)
return self.proj(x)
litgpt.model.LLaMAMLP = ThunderLLaMAMLP
def swiglu_forward_meta(e: TensorProxy, g: TensorProxy) -> TensorProxy:
return TensorProxy(like=e)
litgpt_swiglu = unsloth_ex.register_operator("litgpt_swiglu", meta=swiglu_forward_meta, fn=swiglu, replaces=swiglu)
unsloth_swiglu_forward = unsloth_ex.register_operator(
"unsloth_swiglu_forward", meta=swiglu_forward_meta, fn=lambda *args: kernels.swiglu_fg_kernel(*args)
)
def unsloth_swiglu_backward_meta(DW: TensorProxy, e: TensorProxy, g: TensorProxy) -> tuple[TensorProxy, TensorProxy]:
return TensorProxy(like=g), TensorProxy(like=e)
def unsloth_swiglu_backward_fn(DW: Tensor, e: Tensor, g: Tensor) -> tuple[Tensor, tuple]:
B, T, n_embd = e.shape
e = e.view(-1, n_embd)
g = g.view(-1, n_embd)
DW, e, g = kernels.swiglu_DWf_DW_dfg_kernel(DW, e, g)
e = e.view(B, T, n_embd)
g = g.view(B, T, n_embd)
return g, e
unsloth_swiglu_backward = unsloth_ex.register_operator(
"unsloth_swiglu_backward", meta=unsloth_swiglu_backward_meta, fn=unsloth_swiglu_backward_fn
)
def swiglu_to_unsloth_checker(e: TensorProxy, g: TensorProxy) -> bool:
return e.device.type == "cuda" and g.device.type == "cuda"
def unsloth_swiglu_grad(e: TensorProxy, g: TensorProxy) -> TensorProxy:
h = unsloth_swiglu_forward(**locals())
grad = get_grad(h)
e_grad, g_grad = unsloth_swiglu_backward(grad, e, g)
put_grads((e, g), (e_grad, g_grad))
return h
unsloth_ex.register_implementation(
litgpt_swiglu,
checker=swiglu_to_unsloth_checker,
execution_transform=unsloth_swiglu_forward,
grad_transform=unsloth_swiglu_grad,
)
"""
======
RoPE
======
"""
def apply_rope_meta(x: TensorProxy, cos: TensorProxy, sin: TensorProxy) -> TensorProxy:
return TensorProxy(like=x)
apply_rope = unsloth_ex.register_operator(
"litgpt_apply_rope", like=apply_rope_meta, fn=litgpt.model.apply_rope, replaces=litgpt.model.apply_rope
)
def unsloth_apply_rope_meta(
Q: TensorProxy, cos: TensorProxy, sin: TensorProxy
) -> tuple[TensorProxy, TensorProxy, TensorProxy, int, int, int]:
batch, n_heads, seq_len, head_dim = Q.shape
assert seq_len <= cos.shape[-2]
BLOCK_SIZE, num_warps = kernels.calculate_settings(head_dim // 2)
div, mod = divmod(n_heads, kernels.rope_embedding.ROPE_GROUP_SIZE)
n_groups = div + (mod != 0)
return TensorProxy(like=Q), cos, sin, n_groups, BLOCK_SIZE, num_warps
unsloth_apply_rope = unsloth_ex.register_operator(
"unsloth_apply_rope", meta=unsloth_apply_rope_meta, fn=kernels._rope_embedding_forward_impl
)
def unsloth_apply_rope_backward_meta(
dY: TensorProxy, cos: TensorProxy, sin: TensorProxy, n_groups: int, BLOCK_SIZE: int, num_warps: int
) -> TensorProxy:
return TensorProxy(like=dY)
unsloth_apply_rope_backward = unsloth_ex.register_operator(
"unsloth_apply_rope_backward", meta=unsloth_apply_rope_backward_meta, fn=kernels._rope_embedding_backward_impl
)
def apply_rope_to_unsloth_checker(x: TensorProxy, cos: TensorProxy, sin: TensorProxy) -> bool:
return len(x.shape) == 4 and x.device.type == "cuda" and cos.device.type == "cuda" and sin.device.type == "cuda"
def unsloth_apply_rope_grad(x: TensorProxy, cos: TensorProxy, sin: TensorProxy) -> TensorProxy:
Q, cos, sin, n_groups, BLOCK_SIZE, num_warps = unsloth_apply_rope(x, cos, sin)
dY = get_grad(Q)
dX = unsloth_apply_rope_backward(dY, cos, sin, n_groups, BLOCK_SIZE, num_warps)
put_grads((x,), (dX,))
return Q
unsloth_ex.register_implementation(
apply_rope,
checker=apply_rope_to_unsloth_checker,
execution_transform=lambda *args: unsloth_apply_rope(*args)[0],
grad_transform=unsloth_apply_rope_grad,
)
extensions/thunder/unsloth/kernels/__init__.py
+4
View File
@@ -0,0 +1,4 @@
from .cross_entropy_loss import _cross_entropy_backward_impl, _cross_entropy_forward_impl # noqa: F401
from .rope_embedding import ROPE_GROUP_SIZE, _rope_embedding_backward_impl, _rope_embedding_forward_impl # noqa: F401
from .swiglu import swiglu_DWf_DW_dfg_kernel, swiglu_fg_kernel # noqa: F401
from .utils import calculate_settings # noqa: F401
extensions/thunder/unsloth/kernels/cross_entropy_loss.py
+285
View File
@@ -0,0 +1,285 @@
# Copyright 2023-present Daniel Han-Chen & the Unsloth team. All rights reserved.
#
# Licensed under the Apache License, Version 2.0 (the "License");
# you may not use this file except in compliance with the License.
# You may obtain a copy of the License at
#
# http://www.apache.org/licenses/LICENSE-2.0
#
# Unless required by applicable law or agreed to in writing, software
# distributed under the License is distributed on an "AS IS" BASIS,
# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
# See the License for the specific language governing permissions and
# limitations under the License.
import torch
from litgpt.constants import _TRITON_AVAILABLE
from .utils import MAX_FUSED_SIZE, calculate_settings
if _TRITON_AVAILABLE:
import triton
import triton.language as tl
@triton.jit
def _cross_entropy_forward(
logits_ptr,
logits_row_stride,
loss_ptr,
logsumexp_ptr,
labels_ptr,
VOCAB_SIZE: tl.constexpr,
BLOCK_SIZE: tl.constexpr,
):
"""
Cross Entropy Loss = 1/n sum [ -yi log(Pi) ]
Pi = exp(xi) / sum(exp(xi))
CE_i = -y log(p) = -y log[ exp(x) / sum(exp(x)) ]
= -y [ x - log[sum(exp(x))] ]
= y * (log[sum(exp(x))] - x)
If y == 0: CE_i = 0
If y == 1: CE_i = logsumexp - x
logsumexp is also stable
Take y = log[sum(exp(x))]
exp(y) = sum(exp(x))
exp(y) = sum(exp(x - c)*exp(c)) Since e^(x-c)*e^c = e^x
exp(y) = exp(c)*sum(exp(x - c))
y = log(exp(c)*sum(exp(x - c)))
y = c + log[sum(exp(x - c))]
This means we can set c = max(x) to make sure
exp(x - c) always is exp(x - max(x)).
This ensures exp(x - max(x))'s maximum is 1 as exp(0) = 1.
"""
row_idx = tl.program_id(0)
logits_ptr += row_idx * logits_row_stride.to(tl.int64)
loss_ptr += row_idx
logsumexp_ptr += row_idx
labels_ptr += row_idx
col_offsets = tl.arange(0, BLOCK_SIZE)
mask = col_offsets < VOCAB_SIZE
label_idx = tl.load(labels_ptr).to(tl.int32)
logits = tl.load(logits_ptr + col_offsets, mask=mask, other=-float("inf")).to(tl.float32)
c = tl.max(logits, 0)
logsumexp = c + tl.log(tl.sum(tl.exp(logits - c), 0))
if label_idx != -100:
x = tl.load(logits_ptr + label_idx).to(tl.float32)
loss = logsumexp - x
else:
loss = 0.0
tl.store(logsumexp_ptr, logsumexp)
tl.store(loss_ptr, loss)
pass
@triton.jit
def _chunked_cross_entropy_forward(
logits_ptr,
logits_row_stride,
loss_ptr,
logsumexp_ptr,
labels_ptr,
VOCAB_SIZE: tl.constexpr,
N_CHUNKS: tl.constexpr,
BLOCK_SIZE: tl.constexpr,
):
"""
256K vocab divided in 4 chunks
|-65536-| |-65536-| |-65536-| |-65536-|
|-------| |-------| |-------| |-------|
|-------| |-------| |-------| |-------|
If y == 0: CE_i = 0
If y == 1: CE_i = logsumexp - x
Notice we can do logsumexp for each chunk and then
logsumexp[chunk_sum(logsumexp)] == logsumexp
chunk_sum = log[chunk_sum(logsumexp)]
= log[exp(logsumexp(a)) + ... + exp(logsumexp(z))]
= log[exp(log[sum(exp(a))]) + ... + exp(log[sum(exp(z))])]
= log[sum(exp(a)) + ... + sum(exp(z))]
= logsumexp(x)
This means we can perform a logsumexp for each chunk, then do a
final logsumexp reduction!
Ie do: logsumexp(chunked_logsumexp) - x
"""
row_idx = tl.program_id(0)
chunk_idx = tl.program_id(1)
logits_ptr += row_idx * logits_row_stride.to(tl.int64)
loss_ptr += row_idx
logsumexp_ptr += row_idx * N_CHUNKS + chunk_idx
labels_ptr += row_idx
col_offsets = chunk_idx * BLOCK_SIZE + tl.arange(0, BLOCK_SIZE)
mask = col_offsets < VOCAB_SIZE
label_idx = tl.load(labels_ptr).to(tl.int32)
logits = tl.load(logits_ptr + col_offsets, mask=mask, other=-float("inf")).to(tl.float32)
c = tl.max(logits, 0)
logsumexp = c + tl.log(tl.sum(tl.exp(logits - c), 0))
if chunk_idx == 0:
# logsumexp(chunked_logsumexp) - x
# Do the -x separately
if label_idx != -100:
x = tl.load(logits_ptr + label_idx).to(tl.float32)
loss = -1.0 * x
else:
loss = 0.0
tl.store(loss_ptr, loss)
pass
tl.store(logsumexp_ptr, logsumexp)
pass
@triton.jit
def _cross_entropy_backward(
logits_ptr,
logits_row_stride,
dloss_ptr,
dloss_row_stride,
logsumexp_ptr,
labels_ptr,
VOCAB_SIZE: tl.constexpr,
BLOCK_SIZE: tl.constexpr,
):
"""
CE_i = -y log(P) = y * (log[sum(exp(x))] - x)
dC/dx = d/dx (y * log[sum(exp(x))] - x * y)
From https://en.wikipedia.org/wiki/LogSumExp
d/dx logsumexp = exp(x) / sum(exp(x)) = softmax(x)
dC/dx = y * exp(x) / sum(exp(x)) - d/dx (x * y)
dC/dx = y * exp[ log[exp(x) / sum(exp(x))] ] using x = exp(log(x)) trick
dC/dx = y * exp[x - logsumexp] - d/dx (x * y)
If y == 0: dC/dx = 0
If y == 1 and x == label: dC/dlabel = exp[x - logsumexp] - 1
If y == 1 and x != label: dC/dx = exp[x - logsumexp]
"""
row_idx = tl.program_id(0)
block_idx = tl.program_id(1)
logits_ptr += row_idx * logits_row_stride.to(tl.int64)
dloss_ptr += row_idx * dloss_row_stride
col_offsets = block_idx * BLOCK_SIZE + tl.arange(0, BLOCK_SIZE)
mask = col_offsets < VOCAB_SIZE
label_idx = tl.load(labels_ptr + row_idx).to(tl.int32)
if label_idx != -100:
dloss = tl.load(dloss_ptr)
else:
dloss = 0.0
x = tl.load(logits_ptr + col_offsets, mask=mask, other=-float("inf")).to(tl.float32)
logsumexp = tl.load(logsumexp_ptr + row_idx)
y = tl.exp(x - logsumexp)
y = tl.where(
col_offsets == label_idx,
y - 1.0, # exp(x - logsumexp) - 1
y, # exp(x - logsumexp)
)
# If y == 0: dC/dx = 0 ==> we already masked it to be = 0, so dloss = 0.
tl.store(logits_ptr + col_offsets, dloss * y, mask=mask)
pass
def _cross_entropy_forward_impl(logits, labels):
n_rows, vocab_size = logits.shape
div, mod = divmod(vocab_size, MAX_FUSED_SIZE)
n_chunks = div + (mod != 0)
losses = torch.empty(n_rows, dtype=torch.float32, device="cuda")
if n_chunks == 1:
# For small vocabs <= 65336 like Llama, Mistral
BLOCK_SIZE, num_warps = calculate_settings(vocab_size)
logsumexp = torch.empty(n_rows, dtype=torch.float32, device="cuda")
_cross_entropy_forward[(n_rows,)](
logits,
logits.stride(0),
losses,
logsumexp,
labels,
VOCAB_SIZE=vocab_size,
BLOCK_SIZE=BLOCK_SIZE,
num_warps=num_warps,
)
else:
# For large vocabs > 65336 like Gemma 256K
logsumexp = torch.empty(
(
n_rows,
n_chunks,
),
dtype=torch.float32,
device="cuda",
)
_chunked_cross_entropy_forward[
(
n_rows,
n_chunks,
)
](
logits,
logits.stride(0),
losses,
logsumexp,
labels,
VOCAB_SIZE=vocab_size,
N_CHUNKS=n_chunks,
BLOCK_SIZE=MAX_FUSED_SIZE,
num_warps=32,
)
# logsumexp(chunked_logsumexp) - x
# Do the -x separately
logsumexp = torch.logsumexp(logsumexp, dim=1) # Row sum
losses += logsumexp
losses.masked_fill_(labels == -100, 0) # Don't forget to mask padding out!
return losses, logsumexp
def _cross_entropy_backward_impl(dlosses, logits, logsumexp, labels):
n_rows, vocab_size = logits.shape
BLOCK_SIZE = 4096
div, mod = divmod(vocab_size, BLOCK_SIZE)
n_blocks = div + (mod != 0)
_cross_entropy_backward[
(
n_rows,
n_blocks,
)
](
logits,
logits.stride(0),
dlosses,
dlosses.stride(0),
logsumexp,
labels,
VOCAB_SIZE=vocab_size,
BLOCK_SIZE=BLOCK_SIZE,
num_warps=8,
)
return logits
extensions/thunder/unsloth/kernels/rope_embedding.py
+154
View File
@@ -0,0 +1,154 @@
# Copyright 2023-present Daniel Han-Chen & the Unsloth team. All rights reserved.
#
# Licensed under the Apache License, Version 2.0 (the "License");
# you may not use this file except in compliance with the License.
# You may obtain a copy of the License at
#
# http://www.apache.org/licenses/LICENSE-2.0
#
# Unless required by applicable law or agreed to in writing, software
# distributed under the License is distributed on an "AS IS" BASIS,
# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
# See the License for the specific language governing permissions and
# limitations under the License.
from litgpt.constants import _TRITON_AVAILABLE
from .utils import calculate_settings
if _TRITON_AVAILABLE:
import triton
import triton.language as tl
ROPE_GROUP_SIZE = 4
@triton.heuristics(
{
"BACKWARD_PASS": lambda args: args["BACKWARD_PASS"],
}
)
@triton.jit
def _rope_embedding(
Q,
Q_row_stride,
cos,
cos_row_stride,
sin,
sin_row_stride,
seqlen,
head_dim: tl.constexpr,
n_heads: tl.constexpr,
BACKWARD_PASS: tl.constexpr,
BLOCK_SIZE: tl.constexpr,
ROPE_GROUP_SIZE: tl.constexpr = 4,
):
"""
Calculates the RoPE Embedding quickly
RoPE is Q * cos + rotate_half(Q) * sin
See our blog post for more info
"""
row_position = tl.program_id(0)
group_head_position = tl.program_id(1)
col_offsets = tl.arange(0, BLOCK_SIZE)
half_head_dim = head_dim // 2
mask = col_offsets < half_head_dim
sin1 = tl.load(sin + (row_position % seqlen) * sin_row_stride + half_head_dim * 0 + col_offsets, mask=mask, other=0)
cos1 = tl.load(cos + (row_position % seqlen) * cos_row_stride + half_head_dim * 0 + col_offsets, mask=mask, other=0)
if BACKWARD_PASS:
# See our blog post for more info.
sin1 = -sin1
pass
# [TODO] Autotune ROPE_GROUP_SIZE to be 1, 2, 4, 8
head_start = group_head_position * ROPE_GROUP_SIZE
head_end = min((head_start + ROPE_GROUP_SIZE), n_heads)
# 10% Faster kernel from [HuyNguyen-hust](https://github.com/unslothai/unsloth/pull/238)
for k in range(head_start, head_end):
offs_q1 = row_position * Q_row_stride + k * head_dim + col_offsets
offs_q2 = row_position * Q_row_stride + k * head_dim + col_offsets + half_head_dim
# For Gemma - sometimes RoPE must be done in float32 and not bfloat16
Q1 = tl.load(Q + offs_q1, mask=mask, other=0).to(sin1.dtype)
Q2 = tl.load(Q + offs_q2, mask=mask, other=0).to(sin1.dtype)
tl.store(Q + offs_q1, Q1 * cos1 - Q2 * sin1, mask=mask)
tl.store(Q + offs_q2, Q2 * cos1 + Q1 * sin1, mask=mask)
pass
pass
def _rope_embedding_forward_impl(Q, cos, sin):
Q = Q.transpose(1, 2).clone()
cos, sin = cos.squeeze(), sin.squeeze()
batch, seq_len, n_heads, head_dim = Q.shape
Q = Q.reshape(batch * seq_len, n_heads * head_dim)
n_rows, n_cols = Q.shape
assert seq_len <= cos.shape[0]
# [TODO] Changing blocksize to head_dim//2 seems to have
# some concurrency / un-deterministic issues.
BLOCK_SIZE, num_warps = calculate_settings(head_dim // 2) # (head_dim//2)
# group_size = 4 # 4 or 8, too large group_size can hurt performance.
div, mod = divmod(n_heads, ROPE_GROUP_SIZE)
n_groups = div + (mod != 0)
_rope_embedding[
(
n_rows,
n_groups,
)
](
Q,
Q.stride(0),
cos,
cos.stride(0),
sin,
sin.stride(0),
seq_len,
head_dim,
n_heads,
BACKWARD_PASS=False,
BLOCK_SIZE=BLOCK_SIZE,
num_warps=num_warps,
)
Q = Q.view(batch, seq_len, n_heads, head_dim)
Q = Q.transpose(1, 2)
return Q, cos, sin, n_groups, BLOCK_SIZE, num_warps
def _rope_embedding_backward_impl(dY, cos, sin, n_groups, BLOCK_SIZE, num_warps):
dY = dY.transpose(1, 2)
batch, seq_len, n_heads, head_dim = dY.shape
dY = dY.reshape(batch * seq_len, n_heads * head_dim)
# Must be reshape not view
n_rows, n_cols = dY.shape
_rope_embedding[
(
n_rows,
n_groups,
)
](
dY,
dY.stride(0),
cos,
cos.stride(0),
sin,
sin.stride(0),
seq_len,
head_dim,
n_heads,
BACKWARD_PASS=True,
BLOCK_SIZE=BLOCK_SIZE,
num_warps=num_warps,
)
dY = dY.view(batch, seq_len, n_heads, head_dim)
dY = dY.transpose(1, 2)
return dY
extensions/thunder/unsloth/kernels/swiglu.py
+134
View File
@@ -0,0 +1,134 @@
# Copyright 2023-present Daniel Han-Chen & the Unsloth team. All rights reserved.
#
# Licensed under the Apache License, Version 2.0 (the "License");
# you may not use this file except in compliance with the License.
# You may obtain a copy of the License at
#
# http://www.apache.org/licenses/LICENSE-2.0
#
# Unless required by applicable law or agreed to in writing, software
# distributed under the License is distributed on an "AS IS" BASIS,
# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
# See the License for the specific language governing permissions and
# limitations under the License.
import torch
from litgpt.constants import _TRITON_AVAILABLE
if _TRITON_AVAILABLE:
import triton
import triton.language as tl
@triton.jit
def _fg_kernel(
e,
g,
h,
n_elements,
BLOCK_SIZE: tl.constexpr,
):
block_idx = tl.program_id(0)
offsets = block_idx * BLOCK_SIZE + tl.arange(0, BLOCK_SIZE)
mask = offsets < n_elements
e_row = tl.load(e + offsets, mask=mask, other=0).to(tl.float32)
g_row = tl.load(g + offsets, mask=mask, other=0) # .to(tl.float32)
# f = e * sigmoid(e)
f_row = e_row * tl.sigmoid(e_row) # e_row / (1 + tl.exp(-e_row))
f_row = f_row.to(g_row.dtype) # Exact copy from HF
# h = f * g
h_row = f_row * g_row
# Store h
tl.store(h + offsets, h_row, mask=mask)
pass
def swiglu_fg_kernel(e, g):
batch, seq_len, hd = e.shape
n_elements = e.numel()
h = torch.empty((batch, seq_len, hd), dtype=e.dtype, device="cuda")
grid = lambda meta: (triton.cdiv(n_elements, meta["BLOCK_SIZE"]),)
_fg_kernel[grid](
e,
g,
h,
n_elements,
BLOCK_SIZE=1024,
)
return h
pass
@triton.jit
def _DWf_DW_dfg_kernel(
DW,
e,
g,
n_elements,
BLOCK_SIZE: tl.constexpr,
):
"""
e = e.float()
se = 1.0 / (1.0 + torch.exp(-e))
f = (se * e).to(dtype)
h = f * g
df = DW * f
dg = DW * g
de = (dg.float() * se * (1.0 + e * (1.0 - se))).to(dtype)
"""
block_idx = tl.program_id(0)
offsets = block_idx * BLOCK_SIZE + tl.arange(0, BLOCK_SIZE)
mask = offsets < n_elements
DW_row = tl.load(DW + offsets, mask=mask, other=0) # .to(tl.float32)
e_row = tl.load(e + offsets, mask=mask, other=0).to(tl.float32)
g_row = tl.load(g + offsets, mask=mask, other=0) # .to(tl.float32)
# e = e.float()
# se = 1.0 / (1.0 + torch.exp(-e))
se_row = tl.sigmoid(e_row) # 1.0 / (1.0 + tl.exp(-e_row))
# f = (se * e).to(dtype)
f_row = se_row * e_row
f_row = f_row.to(DW_row.dtype)
# h = f * g
h_row = f_row * g_row
# df = DW * f
df_row = DW_row * f_row
# dg = DW * g
dg_row = DW_row * g_row
# de = (dg.float() * se * (1.0 + e * (1.0 - se))).to(dtype)
de_row = dg_row.to(tl.float32) * se_row * (1.0 + e_row * (1.0 - se_row))
de_row = de_row.to(DW_row.dtype)
# Store derivatives in buffers
tl.store(DW + offsets, h_row, mask=mask) # h = f * g
tl.store(e + offsets, df_row, mask=mask) # df = DW * f
tl.store(g + offsets, de_row, mask=mask) # de
pass
def swiglu_DWf_DW_dfg_kernel(DW, e, g):
batch_seq_len, hd = e.shape
n_elements = e.numel()
grid = lambda meta: (triton.cdiv(n_elements, meta["BLOCK_SIZE"]),)
_DWf_DW_dfg_kernel[grid](
DW,
e,
g,
n_elements,
BLOCK_SIZE=1024,
)
return DW, e, g
pass
extensions/thunder/unsloth/kernels/utils.py
+41
View File
@@ -0,0 +1,41 @@
# Copyright 2023-present Daniel Han-Chen & the Unsloth team. All rights reserved.
#
# Licensed under the Apache License, Version 2.0 (the "License");
# you may not use this file except in compliance with the License.
# You may obtain a copy of the License at
#
# http://www.apache.org/licenses/LICENSE-2.0
#
# Unless required by applicable law or agreed to in writing, software
# distributed under the License is distributed on an "AS IS" BASIS,
# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
# See the License for the specific language governing permissions and
# limitations under the License.
from litgpt.constants import _TRITON_AVAILABLE
if _TRITON_AVAILABLE:
import triton
MAX_FUSED_SIZE = 65536 # 2**16
next_power_of_2 = triton.next_power_of_2
def calculate_settings(n):
BLOCK_SIZE = next_power_of_2(n)
if BLOCK_SIZE > MAX_FUSED_SIZE:
raise RuntimeError(
f"Cannot launch Triton kernel since n = {n} exceeds the maximum CUDA blocksize = {MAX_FUSED_SIZE}."
)
num_warps = 4
if BLOCK_SIZE >= 32768:
num_warps = 32
elif BLOCK_SIZE >= 8192:
num_warps = 16
elif BLOCK_SIZE >= 2048:
num_warps = 8
return BLOCK_SIZE, num_warps
pass
extensions/xla/README.md
+158
View File
@@ -0,0 +1,158 @@
# TPU support
This project utilizes [`Fabric`](https://lightning.ai/docs/fabric/stable), which supports TPUs via [PyTorch XLA](https://github.com/pytorch/xla).
> [!NOTE]
> This guide assumes that you have already set-up your [Google Cloud environment](https://cloud.google.com/run/docs/setup).
To set up a Google Cloud instance with a TPU v4 VM, run the following commands:
```shell
gcloud compute tpus tpu-vm create litgpt --version=tpu-vm-v4-base --accelerator-type=v4-8 --zone=us-central2-b
gcloud compute tpus tpu-vm ssh litgpt --zone=us-central2-b
```
You can also choose a different TPU type. To do so, change the `version`, `accelerator-type`, and `zone` arguments. Find all regions and zones [here](https://cloud.google.com/tpu/docs/regions-zones).
<details>
<summary>Multihost caveats</summary>
TPU v4-8 uses a single host. SSH'ing into the machine and running commands manually will only work when using a single host (1 slice in the TPU pod).
In multi-host environments, such as larger TPU pod slices, it's necessary to launch all commands on all hosts simultaneously to avoid hangs.
For local development, it is advisable to upload a zip file containing all your current changes and execute it inside the VM from your personal computer:
```shell
# Zip the local directory, excluding large directories from the zip. You may want to keep them.
zip -r local_changes.zip . -x ".git/*" "checkpoints/*" "data/*" "out/*"
# Copy the .zip file to the TPU VM
gcloud compute tpus tpu-vm scp --worker=all local_changes.zip "litgpt:~"
# Unzip on each host
gcloud compute tpus tpu-vm ssh litgpt --worker=all --command="cd ~; unzip -q -o local_changes.zip"
# Example of a typical workflow
gcloud compute tpus tpu-vm ssh tmp --worker=all --command="cd ~; bash install_dependencies.sh"
gcloud compute tpus tpu-vm ssh tmp --worker=all --command="cd ~; bash prepare_checkpoints.sh"
gcloud compute tpus tpu-vm ssh tmp --worker=all --command="cd ~; bash run_desired_script.sh"
# This will allow you to kill all python processes on all workers
gcloud compute tpus tpu-vm ssh tmp --worker=all --command="pkill -e python"
```
Notice how the commands to install the environment and prepare checkpoints need to be run on all workers, since the filesystem
for each worker (host) is not shared.
For the rest of this tutorial, it will be assumed that it is being run on a single host for simplicity.
</details>
Once inside the machine, clone the repository and install the dependencies:
```shell
git clone https://github.com/Lightning-AI/litgpt
cd litgpt
pip install .
```
Install Optimized BLAS:
```shell
sudo apt update
sudo apt install libopenblas-dev
```
Since LitGPT requires a torch version newer than torch 2.0.0, manually install nightly builds of torch and torch_xla:
```shell
pip install https://storage.googleapis.com/tpu-pytorch/wheels/tpuvm/torch-nightly-cp38-cp38-linux_x86_64.whl
pip install https://storage.googleapis.com/tpu-pytorch/wheels/tpuvm/torch_xla-nightly-cp38-cp38-linux_x86_64.whl
```
While computations will run by default using the new PjRT runtime, it is recommended to set the following environment variables:
```shell
export ALLOW_MULTIPLE_LIBTPU_LOAD=1
export PJRT_DEVICE=TPU
```
> [!NOTE]
> An extensive guide on setup and available options can be found [here](https://cloud.google.com/tpu/docs/v4-users-guide).
Since a new machine was created, you may need to download pretrained weights.
They can be copied to the machine using `gcloud compute tpus tpu-vm scp`, or you can follow the steps described in our [downloading guide](../../tutorials/download_model_weights.md).
It is also recommended to set up a persistent disk from which to load checkpoints.
Follow [this guide](https://cloud.google.com/tpu/docs/setup-persistent-disk#setting_up_a_tpu_vm_and_a_persistent_disk) to do so.
Read-write disks are not supported in multihost VM setups, so persistent disks cannot be used to save checkpoints in that case.
Persistent disks can still be useful in read-only mode to load pretrained weights before finetuning or inference.
In multihost settings, FSDP will save checkpoint shards per host and consolidate them into a single checkpoint.
For safekeeping, it is recommended to upload the consolidated checkpoints to a Google Cloud bucket.
Alternatively, you can use the `scp` command to transfer these checkpoints from the TPU VM periodically, although this is not implemented in our scripts.
## Inference
This project provides custom versions of the regular recipes to run with XLA in the `xla` directory.
To generate text, use the following command:
```shell
python3 xla/generate/base.py --prompt "Hello, my name is" --num_samples 3
```
For the first generation, this command will take around 17 seconds as XLA needs to compile the graph.
Subsequent generations will take around 2 seconds.
## Fine-tuning
To get started fine-tuning Falcon 7B with adapter, run the following command:
```shell
python3 xla/scripts/prepare_alpaca.py --checkpoint_dir checkpoints/tiiuae/falcon-7b
python3 xla/finetune/adapter.py --checkpoint_dir checkpoints/tiiuae/falcon-7b --precision bf16-true
```
<details>
<summary>Multihost caveats</summary>
This script is configured to save "full" checkpoints, which isn't possible on multihost TPU VMs.
Here's how you can consolidate them together into a single one after training with `state_dict_type="sharded"`:
```shell
path_to_shards="out/adapter/alpaca/lit_model_adapter_finetuned"
mkdir -p $path_to_shards
workers=4 # 4 hosts
for ((i = 0; i < workers; i++)); do
# aggregate all shards locally
gcloud compute tpus tpu-vm scp --worker=$i "litgpt:${path_to_shards}/*" "${path_to_shards}/" --zone us-central2-b
done
# copy all shards to all workers
gcloud compute tpus tpu-vm scp --worker=all ${path_to_shards}/* "litgpt:${path_to_shards}/" --zone us-central2-b
# consolidate the shards in each worker
gcloud compute tpus tpu-vm ssh tmp --worker=all --command="python -m torch_xla.distributed.fsdp.consolidate_sharded_ckpts --ckpt_prefix ${path_to_shards}/checkpoint --ckpt_suffix '_rank-*-of-*.pth' --save_path ${path_to_shards}.pth" --zone us-central2-b
```
</details>
Since the TPU VM host RAM is limited (200 GB), we implement a technique to sequentially load and shard the checkpoint that can be enabled by
setting `reduce_cpu_memory_usage_during_load = True`. This is necessary to load falcon-40b.
To generate text with the adapter fine-tuned model weights, use the following command:
```shell
python3 xla/generate/adapter.py --checkpoint_dir checkpoints/tiiuae/falcon-7b --precision bf16-true --adapter_path out/adapter/alpaca/lit_model_adapter_finetuned.pth
```
> **Warning**
> Remember to delete your instance when you are done.
>
> ```shell
> gcloud compute tpus tpu-vm delete litgpt --zone=us-central2-b
> ```
## Computational Performance
Using the [adapter finetuning script](finetune/adapter.py) and XLA's FSDP implementation, a 49.57% MFU was achieved with Falcon 7B on a v4-32 (micro batch size 7), and a 39.67% MFU was achieved with Falcon 40B on a v4-512 (micro batch size 3) at a fixed 1034 maximum sequence length.
Since the TPU VM host has limited system memory (RAM) compared to device memory (HBM), specific techniques were implemented to limit peak RAM usage when loading the model and pretrained weights before sharding, as well as when saving sharded checkpoints.
A v4 chip has 32 GiB HBM, so with 4 devices per host (4 * 32 = 128 GiB HBM), each host has 188 GiB RAM, which is shared across the devices.
Therefore, any RAM allocation over 188/4 = 47 GiB would exceed the host's RAM capacity.
A ~24B parameter model on CPU (with half precision) would be the largest possible model under this setup without the techniques used in our scripts.
extensions/xla/__init__
+6
View File
@@ -0,0 +1,6 @@
import sys
from pathlib import Path
# support running without installing as a package, adding extensions to the Python path
wd = Path(__file__).parent.parent.resolve()
sys.path.append(str(wd))
extensions/xla/finetune/__init__
View File
extensions/xla/finetune/adapter.py
+284
View File
@@ -0,0 +1,284 @@
# Copyright Lightning AI. Licensed under the Apache License 2.0, see LICENSE file.
import os
import sys
import time
from pathlib import Path
import lightning as L
import torch
import torch_xla.core.xla_model as xm
from lightning.fabric.accelerators import XLAAccelerator
from lightning.fabric.loggers import CSVLogger
from lightning.fabric.strategies import XLAFSDPStrategy
from lightning.fabric.utilities import ThroughputMonitor, measure_flops
from litgpt.adapter import GPT, Block, Config, adapter_filter, mark_only_adapter_as_trainable
from litgpt.tokenizer import Tokenizer
from litgpt.utils import check_valid_checkpoint_dir, chunked_cross_entropy, estimate_flops, lazy_load, num_parameters
# support running without installing as a package
wd = Path(__file__).parents[3].resolve()
sys.path.append(str(wd))
from xla.generate.base import generate # noqa: E402
from xla.scripts.prepare_alpaca import generate_prompt # noqa: E402
from xla.utils import rank_print, sequential_load_and_fsdp_wrap # noqa: E402
eval_interval = 200
save_interval = 200
eval_iters = 100
eval_max_new_tokens = 100
log_interval = 1
devices = XLAAccelerator.auto_device_count()
# the state of very large models will not fit on the system RAM, this flag can alleviate it by loading it on each rank
# sequentially
reduce_cpu_memory_usage_during_load = False
# Hyperparameters
learning_rate = 3e-3
batch_size = 4
micro_batch_size = batch_size
gradient_accumulation_iters = batch_size // micro_batch_size
assert gradient_accumulation_iters > 0
epoch_size = 50000 # train dataset size
num_epochs = 5
max_iters = num_epochs * (epoch_size // micro_batch_size) // devices
weight_decay = 0.02
warmup_steps = 2 * (epoch_size // micro_batch_size) // devices // gradient_accumulation_iters # 2 epochs
hparams = {k: v for k, v in locals().items() if isinstance(v, (int, float, str)) and not k.startswith("_")}
def setup(
*,
data_dir: Path = Path("data/alpaca"),
checkpoint_dir: Path = Path("checkpoints/tiiuae/falcon-7b"),
out_dir: Path = Path("out/adapter/alpaca"),
precision: str = "bf16-true",
) -> None:
if devices > 1:
strategy = XLAFSDPStrategy(
auto_wrap_policy={Block},
activation_checkpointing_policy={Block},
state_dict_type="full", # change to "sharded" in multi-host environments where the filesystem is not shared
sequential_save=True,
)
else:
strategy = "auto"
logger = CSVLogger(out_dir.parent, out_dir.name, flush_logs_every_n_steps=log_interval)
fabric = L.Fabric(devices=devices, strategy=strategy, precision=precision, loggers=logger)
rank_print(fabric, hparams)
fabric.launch(main, data_dir, checkpoint_dir, out_dir)
def main(fabric: L.Fabric, data_dir: Path, checkpoint_dir: Path, out_dir: Path) -> None:
check_valid_checkpoint_dir(checkpoint_dir)
fabric.seed_everything(1337) # same seed for every process to init model (FSDP)
if fabric.global_rank == 0:
os.makedirs(out_dir, exist_ok=True)
train_data = torch.load(data_dir / "train.pt")
val_data = torch.load(data_dir / "test.pt")
config = Config.from_name(name=checkpoint_dir.name, adapter_start_layer=0)
checkpoint_path = checkpoint_dir / "lit_model.pth"
rank_print(fabric, f"Loading model {str(checkpoint_path)!r} with {config.__dict__}")
if reduce_cpu_memory_usage_during_load:
model = sequential_load_and_fsdp_wrap(fabric, lambda: GPT(config), checkpoint_path)
else:
with fabric.init_module(empty_init=False):
model = GPT(config)
checkpoint = lazy_load(checkpoint_path)
# strict=False because missing keys due to adapter weights not contained in state dict
model.load_state_dict(checkpoint, strict=False)
model = fabric.setup_module(model)
# mark as trainable only after sharding due to https://github.com/pytorch/xla/pull/5484
mark_only_adapter_as_trainable(model)
# these are not correct in the sharding case
rank_print(fabric, f"Number of trainable parameters: {num_parameters(model, requires_grad=True):,}")
rank_print(fabric, f"Number of non-trainable parameters: {num_parameters(model, requires_grad=False):,}")
trainable_params = [p for p in model.parameters() if p.requires_grad]
optimizer = torch.optim.SGD(trainable_params, lr=learning_rate)
optimizer = fabric.setup_optimizers(optimizer)
fabric.seed_everything(1337 + fabric.global_rank)
train_time = time.perf_counter()
train(fabric, model, optimizer, train_data, val_data, checkpoint_dir, out_dir)
rank_print(fabric, f"Training time: {(time.perf_counter() - train_time):.2f}s")
# Save the final checkpoint at the end of training
save_path = out_dir / "lit_model_adapter_finetuned.pth"
save_adapter_checkpoint(fabric, model, save_path)
def train(
fabric: L.Fabric,
model: GPT,
optimizer: torch.optim.Optimizer,
train_data: list[dict],
val_data: list[dict],
checkpoint_dir: Path,
out_dir: Path,
) -> None:
tokenizer = Tokenizer(checkpoint_dir)
longest_seq_length = get_longest_seq_length(train_data)
model.max_seq_length = longest_seq_length
# to avoid recompilation, this script is configured to pad batches to the `longest_seq_length`
fabric.print(
f"The longest sequence length in the train data is {longest_seq_length}, the model's maximum sequence length is"
f" {model.max_seq_length} and context length is {model.config.block_size}"
)
with torch.device("meta"):
meta_model = GPT(model.config)
mark_only_adapter_as_trainable(meta_model)
# "estimated" is not as precise as "measured". Estimated is optimistic but widely used in the wild.
# When comparing MFU or FLOP numbers with other projects that use estimated FLOPs,
# consider passing `flops_per_batch=estimated_flops` instead
estimated_flops = estimate_flops(meta_model, training=True) * micro_batch_size
rank_print(fabric, f"Estimated TFLOPs: {estimated_flops * fabric.world_size / 1e12:.2f}")
# this assumes that all samples have a fixed length equal to the longest sequence length
# which is most likely false during finetuning
x = torch.randint(0, 1, (micro_batch_size, longest_seq_length))
forward_fn = lambda: meta_model(x) # noqa: F821
loss_fn = lambda y: chunked_cross_entropy(y, x, chunk_size=0) # noqa: F821
measured_flops = measure_flops(meta_model, forward_fn, loss_fn)
rank_print(fabric, f"Measured TFLOPs: {measured_flops * fabric.world_size / 1e12:.2f}")
del meta_model, x
throughput = ThroughputMonitor(fabric, window_size=50)
step_count = 0
total_t0 = time.perf_counter()
xm.mark_step()
for iter_num in range(1, max_iters + 1):
if step_count <= warmup_steps:
# linear warmup
lr = learning_rate * step_count / warmup_steps
for param_group in optimizer.param_groups:
param_group["lr"] = lr
iter_t0 = time.perf_counter()
input_ids, targets = get_batch(fabric, train_data, longest_seq_length)
is_accumulating = iter_num % gradient_accumulation_iters != 0
with fabric.no_backward_sync(model, enabled=is_accumulating):
logits = model(input_ids, lm_head_chunk_size=128)
xm.mark_step()
# shift the targets such that output n predicts token n+1
logits[-1] = logits[-1][..., :-1, :]
loss = chunked_cross_entropy(logits, targets[..., 1:])
fabric.backward(loss / gradient_accumulation_iters)
xm.mark_step()
if not is_accumulating:
optimizer.step()
optimizer.zero_grad()
step_count += 1
else:
xm.mark_step()
if iter_num % log_interval == 0:
t1 = time.perf_counter()
throughput.update(
time=t1 - total_t0,
batches=iter_num,
samples=iter_num * micro_batch_size,
lengths=iter_num * micro_batch_size * longest_seq_length,
flops=measured_flops * log_interval,
)
throughput.compute_and_log(step=iter_num)
rank_print(
fabric,
f"iter {iter_num} step {step_count}:"
# uncomment to print the loss. this will considerably slow down the iteration times
# + f" loss {loss.item():.4f},"
+ f" iter time: {(t1 - iter_t0) * 1000:.2f}ms"
+ (" (optimizer.step)" if not is_accumulating else ""),
)
if not is_accumulating and step_count % eval_interval == 0:
t0 = time.perf_counter()
val_loss = validate(fabric, model, val_data, tokenizer, longest_seq_length)
t1 = time.perf_counter() - t0
rank_print(fabric, f"step {iter_num}: val loss {val_loss.item():.4f}, val time: {t1 * 1000:.2f}ms")
fabric.barrier()
if not is_accumulating and step_count % save_interval == 0:
checkpoint_path = out_dir / f"iter-{iter_num:06d}-ckpt.pth"
save_adapter_checkpoint(fabric, model, checkpoint_path)
# xla does not support `inference_mode`: RuntimeError: Cannot set version_counter for inference tensor
@torch.no_grad()
def validate(
fabric: L.Fabric, model: GPT, val_data: list[dict], tokenizer: Tokenizer, longest_seq_length: int
) -> torch.Tensor:
rank_print(fabric, "Validating ...")
model.eval()
losses = torch.zeros(eval_iters)
xm.mark_step()
for k in range(eval_iters):
input_ids, targets = get_batch(fabric, val_data, longest_seq_length)
logits = model(input_ids)
xm.mark_step()
losses[k] = chunked_cross_entropy(logits[..., :-1, :], targets[..., 1:], chunk_size=0)
val_loss = losses.mean()
# produce an example:
instruction = "Recommend a movie for me to watch during the weekend and explain the reason."
rank_print(fabric, instruction)
sample = {"instruction": instruction, "input": ""}
prompt = generate_prompt(sample)
encoded = tokenizer.encode(prompt, device=fabric.device)
with fabric.init_tensor():
# do not set `max_seq_length=max_returned_token` because memory is not a concern here
model.set_kv_cache(batch_size=1)
output = generate(model, encoded, max_returned_tokens=len(encoded) + eval_max_new_tokens, temperature=0.8)
model.clear_kv_cache()
output = tokenizer.decode(output)
rank_print(fabric, output)
model.train()
return val_loss
def get_batch(fabric: L.Fabric, data: list[dict], longest_seq_length: int) -> tuple[torch.Tensor, torch.Tensor]:
ix = torch.randint(len(data), (micro_batch_size,))
input_ids = [data[i]["input_ids"].type(torch.int64) for i in ix]
labels = [data[i]["labels"].type(torch.int64) for i in ix]
def pad_right(x, pad_id):
# pad right using a fixed longest sequence length to avoid recompilation
n = longest_seq_length - len(x)
return torch.cat((x, torch.full((n,), pad_id, dtype=x.dtype)))
x = torch.stack([pad_right(x, pad_id=0) for x in input_ids])
y = torch.stack([pad_right(x, pad_id=-1) for x in labels])
x, y = fabric.to_device((x, y))
return x, y
def get_longest_seq_length(data: list[dict]) -> int:
# find out the minimum max_seq_length required during fine-tuning (saves memory!)
return max(len(d["input_ids"]) for d in data)
def save_adapter_checkpoint(fabric: L.Fabric, model: torch.nn.Module, file_path: Path) -> None:
rank_print(fabric, f"Saving adapter weights to {str(file_path)!r}")
fabric.save(file_path, {"model": model}, filter={"model": adapter_filter})
if __name__ == "__main__":
from jsonargparse import CLI
CLI(setup)
extensions/xla/generate/__init__
View File
extensions/xla/generate/adapter.py
+132
View File
@@ -0,0 +1,132 @@
# Copyright Lightning AI. Licensed under the Apache License 2.0, see LICENSE file.
import sys
import time
from pathlib import Path
import lightning as L
from lightning.fabric.accelerators import XLAAccelerator
from lightning.fabric.strategies import XLAFSDPStrategy
from litgpt import Tokenizer
from litgpt.adapter import GPT, Block, Config
from litgpt.prompts import Alpaca
from litgpt.utils import check_valid_checkpoint_dir, lazy_load
# support running without installing as a package
wd = Path(__file__).parents[3].resolve()
sys.path.append(str(wd))
from xla.generate.base import generate # noqa: E402
from xla.utils import rank_print # noqa: E402
def setup(
prompt: str = "What food do llamas eat?",
*,
input: str = "",
sys_prompt: str | None = None,
adapter_path: Path = Path("out/adapter/alpaca/lit_model_adapter_finetuned.pth"),
checkpoint_dir: Path = Path("checkpoints/tiiuae/falcon-7b"),
max_new_tokens: int = 100,
top_k: int | None = 50,
temperature: float = 0.8,
precision: str = "bf16-true",
) -> None:
"""Generates a response based on a given instruction and an optional input.
This script will only work with checkpoints from the instruction-tuned Adapter model.
See `xla/finetune/adapter.py`.
Args:
prompt: The prompt/instruction (Alpaca style).
input: Optional input (Alpaca style).
sys_prompt: Optional system prompt.
adapter_path: Path to the checkpoint with trained adapter weights, which are the output of
`xla/finetune/adapter.py`.
checkpoint_dir: The path to the checkpoint folder with pretrained model weights.
max_new_tokens: The number of generation steps to take.
top_k: The number of top most probable tokens to consider in the sampling process.
temperature: A value controlling the randomness of the sampling process. Higher values result in more random
samples.
precision: Indicates the Fabric precision setting to use.
"""
devices = XLAAccelerator.auto_device_count()
strategy = XLAFSDPStrategy(auto_wrap_policy={Block}) if devices > 1 else "auto"
fabric = L.Fabric(devices=devices, precision=precision, strategy=strategy)
fabric.launch(main, prompt, input, sys_prompt, adapter_path, checkpoint_dir, max_new_tokens, top_k, temperature)
def main(
fabric: L.Fabric,
prompt: str,
input: str,
sys_prompt: str | None,
adapter_path: Path,
checkpoint_dir: Path,
max_new_tokens: int,
top_k: int | None,
temperature: float,
) -> None:
check_valid_checkpoint_dir(checkpoint_dir)
config = Config.from_file(checkpoint_dir / "model_config.yaml", adapter_start_layer=0)
checkpoint_path = checkpoint_dir / "lit_model.pth"
rank_print(fabric, f"Loading model {str(checkpoint_path)!r} with {config.__dict__}", file=sys.stderr)
t0 = time.perf_counter()
with fabric.init_module(empty_init=True):
model = GPT(config)
rank_print(fabric, f"Time to instantiate model: {time.perf_counter() - t0:.02f} seconds.", file=sys.stderr)
t0 = time.perf_counter()
checkpoint = lazy_load(checkpoint_path)
adapter_checkpoint = lazy_load(adapter_path)
checkpoint.update(adapter_checkpoint.get("model", adapter_checkpoint))
model.load_state_dict(checkpoint)
rank_print(fabric, f"Time to load the model weights: {time.perf_counter() - t0:.02f} seconds.", file=sys.stderr)
model.eval()
model = fabric.setup_module(model)
tokenizer = Tokenizer(checkpoint_dir)
# TODO: Load prompt style from checkpoint and apply it here
prompt_style = Alpaca()
prompt = prompt_style.apply(prompt, sys_prompt=sys_prompt, input=input)
encoded = tokenizer.encode(prompt, device=fabric.device)
prompt_length = encoded.size(0)
max_returned_tokens = prompt_length + max_new_tokens
with fabric.init_tensor():
# set the max_seq_length to limit the memory usage to what we need
model.max_seq_length = max_returned_tokens
# enable the kv cache
model.set_kv_cache(batch_size=1)
t0 = time.perf_counter()
y = generate(
model,
encoded,
max_returned_tokens,
max_seq_length=max_returned_tokens,
temperature=temperature,
top_k=top_k,
eos_id=tokenizer.eos_id,
)
t = time.perf_counter() - t0
output = tokenizer.decode(y)
output = output.split("### Response:")[1] if "### Response:" in output else output
output = output.strip()
fabric.print(output)
tokens_generated = y.size(0) - prompt_length
rank_print(
fabric, f"\n\nTime for inference: {t:.02f} sec total, {tokens_generated / t:.02f} tokens/sec", file=sys.stderr
)
if __name__ == "__main__":
from jsonargparse import CLI
CLI(setup)
extensions/xla/generate/base.py
+184
View File
@@ -0,0 +1,184 @@
# Copyright Lightning AI. Licensed under the Apache License 2.0, see LICENSE file.
import sys
import time
from pathlib import Path
import lightning as L
import torch
import torch_xla.core.xla_model as xm
from lightning.fabric.accelerators import XLAAccelerator
from lightning.fabric.strategies import XLAFSDPStrategy
from litgpt import GPT, Config, Tokenizer
from litgpt.model import Block
from litgpt.utils import check_valid_checkpoint_dir, lazy_load
# support running without installing as a package
wd = Path(__file__).parents[3].resolve()
sys.path.append(str(wd))
from xla.utils import rank_print # noqa: E402
# xla does not support `inference_mode`: RuntimeError: Cannot set version_counter for inference tensor
@torch.no_grad()
def generate(
model: GPT,
idx: torch.Tensor,
max_returned_tokens: int,
*,
temperature: float = 1.0,
top_k: int | None = None,
eos_id: int | None = None,
) -> torch.Tensor:
"""Takes a conditioning sequence (prompt) as input and continues to generate as many tokens as requested.
The implementation of this function is modified from A. Karpathy's nanoGPT.
Args:
model: The model to use.
idx: Tensor of shape (T) with indices of the prompt sequence.
max_returned_tokens: The maximum number of tokens to return (given plus generated).
temperature: Scales the predicted logits by 1 / temperature.
top_k: If specified, only sample among the tokens with the k highest probabilities.
eos_id: If specified, stop generating any more token once the <eos> token is triggered.
"""
T = idx.size(0)
assert max_returned_tokens > T
if model.max_seq_length < max_returned_tokens - 1:
# rolling the kv cache based on the `input_pos` value would be necessary. However, doing so would introduce a
# data dependency on the `input_pos` tensor and impact model compilation. Since this setting is uncommon, we do
# not support it to avoid negatively impacting the overall speed
raise NotImplementedError(f"max_seq_length {model.max_seq_length} needs to be >= {max_returned_tokens - 1}")
device, dtype = idx.device, idx.dtype
# create an empty tensor of the expected final shape and fill in the current tokens
empty = torch.empty(max_returned_tokens, dtype=dtype, device=device)
empty[:T] = idx
idx = empty
# TODO: FSDP has an internal broadcasting issue, so we are forced to have this be of length 1 until it's fixed
input_pos = torch.tensor([0], device=device)
xm.mark_step()
# generate up to a fixed number of tokens
for _ in range(max_returned_tokens):
x = idx.index_select(0, input_pos).view(1, -1)
# forward
logits = model(x, input_pos)
logits = logits[0, -1] / temperature
# optionally crop the logits to only the top k options
if top_k is not None:
v, _ = torch.topk(logits, min(top_k, logits.size(-1)))
logits = torch.where(logits < v[[-1]], -float("Inf"), logits)
probs = torch.nn.functional.softmax(logits, dim=-1)
idx_next = torch.multinomial(probs, num_samples=1).to(dtype=dtype)
# advance
input_pos = input_pos[-1:] + 1
xm.mark_step()
# concatenate the new generation
idx = idx.index_copy(0, input_pos, idx_next)
# if <eos> token is triggered, return the output (stop generation)
if idx_next == eos_id:
return idx[:input_pos] # include the EOS token
return idx
def setup(
prompt: str = "What food do llamas eat?",
*,
num_samples: int = 1,
max_new_tokens: int = 100,
top_k: int | None = 50,
temperature: float = 0.8,
checkpoint_dir: Path = Path("checkpoints/tiiuae/falcon-7b"),
precision: str = "bf16-true",
) -> None:
"""Generates text samples based on a pre-trained model and tokenizer.
Args:
prompt: The prompt string to use for generating the samples.
num_samples: The number of text samples to generate.
max_new_tokens: The number of generation steps to take.
top_k: The number of top most probable tokens to consider in the sampling process.
temperature: A value controlling the randomness of the sampling process. Higher values result in more random
samples.
checkpoint_dir: The checkpoint directory to load.
precision: Indicates the Fabric precision setting to use.
"""
devices = XLAAccelerator.auto_device_count()
strategy = XLAFSDPStrategy(auto_wrap_policy={Block}) if devices > 1 else "auto"
fabric = L.Fabric(devices=devices, precision=precision, strategy=strategy)
fabric.launch(main, prompt, num_samples, max_new_tokens, top_k, temperature, checkpoint_dir)
def main(
fabric: L.Fabric,
prompt: str,
num_samples: int,
max_new_tokens: int,
top_k: int | None,
temperature: float,
checkpoint_dir: Path,
) -> None:
check_valid_checkpoint_dir(checkpoint_dir)
config = Config.from_file(checkpoint_dir / "model_config.yaml")
checkpoint_path = checkpoint_dir / "lit_model.pth"
rank_print(fabric, f"Loading model {str(checkpoint_path)!r} with {config.__dict__}", file=sys.stderr)
t0 = time.perf_counter()
with fabric.init_module(empty_init=True):
model = GPT(config)
rank_print(fabric, f"Time to instantiate model: {time.perf_counter() - t0:.02f} seconds.", file=sys.stderr)
t0 = time.perf_counter()
checkpoint = lazy_load(checkpoint_path)
model.load_state_dict(checkpoint.get("model", checkpoint))
rank_print(fabric, f"Time to load the model weights: {time.perf_counter() - t0:.02f} seconds.", file=sys.stderr)
model.eval()
model = fabric.setup_module(model)
tokenizer = Tokenizer(checkpoint_dir)
encoded = tokenizer.encode(prompt, device=fabric.device)
prompt_length = encoded.size(0)
max_returned_tokens = prompt_length + max_new_tokens
with fabric.init_tensor():
# set the max_seq_length to limit the memory usage to what we need
model.max_seq_length = max_returned_tokens
L.seed_everything(1234)
for i in range(num_samples):
with fabric.init_tensor():
# enable the kv cache
model.set_kv_cache(batch_size=1)
t0 = time.perf_counter()
y = generate(model, encoded, max_returned_tokens, temperature=temperature, top_k=top_k)
t = time.perf_counter() - t0
fabric.print(tokenizer.decode(y))
tokens_generated = y.size(0) - prompt_length
rank_print(
fabric,
f"Time for inference {i + 1}: {t:.02f} sec total, {tokens_generated / t:.02f} tokens/sec",
file=sys.stderr,
)
if __name__ == "__main__":
from jsonargparse import CLI
CLI(setup)
extensions/xla/scripts/__init__
View File
extensions/xla/scripts/prepare_alpaca.py
+146
View File
@@ -0,0 +1,146 @@
# Copyright Lightning AI. Licensed under the Apache License 2.0, see LICENSE file.
"""Implementation derived from https://github.com/tloen/alpaca-lora"""
import json
from pathlib import Path
import torch
import yaml
from lightning_utilities.core.imports import RequirementCache
from torch.utils.data import random_split
from tqdm import tqdm
from litgpt.tokenizer import Tokenizer
from litgpt.utils import CLI
def prepare(
destination_path: Path = Path("data/alpaca"),
checkpoint_dir: Path = Path("checkpoints/stabilityai/stablelm-base-alpha-3b"),
val_split_fraction: float = 0.03865, # to get exactly 2000 validation samples,
seed: int = 42,
mask_inputs: bool = False, # as in alpaca-lora
data_file_name: str = "alpaca_data_cleaned_archive.json",
data_file_url: str = "https://raw.githubusercontent.com/tloen/alpaca-lora/main/alpaca_data_cleaned_archive.json",
ignore_index: int = -100,
max_seq_length: int | None = None,
) -> None:
"""Prepare the Alpaca dataset for instruction tuning.
The output is a training and test dataset saved as `train.pt` and `test.pt`,
which stores the preprocessed and tokenized prompts and labels.
"""
if max_seq_length is None:
with open(checkpoint_dir / "model_config.yaml", encoding="utf-8") as file:
config = yaml.safe_load(file)
max_seq_length = config["block_size"]
destination_path.mkdir(parents=True, exist_ok=True)
data_file_path = destination_path / data_file_name
print("Loading data file...")
download_if_missing(data_file_path, data_file_url)
with open(data_file_path, encoding="utf-8") as file:
data = json.load(file)
print("Loading tokenizer...")
tokenizer = Tokenizer(checkpoint_dir)
# Partition the dataset into train and test
train_set, test_set = random_split(
data, [1.0 - val_split_fraction, val_split_fraction], generator=torch.Generator().manual_seed(seed)
)
train_set, test_set = list(train_set), list(test_set)
print(f"train has {len(train_set):,} samples")
print(f"test has {len(test_set):,} samples")
print("Processing train split ...")
train_set = [
prepare_sample(
example=sample,
tokenizer=tokenizer,
max_length=max_seq_length,
mask_inputs=mask_inputs,
ignore_index=ignore_index,
)
for sample in tqdm(train_set)
]
torch.save(train_set, destination_path / "train.pt")
print("Processing test split ...")
test_set = [
prepare_sample(
example=sample,
tokenizer=tokenizer,
max_length=max_seq_length,
mask_inputs=mask_inputs,
ignore_index=ignore_index,
)
for sample in tqdm(test_set)
]
torch.save(test_set, destination_path / "test.pt")
def download_if_missing(file_path: Path, file_url: str) -> None:
"""Downloads the raw json data file and saves it in the given destination."""
if file_path.exists() and file_path.stat().st_size > 0:
return
requests_available = RequirementCache("requests")
if not requests_available:
raise ModuleNotFoundError(str(requests_available))
import requests
with open(file_path, "w", encoding="utf-8") as f:
f.write(requests.get(file_url).text)
def prepare_sample(example: dict, tokenizer: Tokenizer, max_length: int, mask_inputs: bool, ignore_index: int) -> dict:
"""Processes a single sample.
Each sample in the dataset consists of:
- instruction: A string describing the task
- input: A string holding a special input value for the instruction.
This only applies to some samples, and in others this is empty.
- output: The response string
This function processes this data to produce a prompt text and a label for
supervised training. The prompt text is formed as a single message including both
the instruction and the input. The label/target is the same message but with the
response attached.
Finally, both the prompt and the label get tokenized. If desired, all tokens
in the label that correspond to the original input prompt get masked out (default).
"""
full_prompt = generate_prompt(example)
full_prompt_and_response = full_prompt + example["output"]
encoded_full_prompt = tokenizer.encode(full_prompt, max_length=max_length)
encoded_full_prompt_and_response = tokenizer.encode(full_prompt_and_response, eos=True, max_length=max_length)
# The labels are the full prompt with response, but with the prompt masked out
labels = encoded_full_prompt_and_response.clone()
if mask_inputs:
labels[: len(encoded_full_prompt)] = ignore_index
return {**example, "input_ids": encoded_full_prompt_and_response, "labels": labels}
def generate_prompt(example: dict) -> str:
"""Generates a standardized message to prompt the model with an instruction, optional input and a
'response' field."""
if example["input"]:
return (
"Below is an instruction that describes a task, paired with an input that provides further context. "
"Write a response that appropriately completes the request.\n\n"
f"### Instruction:\n{example['instruction']}\n\n### Input:\n{example['input']}\n\n### Response:"
)
return (
"Below is an instruction that describes a task. "
"Write a response that appropriately completes the request.\n\n"
f"### Instruction:\n{example['instruction']}\n\n### Response:"
)
if __name__ == "__main__":
CLI(prepare)
extensions/xla/utils.py
+114
View File
@@ -0,0 +1,114 @@
# Copyright Lightning AI. Licensed under the Apache License 2.0, see LICENSE file.
import itertools
from collections.abc import Callable
from functools import partial
from pathlib import Path
from typing import Any
import lightning as L
import torch
from lightning.fabric.strategies.xla_fsdp import XLAFSDPStrategy, _activation_checkpointing_auto_wrapper
from lightning_utilities.core.rank_zero import rank_prefixed_message
from litgpt import GPT
def rank_print(fabric: L.Fabric, message: object, *, flush: bool = True, **kwargs: Any) -> None:
if fabric.local_rank == 0:
message = str(message)
# let each host print, but only on rank 0
message = rank_prefixed_message(message, fabric.global_rank)
# TPU VM will only print when the script finishes if `flush=False`
print(message, flush=flush, **kwargs)
def materialize_parameters(module: torch.nn.Module, device: torch.device) -> None:
for module_name, module in module.named_modules():
if any(
param.is_meta for param in itertools.chain(module.parameters(recurse=False), module.buffers(recurse=False))
):
module.to_empty(device=device, recurse=False)
module.reset_parameters()
def sequential_load_and_fsdp_wrap(
fabric: L.Fabric, get_model: Callable[[], GPT], checkpoint_path: Path
) -> torch.nn.Module:
assert fabric._launched
# similar logic could be implemented for regular FSDP, but this implementation is specific to XLAFSDP
assert isinstance(fabric.strategy, XLAFSDPStrategy)
with fabric.init_module(empty_init=False), torch.device("meta"):
model = get_model()
# TODO: this could be made faster by broadcasting in separate process groups for each host
if fabric.local_rank == 0:
# load the full checkpoint on a single rank to limit the system memory usage
state_dict = torch.load(checkpoint_path, map_location="cpu", mmap=False) # mmap=True hangs
else:
# XLA cannot broadcast different number of tensors or different shapes in each rank. To get around this
# limitation, we need to load the checkpoint on meta device to get the correct number of tensors and materialize
# them as necessary
state_dict = torch.load(checkpoint_path, map_location="meta", mmap=False)
fsdp_kwargs = fabric.strategy._parse_fsdp_kwargs()
if "auto_wrapper_callable" in fsdp_kwargs:
# includes activation checkpointing if configured
wrap = fsdp_kwargs.pop("auto_wrapper_callable")
else:
wrap = partial(_activation_checkpointing_auto_wrapper, set())
fsdp_kwargs.pop("auto_wrap_policy", None) # this needs to be removed or else root wrapping would error
for i, block in enumerate(model.transformer.h):
rank_print(fabric, f"Broadcasting transformer block {i}")
# get the relevant piece of the state dict
to_load = {}
for param_name, _ in block.named_parameters():
if (key := f"transformer.h.{i}.{param_name}") not in state_dict:
continue
param = state_dict.pop(key)
if not param.is_meta:
to_load[param_name] = param
else:
# materialize this parameter for broadcast to work
to_load[param_name] = torch.empty_like(param, device="cpu")
to_load = fabric.broadcast(to_load)
rank_print(fabric, f"Loading transformer block {i}")
keys = block.load_state_dict(to_load, strict=False, assign=True)
assert not keys.unexpected_keys
# materialize any leftover meta parameters, regular FSDP does it automatically
materialize_parameters(block, torch.device("cpu")) # init on CPU, FSDP will shard and move it
# XLA FSDP only supports fp32 parameters. If the checkpoint had a different dtype, this needs to be converted
# since we are loading with assign=True
block = block.to(torch.float32)
# shard the block
rank_print(fabric, f"Wrapping transformer block {i}")
wrapped_block = wrap(block, **fsdp_kwargs)
model.transformer.h[i] = wrapped_block
# load the rest of the state_dict, this assumes that all keys need to be loaded
# an alternative technique would be to do load the rest of the state dict at once, but we want to materialize
# and move the params to the xla device to reduce the system memory usage
for key in list(state_dict):
rank_print(fabric, f"Loading {key}")
param = state_dict.pop(key)
if param.is_meta:
# materialize this parameter for broadcast to work
param = torch.empty_like(param, device="cpu")
param = fabric.broadcast(param)
param = param.to(device=fabric.device, dtype=torch.float32)
keys = model.load_state_dict({key: param}, strict=False, assign=True)
assert not keys.unexpected_keys
assert not state_dict
# materialize any leftover meta parameters, regular FSDP does it automatically
rank_print(fabric, "Materializing leftover parameters")
materialize_parameters(model, fabric.device)
return model
litgpt/__init__.py
+30
View File
@@ -0,0 +1,30 @@
# Copyright Lightning AI. Licensed under the Apache License 2.0, see LICENSE file.
import logging
import re
from litgpt.api import LLM
from litgpt.config import Config
from litgpt.model import GPT # needs to be imported before config
from litgpt.prompts import PromptStyle
from litgpt.tokenizer import Tokenizer
from litgpt.utils import CheckpointValidationResult, estimate_model_memory, validate_checkpoint
# Suppress excessive warnings, see https://github.com/pytorch/pytorch/issues/111632
pattern = re.compile(".*Profiler function .* will be ignored")
logging.getLogger("torch._dynamo.variables.torch").addFilter(lambda record: not pattern.search(record.getMessage()))
# Avoid printing state-dict profiling output at the WARNING level when saving a checkpoint
logging.getLogger("torch.distributed.fsdp._optim_utils").disabled = True
logging.getLogger("torch.distributed.fsdp._debug_utils").disabled = True
__all__ = [
"LLM",
"GPT",
"Config",
"PromptStyle",
"Tokenizer",
"CheckpointValidationResult",
"validate_checkpoint",
"estimate_model_memory",
]
litgpt/__main__.py
+84
View File
@@ -0,0 +1,84 @@
# Copyright Lightning AI. Licensed under the Apache License 2.0, see LICENSE file.
import warnings
import torch
from jsonargparse import CLI, set_config_read_mode, set_docstring_parse_options
from litgpt.chat.base import main as chat_fn
from litgpt.deploy.serve import run_server as serve_fn
from litgpt.eval.evaluate import convert_and_evaluate as evaluate_fn
from litgpt.finetune.adapter import setup as finetune_adapter_fn
from litgpt.finetune.adapter_v2 import setup as finetune_adapter_v2_fn
from litgpt.finetune.full import setup as finetune_full_fn
from litgpt.finetune.lora import setup as finetune_lora_fn
from litgpt.generate.adapter import main as generate_adapter_fn
from litgpt.generate.adapter_v2 import main as generate_adapter_v2_fn
from litgpt.generate.base import main as generate_base_fn
from litgpt.generate.full import main as generate_full_fn
from litgpt.generate.sequentially import main as generate_sequentially_fn
from litgpt.generate.speculative_decoding import main as generate_speculatively_fn
from litgpt.generate.tp import main as generate_tp_fn
from litgpt.parser_config import parser_commands
from litgpt.pretrain import setup as pretrain_fn
from litgpt.scripts.convert_hf_checkpoint import convert_hf_checkpoint as convert_hf_checkpoint_fn
from litgpt.scripts.convert_lit_checkpoint import convert_lit_checkpoint as convert_lit_checkpoint_fn
from litgpt.scripts.convert_pretrained_checkpoint import (
convert_pretrained_checkpoint as convert_pretrained_checkpoint_fn,
)
from litgpt.scripts.download import download_from_hub as download_fn
from litgpt.scripts.merge_lora import merge_lora as merge_lora_fn
from litgpt.scripts.validate import validate_setup as validate_fn
PARSER_DATA = {
"download": download_fn,
"chat": chat_fn,
"finetune": finetune_lora_fn,
"finetune_lora": finetune_lora_fn,
"finetune_full": finetune_full_fn,
"finetune_adapter": finetune_adapter_fn,
"finetune_adapter_v2": finetune_adapter_v2_fn,
"pretrain": pretrain_fn,
"generate": generate_base_fn,
"generate_full": generate_full_fn,
"generate_adapter": generate_adapter_fn,
"generate_adapter_v2": generate_adapter_v2_fn,
"generate_sequentially": generate_sequentially_fn,
"generate_speculatively": generate_speculatively_fn,
"generate_tp": generate_tp_fn,
"convert_to_litgpt": convert_hf_checkpoint_fn,
"convert_from_litgpt": convert_lit_checkpoint_fn,
"convert_pretrained_checkpoint": convert_pretrained_checkpoint_fn,
"merge_lora": merge_lora_fn,
"evaluate": evaluate_fn,
"serve": serve_fn,
"validate": validate_fn,
}
def _check_commands():
assert set(parser_commands()) == set(PARSER_DATA.keys()), (
"PARSER_DATA has to be kept in sync with litgpt.parser_config.parser_commands()"
)
def main() -> None:
_check_commands()
set_docstring_parse_options(attribute_docstrings=True)
set_config_read_mode(urls_enabled=True)
# PyTorch bug that raises a false-positive warning
# More info: https://github.com/Lightning-AI/litgpt/issues/1561
warning_message = r"The epoch parameter in `scheduler.step\(\)` was not necessary and is being deprecated.*"
warnings.filterwarnings(
action="ignore", message=warning_message, category=UserWarning, module=r".*torch\.optim\.lr_scheduler.*"
)
torch.set_float32_matmul_precision("high")
CLI(PARSER_DATA)
if __name__ == "__main__":
main()
litgpt/adapter.py
+129
View File
@@ -0,0 +1,129 @@
# Copyright Lightning AI. Licensed under the Apache License 2.0, see LICENSE file.
"""Implementation of the paper:
LLaMA-Adapter: Efficient Fine-tuning of Language Models with Zero-init Attention
https://arxiv.org/abs/2303.16199
Port for LitGPT
"""
from dataclasses import dataclass
from typing import Any
import torch
import torch.nn as nn
from typing_extensions import Self
from litgpt.config import Config as BaseConfig
from litgpt.model import GPT as BaseModel
from litgpt.model import Block as BaseBlock
from litgpt.model import CausalSelfAttention as BaseCausalSelfAttention
@dataclass
class Config(BaseConfig):
adapter_prompt_length: int = 10
adapter_start_layer: int = 2
class GPT(BaseModel):
# Copy & paste from :class:`model.GPT`. Note that :class:`Block` is new here.
def __init__(self, config: Config) -> None:
nn.Module.__init__(self)
assert config.padded_vocab_size is not None
self.config = config
self.lm_head = nn.Linear(config.n_embd, config.padded_vocab_size, bias=config.lm_head_bias)
self.transformer = nn.ModuleDict(
dict(
wte=nn.Embedding(config.padded_vocab_size, config.n_embd),
h=nn.ModuleList(Block(config, block_idx) for block_idx in range(config.n_layer)),
ln_f=config.norm_class(config.n_embd, eps=config.norm_eps),
)
)
self.mask_cache: torch.Tensor | None = None
self.max_seq_length = self.config.block_size
@classmethod
def from_name(cls, name: str, **kwargs: Any) -> Self:
return cls(Config.from_name(name, **kwargs))
def _init_weights(self, module: nn.Module) -> None:
"""Meant to be used with `gpt.apply(gpt._init_weights)`. Unused method left for completeness."""
super()._init_weights(module)
if isinstance(module, CausalSelfAttention):
module.reset_parameters()
class Block(BaseBlock):
def __init__(self, config: Config, block_idx: int) -> None:
super().__init__(config, block_idx)
self.attn = CausalSelfAttention(config, block_idx)
class CausalSelfAttention(BaseCausalSelfAttention):
"""A modification of `litgpt.model.CausalSelfAttention` that adds the attention
over the adaption prompt."""
def __init__(self, config: Config, block_idx: int) -> None:
super().__init__(config, block_idx)
if block_idx >= config.adapter_start_layer:
# adapter embedding layer
self.adapter_wte = nn.Embedding(config.adapter_prompt_length, config.n_embd)
# gate for adaption
self.gating_factor = torch.nn.Parameter(torch.zeros(1, 1, config.n_head, 1))
# kv cache for inference
self.adapter_kv_cache: tuple[torch.Tensor, torch.Tensor] | None = None
def scaled_dot_product_attention(
self, q: torch.Tensor, k: torch.Tensor, v: torch.Tensor, mask: torch.Tensor | None = None
) -> torch.Tensor:
y = super().scaled_dot_product_attention(q, k, v, mask)
if self.block_idx < self.config.adapter_start_layer:
return y
aT = self.config.adapter_prompt_length
if self.adapter_kv_cache is not None:
# since this uses the wte weights as the prefix and the kv cache is only used during inference, ak and av
# are the same every call
ak, av = self.adapter_kv_cache
else:
prefix = self.adapter_wte.weight.reshape(1, aT, self.config.n_embd)
aqkv = self.qkv(prefix)
q_per_kv = self.config.n_head // self.config.n_query_groups
aqkv = aqkv.view(1, aT, self.config.n_query_groups, q_per_kv + 2, self.config.head_size)
aqkv = aqkv.permute(0, 2, 3, 1, 4)
_, ak, av = aqkv.split((q_per_kv, 1, 1), dim=2)
if self.config.n_query_groups != 1:
# for MHA this is a no-op
ak = ak.repeat_interleave(q_per_kv, dim=2)
av = av.repeat_interleave(q_per_kv, dim=2)
ak = ak.view(1, -1, aT, self.config.head_size) # (1, nh_ak, aT, hs)
av = av.view(1, -1, aT, self.config.head_size) # (1, nh_av, aT, hs)
self.adapter_kv_cache = (ak, av)
T = q.size(2)
amask = torch.ones(T, aT, dtype=torch.bool, device=q.device)
ay = super().scaled_dot_product_attention(q, ak, av, amask)
return y + self.gating_factor * ay
def reset_parameters(self) -> None:
if hasattr(self, "gating_factor"):
torch.nn.init.zeros_(self.gating_factor)
def _load_from_state_dict(self, state_dict: dict, prefix: str, *args: Any, **kwargs: Any) -> None:
"""For compatibility with older checkpoints."""
if (key := prefix + "gating_factor") in state_dict and state_dict[key].size(1) == self.config.n_head:
state_dict[key] = state_dict[key].permute(0, 2, 1, 3)
super()._load_from_state_dict(state_dict, prefix, *args, **kwargs)
def mark_only_adapter_as_trainable(model: GPT) -> None:
"""Sets `requires_grad=False` for all non-adapter weights."""
for name, param in model.named_parameters():
param.requires_grad = adapter_filter(name, param)
def adapter_filter(key: str, value: Any) -> bool:
return "adapter_wte" in key or "gating_factor" in key
litgpt/adapter_v2.py
+210
View File
@@ -0,0 +1,210 @@
# Copyright Lightning AI. Licensed under the Apache License 2.0, see LICENSE file.
"""Implementation of the paper:
LLaMA-Adapter V2: Parameter-Efficient Visual Instruction Model
https://arxiv.org/abs/2304.15010
Port for LitGPT
"""
from dataclasses import dataclass
from typing import Any
import torch
import torch.nn as nn
from typing_extensions import Self
import litgpt
from litgpt.adapter import GPT as BaseModel
from litgpt.adapter import CausalSelfAttention as BaseCausalSelfAttention
from litgpt.adapter import Config as BaseConfig
from litgpt.model import Block as BaseBlock
from litgpt.scripts.convert_hf_checkpoint import qkv_reassemble
from litgpt.utils import map_old_state_dict_weights
@dataclass
class Config(BaseConfig):
@property
def mlp_class(self) -> type:
return getattr(litgpt.adapter_v2, self.mlp_class_name)
def adapter_filter(key: str, value: Any) -> bool:
adapter_substrings = (
# regular adapter v1 parameters
"adapter_wte",
"gating_factor",
# adapter v2: new bias and scale used in Linear
"adapter_scale",
"adapter_bias",
# adapter v2: Norm parameters are now trainable
"norm_1",
"norm_2",
"ln_f",
)
return any(s in key for s in adapter_substrings)
class AdapterV2Linear(torch.nn.Module):
def __init__(self, in_features: int, out_features: int, **kwargs) -> None:
super().__init__()
self.linear = torch.nn.Linear(in_features, out_features, **kwargs)
self.adapter_bias = torch.nn.Parameter(torch.zeros(out_features), requires_grad=False)
self.adapter_scale = torch.nn.Parameter(torch.ones(out_features), requires_grad=False)
def forward(self, x: torch.Tensor) -> torch.Tensor:
return self.adapter_scale * (self.linear(x) + self.adapter_bias)
def reset_parameters(self) -> None:
nn.init.zeros_(self.adapter_bias)
nn.init.ones_(self.adapter_scale)
class GPT(BaseModel):
# Copy & paste from :class:`model.GPT`. Note that :class:`Block` is new here.
def __init__(self, config: Config) -> None:
nn.Module.__init__(self)
assert config.padded_vocab_size is not None
self.config = config
self.lm_head = AdapterV2Linear(config.n_embd, config.padded_vocab_size, bias=config.lm_head_bias)
self.transformer = nn.ModuleDict(
dict(
wte=nn.Embedding(config.padded_vocab_size, config.n_embd),
h=nn.ModuleList(Block(config, block_idx) for block_idx in range(config.n_layer)),
ln_f=config.norm_class(config.n_embd, eps=config.norm_eps),
)
)
self.mask_cache: torch.Tensor | None = None
self.max_seq_length = self.config.block_size
@classmethod
def from_name(cls, name: str, **kwargs: Any) -> Self:
return cls(Config.from_name(name, **kwargs))
def _init_weights(self, module: nn.Module) -> None:
"""Meant to be used with `gpt.apply(gpt._init_weights)`. Unused method left for completeness."""
super()._init_weights(module)
if isinstance(module, AdapterV2Linear):
module.reset_parameters()
def _load_from_state_dict(self, state_dict: dict, prefix: str, *args: Any, **kwargs: Any) -> None:
"""For compatibility with base checkpoints."""
mapping = {"lm_head.weight": "lm_head.linear.weight", "lm_head.bias": "lm_head.linear.bias"}
state_dict = map_old_state_dict_weights(state_dict, mapping, prefix)
super()._load_from_state_dict(state_dict, prefix, *args, **kwargs)
class Block(BaseBlock):
def __init__(self, config: Config, block_idx: int) -> None:
super().__init__(config, block_idx)
self.attn = CausalSelfAttention(config, block_idx)
self.mlp = config.mlp_class(config)
class CausalSelfAttention(BaseCausalSelfAttention):
"""A modification of `litgpt.adapter.CausalSelfAttention` that uses the Adapter V2 Linear class"""
# Copy&paste from :class:`model.CausalSelfAttention`
def __init__(self, config: Config, block_idx: int) -> None:
super().__init__(config, block_idx)
# key, query, value projections for all heads, but in a batch
shape = (config.n_head + 2 * config.n_query_groups) * config.head_size
self.qkv = AdapterV2Linear(in_features=config.n_embd, out_features=shape, bias=config.bias or config.attn_bias)
# output projection
self.proj = AdapterV2Linear(config.head_size * config.n_head, config.n_embd, bias=config.bias)
def _load_from_state_dict(self, state_dict: dict, prefix: str, *args: Any, **kwargs: Any) -> None:
"""For compatibility with base and/or legacy checkpoints."""
mapping = {
"qkv.weight": "qkv.linear.weight",
"qkv.bias": "qkv.linear.bias",
"proj.weight": "proj.linear.weight",
"proj.bias": "proj.linear.bias",
}
state_dict = map_old_state_dict_weights(state_dict, mapping, prefix)
# For compatibility with older checkpoints
if (key := prefix + "gating_factor") in state_dict and state_dict[key].size(1) == self.config.n_head:
state_dict[key] = state_dict[key].permute(0, 2, 1, 3)
for attr in ("weight", "bias"):
legacy_key = f"{prefix}attn.linear.{attr}"
current_key = f"{prefix}qkv.linear.{attr}"
if legacy_key in state_dict:
state_dict[current_key] = qkv_reassemble(state_dict.pop(legacy_key), self.config)
super()._load_from_state_dict(state_dict, prefix, *args, **kwargs)
class GptNeoxMLP(litgpt.model.GptNeoxMLP):
def __init__(self, config: Config) -> None:
nn.Module.__init__(self)
self.fc = AdapterV2Linear(config.n_embd, config.intermediate_size, bias=config.bias)
self.proj = AdapterV2Linear(config.intermediate_size, config.n_embd, bias=config.bias)
self.config = config
def _load_from_state_dict(self, state_dict: dict, prefix: str, *args: Any, **kwargs: Any) -> None:
"""For compatibility with base checkpoints."""
mapping = {
"fc.weight": "fc.linear.weight",
"fc.bias": "fc.linear.bias",
"proj.weight": "proj.linear.weight",
"proj.bias": "proj.linear.bias",
}
state_dict = map_old_state_dict_weights(state_dict, mapping, prefix)
super()._load_from_state_dict(state_dict, prefix, *args, **kwargs)
class LLaMAMLP(litgpt.model.LLaMAMLP):
def __init__(self, config: Config, intermediate_size: int | None = None) -> None:
nn.Module.__init__(self)
self.intermediate_size = intermediate_size or config.intermediate_size
self.fc_1 = AdapterV2Linear(config.n_embd, self.intermediate_size, bias=config.bias)
self.fc_2 = AdapterV2Linear(config.n_embd, self.intermediate_size, bias=config.bias)
self.proj = AdapterV2Linear(self.intermediate_size, config.n_embd, bias=config.bias)
self.config = config
def _load_from_state_dict(self, state_dict: dict, prefix: str, *args: Any, **kwargs: Any) -> None:
"""For compatibility with base checkpoints."""
mapping = {
"fc_1.weight": "fc_1.linear.weight",
"fc_1.bias": "fc_1.linear.bias",
"fc_2.weight": "fc_2.linear.weight",
"fc_2.bias": "fc_2.linear.bias",
"proj.weight": "proj.linear.weight",
"proj.bias": "proj.linear.bias",
}
state_dict = map_old_state_dict_weights(state_dict, mapping, prefix)
super()._load_from_state_dict(state_dict, prefix, *args, **kwargs)
class GemmaMLP(LLaMAMLP):
def forward(self, x: torch.Tensor) -> torch.Tensor:
x_fc_1 = self.fc_1(x)
x_fc_2 = self.fc_2(x)
x = torch.nn.functional.gelu(x_fc_1, approximate=self.config.gelu_approximate) * x_fc_2
return self.proj(x)
class LLaMAMoE(litgpt.model.LLaMAMoE):
def __init__(self, config: Config) -> None:
nn.Module.__init__(self)
self.gate = AdapterV2Linear(config.n_embd, config.n_expert, bias=False)
self.experts = nn.ModuleList(
LLaMAMLP(config, intermediate_size=config.moe_intermediate_size) for _ in range(config.n_expert)
)
self.config = config
def _load_from_state_dict(self, state_dict: dict, prefix: str, *args: Any, **kwargs: Any) -> None:
"""For compatibility with base checkpoints."""
mapping = {"gate.weight": "gate.linear.weight"}
state_dict = map_old_state_dict_weights(state_dict, mapping, prefix)
super()._load_from_state_dict(state_dict, prefix, *args, **kwargs)
def mark_only_adapter_v2_as_trainable(model: GPT) -> None:
"""Sets requires_grad=False for all non-adapter weights"""
for name, param in model.named_parameters():
param.requires_grad = adapter_filter(name, param)
litgpt/api.py
+735
View File
@@ -0,0 +1,735 @@
# Copyright Lightning AI. Licensed under the Apache License 2.0, see LICENSE file.
#
# This file implements the LitGPT Python API
import sys
import time
from collections.abc import Callable
from pathlib import Path
from typing import Any, Literal
import lightning as L
import numpy as np
import torch
from lightning.fabric.accelerators import CUDAAccelerator
from lightning.fabric.plugins import BitsandbytesPrecision
from tqdm import tqdm
from litgpt.chat.base import generate as stream_generate_fn
from litgpt.config import Config, name_to_config
from litgpt.generate.base import generate as generate_fn
from litgpt.generate.sequentially import sequential
from litgpt.generate.tp import tensor_parallel
from litgpt.model import GPT
from litgpt.prompts import PromptStyle, has_prompt_style, load_prompt_style, save_prompt_style
from litgpt.tokenizer import Tokenizer
from litgpt.utils import (
auto_download_checkpoint,
check_file_size_on_cpu_and_warn,
check_nvlink_connectivity,
chunked_cross_entropy,
copy_config_files,
extend_checkpoint_dir,
get_default_supported_precision,
load_checkpoint,
save_config,
)
class LLM(torch.nn.Module):
def __init__(
self,
model: GPT,
preprocessor=None,
prompt_style: PromptStyle = None,
devices: int | list[int] = None,
config: Config = None,
checkpoint_dir: Path = None,
fabric: L.Fabric = None,
generate_strategy: Literal["sequential", "tensor_parallel"] | None = None,
kv_cache_initialized: bool = False,
fixed_kv_cache_size: int | Literal["max_model_supported"] | None = None,
) -> None:
super().__init__()
self.model = model
self.preprocessor = preprocessor
self.devices = devices
self.prompt_style = prompt_style
self.config = config
self.checkpoint_dir = checkpoint_dir
self.fabric = fabric
self.generate_strategy = generate_strategy
self.kv_cache_initialized = kv_cache_initialized
self.fixed_kv_cache_size = fixed_kv_cache_size
self.prev_generated_seq_length = 0
"""
LLM model class for inference, pretraining, and finetuning.
Example:
from litgpt.api import LLM
llm = LLM.load("microsoft/phi-2")
text = llm.generate("What do Llamas eat?", top_k=1)
print(text)
"""
@property
def tokenizer(self):
return self.preprocessor.tokenizer
def state_dict(self, destination=None, prefix="", keep_vars=False):
return self.model.state_dict(destination=destination, prefix=prefix, keep_vars=keep_vars)
def load_state_dict(self, state_dict, strict=True):
return self.model.load_state_dict(state_dict, strict=strict)
def forward(
self,
input_ids: torch.Tensor,
target_ids: torch.Tensor | None = None,
loss_fn: Callable[[torch.Tensor, torch.Tensor], torch.Tensor] | None = None,
) -> torch.Tensor | tuple[torch.Tensor, torch.Tensor]:
logits = self.model(input_ids)
if target_ids is not None:
if loss_fn is None:
loss_fn = chunked_cross_entropy
loss = loss_fn(logits[..., :-1, :], target_ids[..., 1:])
return logits, loss
else:
return logits
def trainer_setup(self, trainer_ckpt: Path | None = None) -> None:
"""Initializes the model checkpoint for PyTorch Lightning Trainer contexts"""
self.model = GPT(self.config)
if trainer_ckpt is not None:
# strip the object name key from the state_dict
state_dict = torch.load(trainer_ckpt, weights_only=True)["state_dict"]
first_key = next(iter(state_dict))
prefix = first_key.split(".")[0] + "."
keys_to_modify = [key for key in state_dict if key.startswith(prefix)]
for key in keys_to_modify:
new_key = key.replace(prefix, "", 1)
state_dict[new_key] = state_dict.pop(key)
self.load_state_dict(state_dict, strict=True)
elif self.checkpoint_dir is not None:
state_dict = torch.load(self.checkpoint_dir / "lit_model.pth", weights_only=False)
self.load_state_dict(state_dict, strict=False)
else:
raise ValueError(
"No checkpoint found. Either provide a valid path via `trainer_ckpt` "
"or ensure that `self.checkpoint_dir` points to a folder containing a `lit_model.pth` weight file."
)
def save(self, out_dir: Path | None = None, prompt_style: PromptStyle | None = None) -> None:
out_dir = Path(out_dir)
save_path = out_dir / "lit_model.pth"
save_path.parent.mkdir(parents=True, exist_ok=True)
if prompt_style is None:
prompt_style = PromptStyle.from_config(self.config)
if self.fabric is None:
torch.save(self.state_dict(), save_path)
else:
self.fabric.save(save_path, self.state_dict())
if self.fabric is None or self.fabric.global_rank == 0:
# If initialization a model with random weights, the checkpoint dir can be none
if self.checkpoint_dir is not None:
copy_config_files(Path(self.checkpoint_dir), save_path.parent)
else:
save_config(self.config, out_dir)
save_prompt_style(prompt_style, save_path.parent)
@classmethod
def load(
cls,
model: str,
init: Literal["pretrained", "random"] | None = "pretrained",
tokenizer_dir: Path | None = None,
access_token: str | None = None,
distribute: Literal["auto"] | None = "auto",
) -> "LLM":
"""
Loads the LLM from a local directory or model hub.
Arguments
model: A local path to a directory containing the model weights or a valid model name.
You can get a list of valid model names via the `litgpt download list` command line argument.
init: If "pretrained" (default), downloads the model from the HF Hub if a local model can't be found at the `model`
directory name; otherwise loads the model from the local directory.
If "random", initializes the `model` with random weights.
tokenizer_dir: An optional tokenizer directory if `model` is not a checkpoint directory, or if a user
wants to use a different tokenizer instead.
access_token: Optional API token to access models with restrictions when using `init="pretrained"`.
distribute: If "auto" (default), initializes the model on a single GPU if available and otherwise on the CPU.
To have more control over the model distribution strategy and utilize multiple GPUs, you can set
`llm = LLM.load(..., distribute=None)` and call `llm.distribute(...)` manually.
"""
allowed_init = {"pretrained", "random"}
if init == "pretrained":
checkpoint_dir = auto_download_checkpoint(
model_name=model, access_token=access_token, ignore_tokenizer_files=tokenizer_dir is not None
)
config = Config.from_file(checkpoint_dir / "model_config.yaml")
elif init == "random":
checkpoint_dir = None
try:
config = Config.from_name(model)
except ValueError:
print(f"Model name {model} is not supported.\n")
available_models = "\n".join(sorted(name_to_config))
print(f"Available values:\n{available_models}")
return
else:
raise ValueError(f"Invalid init option: {init}. Must be one of {allowed_init}")
torch.set_float32_matmul_precision("high")
if tokenizer_dir is not None:
tokenizer_dir = extend_checkpoint_dir(Path(tokenizer_dir))
tokenizer = Tokenizer(tokenizer_dir)
elif checkpoint_dir is not None:
tokenizer = Tokenizer(checkpoint_dir)
else:
raise ValueError("Provide a path to a tokenizer directory via the `tokenizer_dir` setting.")
if checkpoint_dir is not None:
prompt_style = (
load_prompt_style(checkpoint_dir)
if has_prompt_style(checkpoint_dir)
else PromptStyle.from_config(config)
)
else:
prompt_style = PromptStyle.from_config(config)
if distribute == "auto":
if torch.cuda.is_available():
accelerator = "cuda"
elif torch.backends.mps.is_available():
accelerator = "mps"
else:
accelerator = "cpu"
fabric = L.Fabric(
accelerator=accelerator,
devices=1,
precision=get_default_supported_precision(training=False),
)
with fabric.init_module(empty_init=False):
model = GPT(config)
model.eval()
preprocessor = Preprocessor(tokenizer, device=fabric.device)
if checkpoint_dir is not None:
checkpoint_path = checkpoint_dir / "lit_model.pth"
check_file_size_on_cpu_and_warn(checkpoint_path, fabric.device)
load_checkpoint(fabric, model, checkpoint_path)
model = fabric.setup_module(model)
else:
preprocessor = Preprocessor(tokenizer, device="cuda" if torch.cuda.is_available() else "cpu")
model = None
fabric = None
return cls(
model=model,
preprocessor=preprocessor,
prompt_style=prompt_style,
config=config,
checkpoint_dir=checkpoint_dir,
fabric=fabric,
generate_strategy=None,
kv_cache_initialized=False,
fixed_kv_cache_size=False,
)
def distribute(
self,
accelerator: Literal["cpu", "cuda", "auto"] = "auto",
devices: int | Literal["auto"] = "auto",
precision: Any | None = None,
quantize: Literal["bnb.nf4", "bnb.nf4-dq", "bnb.fp4", "bnb.fp4-dq", "bnb.int8"] | None = None,
generate_strategy: Literal["sequential", "tensor_parallel"] | None = None,
fixed_kv_cache_size: int | Literal["max_model_supported"] | None = None,
) -> None:
"""
Moves the model onto specified devices for single-GPU or multi-GPU inference
accelerator: Which device type to load the model on ("cpu", "gpu", "mps", "cuda", or "auto")
devices: The number of devices (1, 2, etc.) or "auto", which uses all available devices
quantize: Whether to quantize the model and using which method:
- bnb.nf4, bnb.nf4-dq, bnb.fp4, bnb.fp4-dq: 4-bit quantization from bitsandbytes
- bnb.int8: 8-bit quantization from bitsandbytes
for more details, see https://github.com/Lightning-AI/litgpt/blob/main/tutorials/quantize.md
precision: Indicates the Fabric precision setting to use.
For instance, "32-true", "16-mixed", "16-true", "bf16-mixed", "bf16-true".
For more details, see https://lightning.ai/docs/fabric/stable/api/fabric_args.html#precision
generate_strategy: Whether to use a sequential model generation strategy. The "sequential" settings allows running
models that wouldn't fit in a single card by partitioning the transformer blocks across
all devices and running them sequentially. Sequential generation may be slower but allows using larger models.
Note that sequential generation sets `fixed_kv_cache_size="max_model_supported"`. You can set it to a lower integer
value, `fixed_kv_cache_size=256` to reduce memory. The `fixed_kv_cache_size` value determines the maximum number
of tokens that can be returned via `llm.generate(...)`.
fixed_kv_cache_size: If set to an integer value or "max_model_supported" is set, the kv-cache won't be resized dynamically
during `llm.generate` calls. Use this setting if you plan to compile the model or use `generate_strategy="sequential`.
Note that the chosen `fixed_kv_cache_size` value determines the maximum number of tokens that can be returned in `llm.generate(...)`.
"""
if self.checkpoint_dir is None:
raise NotImplementedError(
"The LLM was initialized with init='random' but .distribute() "
"currently only supports pretrained weights."
)
allowed_accelerators = {"cpu", "gpu", "cuda", "mps", "auto"}
if accelerator not in allowed_accelerators:
raise ValueError(f"Invalid accelerator: {accelerator}. Must be one of {allowed_accelerators}.")
if accelerator == "auto":
if torch.cuda.is_available():
accelerator = "cuda"
elif torch.backends.mps.is_available():
accelerator = "mps"
else:
accelerator = "cpu"
if generate_strategy in ("sequential", "tensor_parallel") and accelerator not in ("cuda", "gpu"):
raise NotImplementedError(
f"generate_strategy='{generate_strategy}' is only supported for accelerator='cuda'|'gpu'."
)
if devices == "auto":
if generate_strategy in ("sequential", "tensor_parallel"):
total_devices = CUDAAccelerator.auto_device_count()
else:
total_devices = 1
elif isinstance(devices, int) and accelerator == "cuda":
use_devices = calculate_number_of_devices(devices)
total_devices = CUDAAccelerator.auto_device_count()
if use_devices > total_devices:
raise ValueError(
f"You selected more devices ({use_devices}) than available in your system ({total_devices})."
)
else:
total_devices = use_devices
if total_devices > 1 and generate_strategy not in ("sequential", "tensor_parallel"):
raise NotImplementedError(
"Support for multiple devices is currently only implemented for generate_strategy='sequential'|'tensor_parallel'."
)
elif accelerator == "cpu" or accelerator == "mps":
total_devices = 1
else:
raise ValueError(f"devices argument must be an integer or 'auto', got {devices}")
print(f"Using {total_devices} device(s)", file=sys.stderr)
if precision is None:
precision = get_default_supported_precision(training=False)
print("Precision set", file=sys.stderr)
plugins = None
if quantize is not None and quantize.startswith("bnb."):
if "mixed" in precision:
raise ValueError("The combination of quantization and mixed precision is not supported.")
dtype = {"16-true": torch.float16, "bf16-true": torch.bfloat16, "32-true": torch.float32}[precision]
plugins = BitsandbytesPrecision(quantize[4:], dtype)
precision = None
# set "ddp" as the strategy for the launching functionality, but there's no data-parallelism
if generate_strategy != "tensor_parallel":
fabric = L.Fabric(
accelerator=accelerator,
devices=1, # Otherwise sequential wouldn't work, see litgpt/generate/sequentially.py
# devices=devices,
precision=precision,
plugins=plugins,
)
else:
fabric = L.Fabric(
accelerator=accelerator, devices=total_devices, strategy="ddp", precision=precision, plugins=plugins
)
if torch.cuda.is_available() and fabric.accelerator.auto_device_count() > 1:
check_nvlink_connectivity(fabric)
fabric.launch()
print("Fabric launched", file=sys.stderr)
self.kv_cache_initialized = False
if generate_strategy is None:
with fabric.init_module(empty_init=(total_devices > 1)):
model = GPT(self.config)
model.eval()
if self.checkpoint_dir is not None:
load_checkpoint(fabric, model, self.checkpoint_dir / "lit_model.pth")
model = fabric.setup_module(model)
if fixed_kv_cache_size is not None:
if fixed_kv_cache_size is None or fixed_kv_cache_size == "max_model_supported":
kv_cache_size = model.max_seq_length
else:
kv_cache_size = fixed_kv_cache_size
model.set_kv_cache(batch_size=1, max_seq_length=kv_cache_size, device=fabric.device)
self.kv_cache_initialized = True
self.fixed_kv_cache_size = fixed_kv_cache_size
elif generate_strategy in ("sequential", "tensor_parallel"):
with fabric.init_tensor(), torch.device("meta"):
model = GPT(self.config)
model.eval()
if generate_strategy == "sequential":
state_dict = torch.load(
str(self.checkpoint_dir / "lit_model.pth"), mmap=True, map_location="cpu", weights_only=False
)
model.load_state_dict(state_dict, assign=True)
model = fabric.setup_module(model, move_to_device=False)
if fixed_kv_cache_size is None:
fixed_kv_cache_size = "max_model_supported"
if fixed_kv_cache_size == "max_model_supported":
kv_cache_size = model.max_seq_length
else:
kv_cache_size = fixed_kv_cache_size
model = sequential(model, fabric.device, kv_cache_size, total_devices)
self.fixed_kv_cache_size = fixed_kv_cache_size
elif generate_strategy == "tensor_parallel":
if fabric.global_rank == 0:
pbar = tqdm(total=fabric.world_size, desc="Loading model weights")
for rank in range(fabric.world_size):
if fabric.global_rank == rank:
state_dict = torch.load(
str(self.checkpoint_dir / "lit_model.pth"),
mmap=True,
map_location="cpu",
weights_only=False,
)
model.load_state_dict(state_dict, assign=True)
# cannot use `.setup_module` because it will wrap with DDP
model = fabric._precision.convert_module(model)
model = tensor_parallel(fabric, model)
with fabric.init_tensor():
if fixed_kv_cache_size is None:
fixed_kv_cache_size = "max_model_supported"
if fixed_kv_cache_size == "max_model_supported":
kv_cache_size = model.max_seq_length
else:
kv_cache_size = fixed_kv_cache_size
model.max_seq_length = kv_cache_size
# the rope cache which is on meta device
model.cos, model.sin = model.rope_cache()
# enable the kv cache
model.set_kv_cache(batch_size=1)
model.eval()
model = fabric.to_device(model)
fabric.barrier()
if fabric.global_rank == 0:
pbar.update(1)
if fabric.global_rank == 0:
pbar.close()
self.kv_cache_initialized = True
else:
raise ValueError(f"Unsupported generate_strategy: {generate_strategy}")
self.model = model
self.fabric = fabric
self.preprocessor.device = fabric.device
@torch.inference_mode()
def generate(
self,
prompt: str,
sys_prompt: str | None = None,
max_new_tokens: int = 50,
temperature: float = 1.0,
top_k: int | None = None,
top_p: float = 1.0,
return_as_token_ids: bool = False,
stream: bool = False,
) -> str | torch.Tensor:
"""
Takes a conditioning sequence (prompt) as input and continues to generate as many tokens as requested.
Arguments:
model: The model to use.
prompt: The prompt string to use for generating the samples.
sys_prompt: The system prompt string to use for generating the samples.
The system prompt allows the user to provide additional instructions to shape all responses by providing additional context, behavioral guidelines, style, and constraints.
max_new_tokens: The maximum number of new tokens to return.
temperature: Scales the predicted logits by 1 / temperature.
top_k: If specified, only sample among the tokens with the k highest probabilities.
top_p: If specified, it represents the cumulative probability threshold to consider in the sampling process.
In top-p sampling, the next token is sampled from the highest probability tokens
whose cumulative probability exceeds the threshold `top_p`. When specified,
it must be `0 <= top_p <= 1`. Here, `top_p=0` is equivalent
to sampling the most probable token, while `top_p=1` samples from the whole distribution.
It can be used in conjunction with `top_k` and `temperature` with the following order
of application:
1. `top_k` sampling
2. `temperature` scaling
3. `top_p` sampling
For more details, see https://arxiv.org/abs/1904.09751
or https://huyenchip.com/2024/01/16/sampling.html#top_p
return_as_token_ids: If True, returns the token IDs as a torch.Tensor. Otherwise, returns the decoded text as a string.
stream: If True, returns a generator that yields tokens as they are generated.
At the moment, this setting is slower and may use more memory than the non-streaming version.
We plan to resolve this in the future.
"""
if self.model is None:
raise AttributeError(
"The model is not initialized yet; use the .distribute() "
"or .trainer_setup() method to initialize the model."
)
input_ids = self._text_to_token_ids(prompt, sys_prompt)
prompt_length = input_ids.size(0)
max_returned_tokens = prompt_length + max_new_tokens
if not self.kv_cache_initialized:
if self.fabric is not None:
device = self.fabric.device
else:
device = self.preprocessor.device
self.model.set_kv_cache(batch_size=1, max_seq_length=max_returned_tokens, device=device)
self.kv_cache_initialized = True
# Dynamically grow the kv cache size if necessary
if not self.fixed_kv_cache_size and self.prev_generated_seq_length < max_returned_tokens:
tmp_device = self.model.mask_cache.device
self.model.clear_kv_cache()
self.model.set_kv_cache(batch_size=1, max_seq_length=max_returned_tokens, device=tmp_device)
else:
for block in self.model.transformer.h:
block.attn.kv_cache.reset_parameters()
self.prev_generated_seq_length = max_returned_tokens
self.model.eval()
def iterator():
outputs = stream_generate_fn(
model=self.model,
prompt=input_ids,
max_returned_tokens=max_returned_tokens,
temperature=temperature,
top_k=top_k,
top_p=top_p,
stop_tokens=([self.preprocessor.tokenizer.eos_id],),
)
if return_as_token_ids:
yield from outputs
else:
for output in outputs:
yield self.preprocessor.decode(output)
return
if stream:
outputs = iterator()
else:
outputs = generate_fn(
model=self.model,
prompt=input_ids,
max_returned_tokens=max_returned_tokens,
temperature=temperature,
top_k=top_k,
top_p=top_p,
eos_id=self.preprocessor.tokenizer.eos_id,
include_prompt=False,
)
if stream:
return outputs
elif return_as_token_ids:
return outputs
else:
return self.preprocessor.decode(outputs)
def _text_to_token_ids(self, prompt: str, sys_prompt: str | None = None) -> torch.Tensor:
"""Utility method to convert a prompt text to token IDs"""
prompt = self.prompt_style.apply(prompt, sys_prompt=sys_prompt)
input_ids = self.preprocessor.encode(prompt)
return input_ids
def benchmark(self, num_iterations=1, **kwargs):
"""
A wrapper around the .generate() method to calculate runtime performance.
Arguments:
num_iterations: How often the `.generate()` call is repeated.
kwargs: Keyword arguments that are passed to the .generate() method.
"""
benchmark_dict = {}
for i in range(num_iterations):
time_to_first_token = None
t0 = time.perf_counter()
outputs = self.generate(**kwargs)
if kwargs.get("stream", False):
gen_outputs = []
for e in outputs:
if time_to_first_token is None:
t1 = time.perf_counter()
time_to_first_token = t1 - t0
gen_outputs.append(e)
outputs = "".join(gen_outputs)
else:
outputs = self.generate(
**kwargs,
)
benchmark_dict.setdefault("Seconds total", []).append(time.perf_counter() - t0)
benchmark_dict.setdefault("Seconds to first token", []).append(time_to_first_token)
tokens_generated = self.preprocessor.encode(outputs).size(0)
benchmark_dict.setdefault("Tokens generated", []).append(tokens_generated)
benchmark_dict.setdefault("Inference speed in tokens/sec", []).append(
benchmark_dict["Tokens generated"][-1] / benchmark_dict["Seconds total"][-1]
)
if self.fabric is not None and self.fabric.device.type == "cuda":
benchmark_dict.setdefault("Total GPU memory allocated in GB", []).append(
torch.cuda.max_memory_allocated() / 1e9
)
return outputs, benchmark_dict
class Preprocessor:
"""
Preprocessor class for tokenization and de-tokenization.
"""
def __init__(self, tokenizer: Tokenizer, device: str = "cpu") -> None:
self.tokenizer = tokenizer
self.device = device
def encode(self, text: str) -> torch.Tensor:
return self.tokenizer.encode(text, device=self.device)
def decode(self, token_ids: torch.Tensor) -> str:
return self.tokenizer.decode(token_ids)
def calculate_number_of_devices(devices):
"""
Utility function to calculate the number of devices.
"""
num_devices = devices if isinstance(devices, int) else len(devices) if isinstance(devices, list) else 0
return num_devices
def benchmark_dict_to_markdown_table(data):
"""
Converts .benchmark() outputs to a markdown table
"""
markdown_table = (
"| Metric | Mean | Std Dev |\n"
)
markdown_table += (
"|-------------------------------------|-----------------------------|-----------------------------|\n"
)
for key, values in data.items():
mean_value = np.mean(values)
std_dev_value = np.std(values, ddof=1)
formatted_mean = f"{mean_value:.2f}"
formatted_std_dev = f"{std_dev_value:.2f}"
markdown_table += f"| {key.ljust(35)} | {formatted_mean.ljust(27)} | {formatted_std_dev.ljust(27)} |\n"
return markdown_table
def pull_request_benchmark_util(model_name="microsoft/phi-2", num_iterations=6):
def print_table(header, data):
print(f"\n### {header}\n")
markdown_table = (
f"| Metric | First Iteration | "
f"Iter 2-{num_iterations} Mean | Iter 2-{num_iterations} Standard Dev. |\n"
f"|--------------------------------------|-----------------|"
f"-------------------|-------------------------|\n"
)
for key, value in data.items():
first_iteration = f"{value[0]:.2f}" if value[0] is not None else "N/A"
clean_values = [v for v in value[1:] if v is not None]
if clean_values:
mean_value = np.mean(clean_values)
std_dev_value = np.std(clean_values, ddof=1)
mean_str = f"{mean_value:.2f}"
std_dev_str = f"{std_dev_value:.2f}"
else:
mean_str = "N/A"
std_dev_str = "N/A"
markdown_table += f"| {key:<36} | {first_iteration:<15} | {mean_str:<17} | {std_dev_str:<23} |\n"
print(markdown_table)
import subprocess
try:
g_hash = subprocess.run(
["git", "rev-parse", "--short", "HEAD"], capture_output=True, text=True, check=True
).stdout.strip()
print(f"Git Commit Hash: {g_hash}")
except subprocess.CalledProcessError:
print("Git Commit Hash: N/A")
print(f"PyTorch version: {torch.__version__}")
device = torch.device("cuda" if torch.cuda.is_available() else "cpu")
print(f"Device: {device}\n")
# 1st table
llm = LLM.load(
model=model_name,
)
text, bench_d = llm.benchmark(num_iterations=num_iterations, prompt="What do llamas eat?", top_k=1)
print_table(f"Defaults ({model_name}), 1st time", bench_d)
del llm
# 2nd table
llm = LLM.load(
model=model_name,
)
text, bench_d = llm.benchmark(num_iterations=num_iterations, prompt="What do llamas eat?", top_k=1)
print_table(f"Defaults ({model_name}), 2nd time", bench_d)
del llm
# 3rd table
llm = LLM.load(
model=model_name,
)
text, bench_d = llm.benchmark(num_iterations=num_iterations, prompt="What do llamas eat?", top_k=1, stream=True)
print_table("stream=True", bench_d)
del llm
# 4th table
llm = LLM.load(model=model_name, distribute=None)
llm.distribute(fixed_kv_cache_size=500)
text, bench_d = llm.benchmark(num_iterations=num_iterations, prompt="What do llamas eat?", top_k=1, stream=True)
print_table("stream=True + fixed_kv_cache=500", bench_d)
litgpt/args.py
+118
View File
@@ -0,0 +1,118 @@
# Copyright Lightning AI. Licensed under the Apache License 2.0, see LICENSE file.
import math
import warnings
from dataclasses import dataclass
@dataclass
class TrainArgs:
"""Training-related arguments"""
save_interval: int | None = 1000
"""Number of optimizer steps between saving checkpoints"""
log_interval: int = 1
"""Number of iterations between logging calls"""
global_batch_size: int = 64
"""Number of samples between optimizer steps across data-parallel ranks"""
micro_batch_size: int = 4
"""Number of samples per data-parallel rank"""
lr_warmup_steps: int | None = 100
"""Number of iterations with learning rate warmup active"""
lr_warmup_fraction: float | None = None
"""The fraction of an epoch to use for learning rate warmup"""
epochs: int | None = None
"""Number of epochs to train on"""
# TODO: `pretrain` is the only script using `max_tokens` explicitly. replace it with epoch_size*epochs?
max_tokens: int | None = None
"""Total number of tokens to train on"""
max_steps: int | None = None
"""Limits the number of optimizer steps to run"""
max_time: float | None = None
"""Limits the number of seconds to train for"""
max_seq_length: int | None = None
"""Limits the length of samples"""
tie_embeddings: bool | None = None
"""Whether to tie the embedding weights with the language modeling head weights"""
# Optimization args
max_norm: float | None = None
min_lr: float = 6e-5
def __post_init__(self) -> None:
if self.lr_warmup_fraction and self.lr_warmup_steps:
raise ValueError(
"Can't provide both `--train.lr_warmup_fraction` and `--train.lr_warmup_steps`. Choose one."
)
if self.lr_warmup_fraction and not (0 <= self.lr_warmup_fraction <= 1):
raise ValueError("`--train.lr_warmup_fraction` must be between 0 and 1.")
if self.lr_warmup_steps and self.max_steps and (self.lr_warmup_steps >= self.max_steps):
warnings.warn(
"`--train.lr_warmup_steps` should be less than `--train.max_steps`."
f" Got {self.lr_warmup_steps} lr_warmup_steps and {self.max_steps} max_steps.",
UserWarning,
)
def gradient_accumulation_iters(self, devices: int, num_nodes: int = 1) -> int:
"""Number of iterations between gradient synchronizations"""
gradient_accumulation_iters = self.batch_size(devices, num_nodes) // self.micro_batch_size
assert gradient_accumulation_iters > 0
return gradient_accumulation_iters
def batch_size(self, devices: int, num_nodes: int = 1) -> int:
"""Number of samples between optimizer steps per data-parallel rank"""
batch_size = self.global_batch_size // (devices * num_nodes)
assert batch_size > 0
return batch_size
def warmup_iters(self, devices: int, num_nodes: int, max_iters: int, train_dataloader) -> int:
"""Number of iterations to warm up the learning rate."""
if self.lr_warmup_fraction:
return min(max_iters, math.ceil(self.lr_warmup_fraction * len(train_dataloader)))
if self.lr_warmup_steps:
return min(max_iters, self.lr_warmup_steps * self.gradient_accumulation_iters(devices, num_nodes))
return 0
@dataclass
class EvalArgs:
"""Evaluation-related arguments"""
interval: int = 600
"""Number of optimizer steps between evaluation calls"""
max_new_tokens: int | None = None
"""Number of tokens to generate"""
max_iters: int = 100
"""Number of iterations"""
initial_validation: bool = False
"""Whether to evaluate on the validation set at the beginning of the training"""
final_validation: bool = True
"""Whether to evaluate on the validation set at the end of the training"""
evaluate_example: str | int = "first"
"""How to pick an example instruction to evaluate periodically during training.
Can be "first", "random", or an integer index to pick a specific example."""
@dataclass
class LogArgs:
"""Logging-related arguments. Different loggers use different fields."""
# === WandB Fields ===
project: str | None = None
"""WandB project name"""
run: str | None = None
"""WandB run name (defaults to generated name)"""
group: str | None = None
"""WandB group name"""
# === LitLogger Fields (Lightning.ai) ===
teamspace: str | None = None
"""Teamspace name where charts and artifacts will appear"""
metadata: dict | None = None
"""Extra metadata to associate with the experiment as tags"""
log_model: bool = False
"""If True, automatically log model checkpoints as artifacts"""
save_logs: bool = True
"""If True, capture and upload terminal logs"""
checkpoint_name: str | None = None
"""Override the base name for logged checkpoints"""
litgpt/chat/__init__.py
View File
litgpt/chat/base.py
+273
View File
@@ -0,0 +1,273 @@
# Copyright Lightning AI. Licensed under the Apache License 2.0, see LICENSE file.
import sys
import time
from collections.abc import Iterator
from pathlib import Path
from pprint import pprint
from typing import Literal
import lightning as L
import torch
from lightning.fabric.plugins import BitsandbytesPrecision
from litgpt.config import Config
from litgpt.model import GPT
from litgpt.prompts import PromptStyle, has_prompt_style, load_prompt_style
from litgpt.scripts.merge_lora import merge_lora
from litgpt.tokenizer import Tokenizer
from litgpt.utils import (
auto_download_checkpoint,
check_file_size_on_cpu_and_warn,
extend_checkpoint_dir,
get_default_supported_precision,
load_checkpoint,
)
@torch.inference_mode()
def generate(
model: GPT,
prompt: torch.Tensor,
max_returned_tokens: int,
*,
temperature: float = 1.0,
top_k: int | None = None,
top_p: float = 1.0,
stop_tokens: tuple[list[int], ...] = (),
) -> Iterator[torch.Tensor]:
"""Takes a conditioning sequence (prompt) as input and continues to generate as many tokens as possible.
Arguments:
model: The model to use.
prompt: Tensor of shape (T) with indices of the prompt sequence.
max_returned_tokens: The maximum number of tokens to return (given plus generated).
temperature: Scales the predicted logits by 1 / temperature
top_k: If specified, only sample among the tokens with the k highest probabilities.
top_p: If specified, it represents the cumulative probability threshold to consider in the sampling process.
In top-p sampling, the next token is sampled from the highest probability tokens
whose cumulative probability exceeds the threshold `top_p`. When specified,
it must be `0 <= top_p <= 1`. Here, `top_p=0` is equivalent
to sampling the most probable token, while `top_p=1` samples from the whole distribution.
It can be used in conjunction with `top_k` and `temperature` with the following order
of application:
1. `top_k` sampling
2. `temperature` scaling
3. `top_p` sampling
For more details, see https://arxiv.org/abs/1904.09751
or https://huyenchip.com/2024/01/16/sampling.html#top_p
stop_tokens: If specified, stop generating any more token once one of this list is generated.
"""
from litgpt.generate.base import generate_fn
return generate_fn(
include_prompt=False,
include_eos=False,
model=model,
prompt=prompt,
max_returned_tokens=max_returned_tokens,
temperature=temperature,
top_k=top_k,
top_p=top_p,
stop_tokens=stop_tokens,
)
def process_prompt(
prompt, model, tokenizer, prompt_style, fabric, temperature, max_new_tokens, top_k, top_p, stop_tokens
):
prompt = prompt_style.apply(prompt=prompt)
encoded_prompt = tokenizer.encode(prompt, device=fabric.device)
if max_new_tokens is None:
max_returned_tokens = model.max_seq_length
else:
first_turn = model.mask_cache is None
max_returned_tokens = encoded_prompt.size(0) + max_new_tokens
if first_turn or max_returned_tokens > model.max_seq_length:
model.max_seq_length = max_returned_tokens
model.set_kv_cache(batch_size=1, device=fabric.device)
y: Iterator[torch.Tensor] = generate(
model,
encoded_prompt,
max_returned_tokens,
temperature=temperature,
top_k=top_k,
top_p=top_p,
stop_tokens=stop_tokens,
)
token_generator: Iterator[str] = tokenizer.decode_stream(y, device=fabric.device)
fabric.print(">> Reply: ", end="")
t0 = time.perf_counter()
tokens_generated = 0
for tok in token_generator:
tokens_generated += 1
fabric.print(tok, end="", flush=True)
t = time.perf_counter() - t0
for block in model.transformer.h:
block.attn.kv_cache.reset_parameters()
fabric.print(
f"\nTime for inference: {t:.02f} sec total, {tokens_generated / t:.02f} tokens/sec, {tokens_generated} tokens",
file=sys.stderr,
)
fabric.print()
def interact(multiline, model, tokenizer, prompt_style, fabric, temperature, max_new_tokens, top_k, top_p, stop_tokens):
while True:
try:
if not multiline:
prompt = input(">> Prompt: ")
else:
print(">> Prompt: (Type '!submit' on a new line to end input).")
prompt_lines = []
while True:
line = input()
if line.strip().lower() in ("!submit", "!quit", "!exit"):
break
prompt_lines.append(line)
prompt = "\n".join(prompt_lines)
except KeyboardInterrupt:
break
prompt = prompt.strip()
if not prompt or prompt.lower() in ("!quit", "!exit"):
break
process_prompt(
prompt, model, tokenizer, prompt_style, fabric, temperature, max_new_tokens, top_k, top_p, stop_tokens
)
@torch.inference_mode()
def main(
checkpoint_dir: Path,
*,
max_new_tokens: int = 50,
top_k: int | None = 50,
top_p: float = 1.0,
temperature: float = 0.8,
quantize: Literal["bnb.nf4", "bnb.nf4-dq", "bnb.fp4", "bnb.fp4-dq", "bnb.int8"] | None = None,
precision: str | None = None,
compile: bool = False,
multiline: bool = False,
access_token: str | None = None,
) -> None:
"""Chat with a model.
Args:
checkpoint_dir: A local path to a directory containing the model weights or a valid model name.
You can get a list of valid model names via the `litgpt download list` command line argument.
max_new_tokens: The number of generation steps to take.
top_k: The number of top most probable tokens to consider in the sampling process.
top_p: If specified, it represents the cumulative probability threshold to consider in the sampling process.
In top-p sampling, the next token is sampled from the highest probability tokens
whose cumulative probability exceeds the threshold `top_p`. When specified,
it must be `0 <= top_p <= 1`. Here, `top_p=0` is equivalent
to sampling the most probable token, while `top_p=1` samples from the whole distribution.
It can be used in conjunction with `top_k` and `temperature` with the following order
of application:
1. `top_k` sampling
2. `temperature` scaling
3. `top_p` sampling
For more details, see https://arxiv.org/abs/1904.09751
or https://huyenchip.com/2024/01/16/sampling.html#top_p
temperature: A value controlling the randomness of the sampling process. Higher values result in more random
samples.
quantize: Whether to quantize the model and using which method:
- bnb.nf4, bnb.nf4-dq, bnb.fp4, bnb.fp4-dq: 4-bit quantization from bitsandbytes
- bnb.int8: 8-bit quantization from bitsandbytes
for more details, see https://github.com/Lightning-AI/litgpt/blob/main/tutorials/quantize.md
precision: Indicates the Fabric precision setting to use.
compile: Whether to use compilation to speed up token generation. Will increase startup time.
multiline: Whether to support multiline input prompts.
access_token: Optional API token to access models with restrictions.
"""
checkpoint_dir = extend_checkpoint_dir(checkpoint_dir)
pprint(locals())
precision = precision or get_default_supported_precision(training=False)
plugins = None
if quantize is not None and quantize.startswith("bnb."):
if "mixed" in precision:
raise ValueError("Quantization and mixed precision is not supported.")
dtype = {"16-true": torch.float16, "bf16-true": torch.bfloat16, "32-true": torch.float32}[precision]
plugins = BitsandbytesPrecision(quantize[4:], dtype)
precision = None
fabric = L.Fabric(devices=1, precision=precision, plugins=plugins)
# Merge if this is a raw LoRA checkpoint
checkpoint_path = checkpoint_dir / "lit_model.pth"
if (checkpoint_dir / "lit_model.pth.lora").is_file() and not checkpoint_path.is_file():
print("Merging LoRA weights with the base model. This won't take long and is a one-time-only thing.")
merge_lora(checkpoint_dir)
if not checkpoint_path.is_file():
checkpoint_dir = auto_download_checkpoint(model_name=checkpoint_dir, access_token=access_token)
checkpoint_path = checkpoint_dir / "lit_model.pth"
check_file_size_on_cpu_and_warn(checkpoint_path, fabric.device)
config = Config.from_file(checkpoint_dir / "model_config.yaml")
with fabric.init_module(empty_init=True):
model = GPT(config)
if compile:
print(
"IMPORTANT: with enabled compilation the KV-cache size is determined by model's maximum context size, which leads to "
"a higher memory consumption. In case of an OOM error, try to set `--compile=False`."
)
model.set_kv_cache(batch_size=1)
load_checkpoint(fabric, model, checkpoint_path)
model.eval()
if compile:
torch._dynamo.config.automatic_dynamic_shapes = True
torch._inductor.config.triton.unique_kernel_names = True
torch._inductor.config.coordinate_descent_tuning = True
global next_token
next_token = torch.compile(next_token, mode="reduce-overhead", dynamic=True)
model = fabric.setup_module(model)
tokenizer = Tokenizer(checkpoint_dir)
prompt_style = (
load_prompt_style(checkpoint_dir) if has_prompt_style(checkpoint_dir) else PromptStyle.from_config(config)
)
stop_tokens = prompt_style.stop_tokens(tokenizer)
if multiline:
exit_instruction = "To exit, enter '!quit' or '!exit' on an empty prompt and press 'Enter'."
else:
exit_instruction = "To exit, press 'Enter' on an empty prompt."
print(f"Now chatting with {config.name}.\n{exit_instruction}\n")
L.seed_everything(1234)
interact(
multiline=multiline,
model=model,
tokenizer=tokenizer,
prompt_style=prompt_style,
fabric=fabric,
temperature=temperature,
max_new_tokens=(None if compile else max_new_tokens),
top_k=top_k,
top_p=top_p,
stop_tokens=stop_tokens,
)
if fabric.device.type == "cuda":
fabric.print(f"\nMemory used: {torch.cuda.max_memory_allocated() / 1e9:.02f} GB", file=sys.stderr)
litgpt/config.py
+3194
View File
File diff suppressed because it is too large Load Diff
litgpt/constants.py
+30
View File
@@ -0,0 +1,30 @@
# Copyright Lightning AI. Licensed under the Apache License 2.0, see LICENSE file.
"""Centralized package availability constants for optional dependencies."""
from lightning_utilities.core.imports import RequirementCache
# Logger-related constants
_SUPPORTED_LOGGERS: tuple[str, ...] = ("csv", "tensorboard", "wandb", "mlflow", "litlogger")
# Logger-related optional dependencies
_LITLOGGER_AVAILABLE = RequirementCache("litlogger>=0.1.7")
_TENSORBOARD_AVAILABLE = RequirementCache("tensorboard")
_WANDB_AVAILABLE = RequirementCache("wandb")
_MLFLOW_AVAILABLE = RequirementCache("mlflow")
_MLFLOW_SKINNY_AVAILABLE = RequirementCache("mlflow-skinny")
# PyTorch version-specific constants
_TORCH_EQUAL_2_7 = RequirementCache("torch>=2.7.0,<2.8")
_TORCH_EQUAL_2_8 = RequirementCache("torch>=2.8.0,<2.9")
# Other optional dependencies
_REQUESTS_AVAILABLE = RequirementCache("requests")
_THUNDER_AVAILABLE = RequirementCache("thunder")
_TRITON_AVAILABLE = RequirementCache("triton")
_BITANDBYTES_AVAILABLE = RequirementCache("bitsandbytes")
_BITANDBYTES_AVAILABLE_NOT_EQUAL_0_42_0 = RequirementCache("bitsandbytes != 0.42.0")
_LITDATA_AVAILABLE = RequirementCache("litdata")
_LITSERVE_AVAILABLE = RequirementCache("litserve")
_JINJA2_AVAILABLE = RequirementCache("jinja2")
_SAFETENSORS_AVAILABLE = RequirementCache("safetensors")
_HF_TRANSFER_AVAILABLE = RequirementCache("hf_transfer")

Some files were not shown because too many files have changed in this diff Show More