fix the patch to work properly and work with FSDP

enable eval dataloader using multipack again
update forwards so we only calculate cu_seqlens once
2023-08-17 20:33:56 -04:00 · 2023-08-17 18:17:31 -04:00 · 2023-08-17 18:17:31 -04:00 · 2023-08-17 18:17:30 -04:00 · 2023-08-17 18:17:30 -04:00 · 2023-08-17 18:17:30 -04:00
34 changed files with 1567 additions and 495 deletions
--- a/.github/CODE_OF_CONDUCT.md
+++ b/.github/CODE_OF_CONDUCT.md
@@ -0,0 +1,129 @@
+# Contributor Covenant Code of Conduct
+
+## Our Pledge
+
+We as members, contributors, and leaders pledge to make participation in our
+community a harassment-free experience for everyone, regardless of age, body
+size, visible or invisible disability, ethnicity, sex characteristics, gender
+identity and expression, level of experience, education, socio-economic status,
+nationality, personal appearance, race, religion, or sexual identity
+and orientation.
+
+We pledge to act and interact in ways that contribute to an open, welcoming,
+diverse, inclusive, and healthy community.
+
+## Our Standards
+
+Examples of behavior that contributes to a positive environment for our
+community include:
+
+* Demonstrating empathy and kindness toward other people
+* Being respectful of differing opinions, viewpoints, and experiences
+* Giving and gracefully accepting constructive feedback
+* Accepting responsibility and apologizing to those affected by our mistakes,
+  and learning from the experience
+* Focusing on what is best not just for us as individuals, but for the
+  overall community
+
+Examples of unacceptable behavior include:
+
+* The use of sexualized language or imagery, and sexual attention or
+  advances of any kind
+* Trolling, insulting or derogatory comments, and personal or political attacks
+* Public or private harassment
+* Publishing others' private information, such as a physical or email
+  address, without their explicit permission
+* Other conduct which could reasonably be considered inappropriate in a
+  professional setting
+
+## Enforcement Responsibilities
+
+Community leaders are responsible for clarifying and enforcing our standards of
+acceptable behavior and will take appropriate and fair corrective action in
+response to any behavior that they deem inappropriate, threatening, offensive,
+or harmful.
+
+Community leaders 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, and will communicate reasons for moderation
+decisions when appropriate.
+
+## Scope
+
+This Code of Conduct applies within all community spaces, and also applies when
+an individual is officially representing the community in public spaces.
+Examples of representing our community include using an official e-mail address,
+posting via an official social media account, or acting as an appointed
+representative at an online or offline event.
+
+## Enforcement
+
+Instances of abusive, harassing, or otherwise unacceptable behavior may be
+reported to the community leaders responsible for enforcement on Discord
+at https://discord.gg/QYF8QrtEUm
+
+All complaints will be reviewed and investigated promptly and fairly.
+
+All community leaders are obligated to respect the privacy and security of the
+reporter of any incident.
+
+## Enforcement Guidelines
+
+Community leaders will follow these Community Impact Guidelines in determining
+the consequences for any action they deem in violation of this Code of Conduct:
+
+### 1. Correction
+
+**Community Impact**: Use of inappropriate language or other behavior deemed
+unprofessional or unwelcome in the community.
+
+**Consequence**: A private, written warning from community leaders, providing
+clarity around the nature of the violation and an explanation of why the
+behavior was inappropriate. A public apology may be requested.
+
+### 2. Warning
+
+**Community Impact**: A violation through a single incident or series
+of actions.
+
+**Consequence**: A warning with consequences for continued behavior. No
+interaction with the people involved, including unsolicited interaction with
+those enforcing the Code of Conduct, for a specified period of time. This
+includes avoiding interactions in community spaces as well as external channels
+like social media. Violating these terms may lead to a temporary or
+permanent ban.
+
+### 3. Temporary Ban
+
+**Community Impact**: A serious violation of community standards, including
+sustained inappropriate behavior.
+
+**Consequence**: A temporary ban from any sort of interaction or public
+communication with the community for a specified period of time. No public or
+private interaction with the people involved, including unsolicited interaction
+with those enforcing the Code of Conduct, is allowed during this period.
+Violating these terms may lead to a permanent ban.
+
+### 4. Permanent Ban
+
+**Community Impact**: Demonstrating a pattern of violation of community
+standards, including sustained inappropriate behavior,  harassment of an
+individual, or aggression toward or disparagement of classes of individuals.
+
+**Consequence**: A permanent ban from any sort of public interaction within
+the community.
+
+## Attribution
+
+This Code of Conduct is adapted from the [Contributor Covenant][homepage],
+version 2.0, available at
+https://www.contributor-covenant.org/version/2/0/code_of_conduct.html.
+
+Community Impact Guidelines were inspired by [Mozilla's code of conduct
+enforcement ladder](https://github.com/mozilla/diversity).
+
+[homepage]: https://www.contributor-covenant.org
+
+For answers to common questions about this code of conduct, see the FAQ at
+https://www.contributor-covenant.org/faq. Translations are available at
+https://www.contributor-covenant.org/translations.
--- a/.github/CONTRIBUTING.md
+++ b/.github/CONTRIBUTING.md
@@ -0,0 +1,76 @@
+# Contributing to axolotl
+
+First of all, thank you for your interest in contributing to axolotl! We appreciate the time and effort you're willing to invest in making our project better. This document provides guidelines and information to make the contribution process as smooth as possible.
+
+## Table of Contents
+
+- [Code of Conduct](#code-of-conduct)
+- [Getting Started](#getting-started)
+- [How to Contribute](#how-to-contribute)
+  - [Reporting Bugs](#reporting-bugs)
+  - [Suggesting Enhancements](#suggesting-enhancements)
+  - [Submitting Pull Requests](#submitting-pull-requests)
+- [Style Guidelines](#style-guidelines)
+  - [Code Style](#code-style)
+  - [Commit Messages](#commit-messages)
+- [Additional Resources](#additional-resources)
+
+## Code of Conductcode
+
+All contributors are expected to adhere to our [Code of Conduct](CODE_OF_CONDUCT.md). Please read it before participating in the axolotl community.
+
+## Getting Started
+
+Bugs? Please check for open issue else create a new [Issue](https://github.com/OpenAccess-AI-Collective/axolotl/issues/new).
+
+PRs are **greatly welcome**!
+
+1. Fork the repository and clone it to your local machine.
+2. Set up the development environment by following the instructions in the [README.md](https://github.com/OpenAccess-AI-Collective/axolotl/tree/main/README.md) file.
+3. Explore the codebase, run tests, and verify that everything works as expected.
+
+Please run below to setup env
+```bash
+pip3 install -r requirements-dev.txt -r requirements-tests.txt
+pre-commit install
+
+# test
+pytest tests/
+```
+
+## How to Contribute
+
+### Reporting Bugs
+
+If you encounter a bug or issue while using axolotl, please open a new issue on the [GitHub Issues](https://github.com/OpenAccess-AI-Collective/axolotl/issues) page. Provide a clear and concise description of the problem, steps to reproduce it, and any relevant error messages or logs.
+
+### Suggesting Enhancements
+
+We welcome ideas for improvements and new features. To suggest an enhancement, open a new issue on the [GitHub Issues](https://github.com/OpenAccess-AI-Collective/axolotl/issues) page. Describe the enhancement in detail, explain the use case, and outline the benefits it would bring to the project.
+
+### Submitting Pull Requests
+
+1. Create a new branch for your feature or bugfix. Use a descriptive name like `feature/your-feature-name` or `fix/your-bugfix-name`.
+2. Make your changes, following the [Style Guidelines](#style-guidelines) below.
+3. Test your changes and ensure that they don't introduce new issues or break existing functionality.
+4. Commit your changes, following the [commit message guidelines](#commit-messages).
+5. Push your branch to your fork on GitHub.
+6. Open a new pull request against the `main` branch of the axolotl repository. Include a clear and concise description of your changes, referencing any related issues.
+
+## Style Guidelines
+
+### Code Style
+
+axolotl uses [{codestyle}]({URLofCodestyle}) as its code style guide. Please ensure that your code follows these guidelines.
+
+### Commit Messages
+
+Write clear and concise commit messages that briefly describe the changes made in each commit. Use the imperative mood and start with a capitalized verb, e.g., "Add new feature" or "Fix bug in function".
+
+## Additional Resources
+
+- [GitHub Help](https://help.github.com/)
+- [GitHub Pull Request Documentation](https://docs.github.com/en/github/collaborating-with-issues-and-pull-requests)
+- [{codestyle}]({URLofCodestyle})
+
+Thank you once again for your interest in contributing to axolotl. We look forward to collaborating with you and creating an even better project together!
--- a/.github/FUNDING.yml
+++ b/.github/FUNDING.yml
@@ -0,0 +1,13 @@
+# These are supported funding model platforms
+
+github: OpenAccess-AI-Collective # Replace with up to 4 GitHub Sponsors-enabled usernames e.g., [user1, user2]
+patreon: # Replace with a single Patreon username
+open_collective: # Replace with a single Open Collective username
+ko_fi: # Replace with a single Ko-fi username
+tidelift: # Replace with a single Tidelift platform-name/package-name e.g., npm/babel
+community_bridge: # Replace with a single Community Bridge project-name e.g., cloud-foundry
+liberapay: # Replace with a single Liberapay username
+issuehunt: # Replace with a single IssueHunt username
+otechie: # Replace with a single Otechie username
+lfx_crowdfunding: # Replace with a single LFX Crowdfunding project-name e.g., cloud-foundry
+custom: # Replace with up to 4 custom sponsorship URLs e.g., ['link1', 'link2']
--- a/.github/ISSUE_TEMPLATE/bug-report.yaml
+++ b/.github/ISSUE_TEMPLATE/bug-report.yaml
@@ -0,0 +1,105 @@
+name: Bug Report
+description: File a bug report
+labels: ["bug", "needs triage"]
+body:
+  - type: markdown
+    attributes:
+      value: |
+        ## Before you start
+        Please **make sure you are on the latest version.**
+        If you encountered the issue after you installed, updated, or reloaded, **please try restarting before reporting the bug**.
+
+  - type: checkboxes
+    id: no-duplicate-issues
+    attributes:
+      label: "Please check that this issue hasn't been reported before."
+      description: "The **Label filters** may help make your search more focussed."
+      options:
+        - label: "I searched previous [Bug Reports](https://github.com/OpenAccess-AI-Collective/axolotl/labels/bug) didn't find any similar reports."
+          required: true
+
+  - type: textarea
+    id: expected
+    attributes:
+      label: Expected Behavior
+      description: Tell us what **should** happen.
+    validations:
+      required: true
+
+  - type: textarea
+    id: what-happened
+    attributes:
+      label: Current behaviour
+      description: |
+        Tell us what happens instead of the expected behavior.
+        Provide stacktrace and/or screenshots.
+    validations:
+      required: true
+
+  - type: textarea
+    id: reproduce
+    attributes:
+      label: Steps to reproduce
+      description: |
+        Which exact steps can a developer take to reproduce the issue?
+        The more detail you provide, the easier it will be to narrow down and fix the bug.
+        Please paste in tasks and/or queries **as text, not screenshots**.
+      placeholder: |
+        Example of the level of detail needed to reproduce any bugs efficiently and reliably.
+        1. Go to the '...' page.
+        2. Click on the '...' button.
+        3. Scroll down to '...'.
+        4. Observe the error.
+    validations:
+      required: true
+
+  - type: textarea
+    id: possible-solution
+    attributes:
+      label: Possible solution
+      description: |
+        Not obligatory, but please suggest a fix or reason for the bug, if you have an idea.
+
+
+  - type: checkboxes
+    id: operating-systems
+    attributes:
+      label: Which Operating Systems are you using?
+      description: You may select more than one.
+      options:
+        - label: Linux
+        - label: macOS
+        - label: Windows
+
+  - type: input
+    id: Python-version
+    attributes:
+      label: Python Version
+      description: Which {Programming} version are you using?
+      placeholder: 3.10 / please change accordingly
+    validations:
+      required: true
+
+  - type: input
+    id: axolotl-branch-commit
+    attributes:
+      label: axolotl branch-commit
+      description: On which branch/commit are you?
+      placeholder: main/4d6490b
+    validations:
+      required: true
+
+  - type: checkboxes
+    id: acknowledgements
+    attributes:
+      label: 'Acknowledgements'
+      description: 'Please confirm the following:'
+      options:
+        - label: 'My issue title is concise, descriptive, and in title casing.'
+          required: true
+        - label: 'I have searched the existing issues to make sure this bug has not been reported yet.'
+          required: true
+        - label: 'I am using the latest version of axolotl.'
+          required: true
+        - label: 'I have provided enough information for the maintainers to reproduce and diagnose the issue.'
+          required: true
--- a/.github/ISSUE_TEMPLATE/config.yml
+++ b/.github/ISSUE_TEMPLATE/config.yml
@@ -0,0 +1,7 @@
+blank_issues_enabled: false
+contact_links:
+  - name: Ask a question
+    url: https://github.com/OpenAccess-AI-Collective/axolotl/discussions/categories/q-a
+    about: Ask questions and discuss with other community members
+  - name: Discuss the Project in Discord
+    url: https://discord.gg/HhrNrHJPRb
--- a/.github/ISSUE_TEMPLATE/docs.yml
+++ b/.github/ISSUE_TEMPLATE/docs.yml
@@ -0,0 +1,46 @@
+name: Documentation Improvement / Clarity
+description: Make a suggestion to improve the project documentation.
+labels: ['needs triage', 'docs']
+body:
+  - type: markdown
+    attributes:
+      value: '## :book: Documentation :book:'
+  - type: markdown
+    attributes:
+      value: |
+        * Ask questions in [Discord](https://discord.gg/HhrNrHJPRb).
+        * Before you file an issue read the [Contributing guide](./CONTRIBUTING.md).
+        * Check to make sure someone hasn't already opened a [similar issue](https://github.com/OpenAccess-AI-Collective/axolotl/issues).
+  - type: textarea
+    attributes:
+      label: What piece of documentation is affected?
+      description: Please link to the article you'd like to see updated.
+    validations:
+      required: true
+  - type: textarea
+    attributes:
+      label: What part(s) of the article would you like to see updated?
+      description: |
+        - Give as much detail as you can to help us understand the change you want to see.
+        - Why should the docs be changed? What use cases does it support?
+        - What is the expected outcome?
+    validations:
+      required: true
+  - type: textarea
+    attributes:
+      label: Additional Information
+      description: Add any other context or screenshots about the feature request here.
+    validations:
+      required: false
+  - type: checkboxes
+    id: acknowledgements
+    attributes:
+      label: 'Acknowledgements'
+      description: 'Please confirm the following:'
+      options:
+        - label: 'My issue title is concise, descriptive, and in title casing.'
+          required: true
+        - label: 'I have searched the existing issues to make sure this feature has not been requested yet.'
+          required: true
+        - label: 'I have provided enough information for the maintainers to understand and evaluate this request.'
+          required: true
--- a/.github/ISSUE_TEMPLATE/feature-request.yaml
+++ b/.github/ISSUE_TEMPLATE/feature-request.yaml
@@ -0,0 +1,63 @@
+name: Feature Request / Enhancement
+description: Suggest a new feature or feature enhancement for the project
+labels: ["enhancement", "needs triage"]
+body:
+  - type: checkboxes
+    id: no-duplicate-issues
+    attributes:
+      label: "⚠️ Please check that this feature request hasn't been suggested before."
+      description: "There are two locations for previous feature requests. Please search in both. Thank you. The **Label filters** may help make your search more focussed."
+      options:
+        - label: "I searched previous [Ideas in Discussions](https://github.com/OpenAccess-AI-Collective/axolotl/discussions/categories/ideas) didn't find any similar feature requests."
+          required: true
+        - label: "I searched previous [Issues](https://github.com/OpenAccess-AI-Collective/axolotl/labels/enhancement) didn't find any similar feature requests."
+          required: true
+
+  - type: textarea
+    id: feature-description
+    validations:
+      required: true
+    attributes:
+      label: "🔖 Feature description"
+      description: "A clear and concise description of what the feature request is."
+      placeholder: "You should add ..."
+
+  - type: textarea
+    id: solution
+    validations:
+      required: true
+    attributes:
+      label: "✔️ Solution"
+      description: "A clear and concise description of what you want to happen, and why."
+      placeholder: "In my use-case, ..."
+
+  - type: textarea
+    id: alternatives
+    validations:
+      required: false
+    attributes:
+      label: "❓ Alternatives"
+      description: "A clear and concise description of any alternative solutions or features you've considered."
+      placeholder: "I have considered ..."
+
+  - type: textarea
+    id: additional-context
+    validations:
+      required: false
+    attributes:
+      label: "📝 Additional Context"
+      description: "Add any other context or screenshots about the feature request here."
+      placeholder: "..."
+
+  - type: checkboxes
+    id: acknowledgements
+    attributes:
+      label: 'Acknowledgements'
+      description: 'Please confirm the following:'
+      options:
+        - label: 'My issue title is concise, descriptive, and in title casing.'
+          required: true
+        - label: 'I have searched the existing issues to make sure this feature has not been requested yet.'
+          required: true
+        - label: 'I have provided enough information for the maintainers to understand and evaluate this request.'
+          required: true
--- a/.github/PULL_REQUEST_TEMPLATE/pull_request_template_simple.md
+++ b/.github/PULL_REQUEST_TEMPLATE/pull_request_template_simple.md
@@ -0,0 +1,22 @@
+<!--- Provide a general summary of your changes in the Title above -->
+
+# Description
+
+<!--- Describe your changes in detail -->
+
+## Motivation and Context
+
+<!--- Why is this change required? What problem does it solve? -->
+<!--- If it fixes an open issue, please link to the issue here. -->
+
+## How has this been tested?
+
+<!--- Please describe in detail how you tested your changes. -->
+<!--- Include details of your testing environment, tests ran to see how -->
+<!--- your change affects other areas of the code, etc. -->
+
+## Screenshots (if appropriate)
+
+## Types of changes
+
+<!--- What types of changes does your code introduce? Put an `x` in all the boxes that apply: -->
--- a/.github/SECURITY.md
+++ b/.github/SECURITY.md
@@ -0,0 +1,9 @@
+# Security Policy
+
+## Supported Versions
+
+Due to the nature of the fast development that is happening in this project, only the latest released version can be supported.
+
+## Reporting a Vulnerability
+
+If you find a vulnerability, please contact us on  [Discord](https://discord.gg/xcu3ECkH9a) rather than creating a GitHub issue to allow us some time to fix it before it is a known vulnerability to others.
--- a/.github/SUPPORT.md
+++ b/.github/SUPPORT.md
@@ -0,0 +1,10 @@
+# Support
+
+If you need help with this project or have questions, please:
+
+1. Check the documentation.
+2. Search the existing issues and pull requests.
+3. Create a new issue if your question is not answered or your problem is not solved.
+4. Have a look in the [Discord server](https://discord.gg/HhrNrHJPRb)
+
+Please note that this project is maintained by volunteers who have limited availability. We'll do our best to address your questions and concerns in a timely manner.
--- a/.github/workflows/main.yml
+++ b/.github/workflows/main.yml
@@ -4,7 +4,6 @@ on:
  push:
    branches:
      - "main"
-      - "dev"

 jobs:
  build-axolotl:
@@ -72,6 +71,7 @@ jobs:
            python_version: "3.10"
            pytorch: 2.0.1
            axolotl_extras:
+            is_latest: true
          - cuda: 118
            cuda_version: 11.8.0
            python_version: "3.9"
@@ -102,5 +102,7 @@ jobs:
            CUDA=${{ matrix.cuda }}
          file: ./docker/Dockerfile-runpod
          push: ${{ github.event_name != 'pull_request' }}
-          tags: ${{ steps.metadata.outputs.tags }}-py${{ matrix.python_version }}-cu${{ matrix.cuda }}-${{ matrix.pytorch }}${{ matrix.axolotl_extras != '' && '-' || '' }}${{ matrix.axolotl_extras }}
+          tags: |
+             ${{ steps.metadata.outputs.tags }}-py${{ matrix.python_version }}-cu${{ matrix.cuda }}-${{ matrix.pytorch }}${{ matrix.axolotl_extras != '' && '-' || '' }}${{ matrix.axolotl_extras }}
+             ${{ (matrix.is_latest) && format('{0}-latest', steps.metadata.outputs.tags) || '' }}
          labels: ${{ steps.metadata.outputs.labels }}
--- a/README.md
+++ b/README.md
@@ -1,10 +1,39 @@
 # Axolotl

+Axolotl is a tool designed to streamline the fine-tuning of various AI models, offering support for multiple configurations and architectures.
+
+<table>
+<tr>
+<td>
+
+## Table of Contents
+- [Introduction](#axolotl)
+- [Supported Features](#axolotl-supports)
+- [Quickstart](#quickstart-)
+- [Installation](#installation)
+  - [Docker Installation](#environment)
+  - [Conda/Pip venv Installation](#condapip-venv)
+  - [LambdaLabs Installation](#lambdalabs)
+- [Dataset](#dataset)
+  - [How to Add Custom Prompts](#how-to-add-custom-prompts)
+- [Config](#config)
+  - [Train](#train)
+  - [Inference](#inference)
+  - [Merge LORA to Base](#merge-lora-to-base)
+- [Common Errors](#common-errors-)
+- [Need Help?](#need-help-)
+- [Badge](#badge-)
+- [Community Showcase](#community-showcase)
+- [Contributing](#contributing-)
+
+</td>
+<td>
+
 <div align="center">
  <img src="image/axolotl.png" alt="axolotl" width="160">
  <div>
    <p>
-      <b>One repo to finetune them all! </b>
+      <b>Axolotl provides a unified repository for fine-tuning <br />a variety of AI models with ease</b>
    </p>
    <p>
      Go ahead and axolotl questions!!
@@ -14,21 +43,27 @@
  </div>
 </div>

+</td>
+</tr>
+</table>
+
 ## Axolotl supports

-|          | fp16/fp32 | lora | qlora | gptq | gptq w/ lora | gptq w/flash attn | flash attn | xformers attn |
-|----------|:----------|:-----|-------|------|:-------------|-------------------|------------|---------------|
-| llama    | ✅         | ✅    | ✅     | ✅    | ✅             | ✅                 | ✅          | ✅             |
-| Pythia   | ✅         | ✅    | ✅     | ❌    | ❓            | ❌                 | ❌          | ❓             |
-| cerebras | ✅         | ✅    | ✅     | ❌    | ❓            | ❌                 | ❌          | ✅             |
-| mpt      | ✅         | ❌    | ❓     | ❌    | ❓            | ❌                 | ❌          | ❓             |
-| falcon   | ✅         | ✅    | ✅     | ❌    | ❓            | ❌                 | ❌          | ✅             |
-| gpt-j    | ✅         | ✅    | ✅     | ❌    | ❓            | ❌                 | ❓          | ✅             |
-| XGen     | ✅         | ❓    | ✅     | ❓    | ❓            | ❓                 | ❓          | ✅
+|          | fp16/fp32 | lora | qlora | gptq | gptq w/flash attn | flash attn | xformers attn |
+|----------|:----------|:-----|-------|------|-------------------|------------|---------------|
+| llama    | ✅         | ✅    | ✅     | ✅             | ✅                 | ✅          | ✅             |
+| Pythia   | ✅         | ✅    | ✅     | ❌             | ❌                 | ❌          | ❓             |
+| cerebras | ✅         | ✅    | ✅     | ❌             | ❌                 | ❌          | ❓             |
+| mpt      | ✅         | ❌    | ❓     | ❌             | ❌                 | ❌          | ❓             |
+| falcon   | ✅         | ✅    | ✅     | ❌             | ❌                 | ❌          | ❓             |
+| gpt-j    | ✅         | ✅    | ✅     | ❌             | ❌                 | ❓          | ❓             |
+| XGen     | ✅         | ❓    | ✅     | ❓             | ❓                 | ❓          | ✅             |


 ## Quickstart ⚡

+Get started with Axolotl in just a few steps! This quickstart guide will walk you through setting up and running a basic fine-tuning task.
+
 **Requirements**: Python >=3.9 and Pytorch >=2.0.

 ```bash
@@ -130,13 +165,14 @@ accelerate launch scripts/finetune.py examples/openllama-3b/lora.yml \

 ### Dataset

+Axolotl supports a variety of dataset formats. Below are some of the formats you can use.
 Have dataset(s) in one of the following format (JSONL recommended):

 - `alpaca`: instruction; input(optional)
  ```json
  {"instruction": "...", "input": "...", "output": "..."}
  ```
- `sharegpt:chat`: conversations
+- `sharegpt:chat`: conversations where `from` is `human`/`gpt`
  ```json
  {"conversations": [{"from": "...", "value": "..."}]}
  ```
@@ -225,6 +261,10 @@ Have dataset(s) in one of the following format (JSONL recommended):
  ```json
  {"conversations": [{"role": "...", "value": "..."}]}
  ```
+- `sharegpt_simple.load_guanaco`: conversations where `from` is `prompter`/`assistant` instead of default sharegpt
+  ```json
+  {"conversations": [{"from": "...", "value": "..."}]}
+  ```
 - `sharegpt_jokes`: creates a chat where bot is asked to tell a joke, then explain why the joke is funny
  ```json
  {"conversations": [{"title": "...", "text": "...", "explanation": "..."}]}
@@ -322,6 +362,8 @@ tokenizer_type: AutoTokenizer
 trust_remote_code:
 # use_fast option for tokenizer loading from_pretrained, default to True
 tokenizer_use_fast:
+# Whether to use the legacy tokenizer setting, defaults to True
+tokenizer_legacy:
 # resize the model embeddings when new tokens are added to multiples of 32
 # this is reported to improve training speed on some models
 resize_token_embeddings_to_32x:
@@ -360,10 +402,13 @@ dataset_prepared_path: data/last_run_prepared
 push_dataset_to_hub: # repo path
 # push checkpoints to hub
 hub_model_id: # repo path to push finetuned model
+# how to push checkpoints to hub
+# https://huggingface.co/docs/transformers/v4.31.0/en/main_classes/trainer#transformers.TrainingArguments.hub_strategy
+hub_strategy:
 # whether to use hf `use_auth_token` for loading datasets. Useful for fetching private datasets
 # required to be true when used in combination with `push_dataset_to_hub`
 hf_use_auth_token: # boolean
-# How much of the dataset to set aside as evaluation. 1 = 100%, 0.50 = 50%, etc
+# How much of the dataset to set aside as evaluation. 1 = 100%, 0.50 = 50%, etc. 0 for no eval.
 val_set_size: 0.04
 # Num shards for whole dataset
 dataset_shard_num:
@@ -425,10 +470,12 @@ eval_batch_size: 2
 num_epochs: 3
 warmup_steps: 100
 learning_rate: 0.00003
+lr_quadratic_warmup:
 logging_steps:
-save_steps:
+save_steps: # leave empty to save at each epoch
 eval_steps:
-save_total_limit:
+save_total_limit: # checkpoints saved at a time
+max_steps:

 # save model as safetensors (require safetensors package)
 save_safetensors:
@@ -473,8 +520,8 @@ max_grad_norm:
 flash_optimum:
 # whether to use xformers attention patch https://github.com/facebookresearch/xformers:
 xformers_attention:
-# whether to use flash attention patch https://github.com/HazyResearch/flash-attention:
-flash_attention:  # require a100 for llama
+# whether to use flash attention patch https://github.com/Dao-AILab/flash-attention:
+flash_attention:
 # whether to use scaled-dot-product attention
 # https://pytorch.org/docs/stable/generated/torch.nn.functional.scaled_dot_product_attention.html
 sdp_attention:
@@ -509,7 +556,7 @@ tokens:
 fsdp:
 fsdp_config:

-# Deepspeed
+# Deepspeed config path
 deepspeed:

 # Path to torch distx for optim 'adamw_anyprecision'
@@ -560,7 +607,10 @@ fsdp_config:
  fsdp_transformer_layer_cls_to_wrap: LlamaDecoderLayer
 ```

- llama Deepspeed: append `ACCELERATE_USE_DEEPSPEED=true` in front of finetune command
+- llama Deepspeed
+```yaml
+deepspeed: deepspeed/zero3.json
+```

 ##### Weights & Biases Logging

@@ -608,7 +658,7 @@ CUDA_VISIBLE_DEVICES="" python3 scripts/finetune.py ...

 ## Common Errors 🧰

-> Cuda out of memory
+> If you encounter a 'Cuda out of memory' error, it means your GPU ran out of memory during the training process. Here's how to resolve it:

 Please reduce any below
  - `micro_batch_size`
@@ -616,6 +666,10 @@ Please reduce any below
  - `gradient_accumulation_steps`
  - `sequence_len`

+> `failed (exitcode: -9)` usually means your system has run out of system memory.
+Similarly, you should consider reducing the same settings as when you run out of VRAM.
+Additionally, look into upgrading your system RAM which should be simpler than GPU upgrades.
+
 > RuntimeError: expected scalar type Float but found Half

 Try set `fp16: true`
@@ -644,6 +698,8 @@ Building something cool with Axolotl? Consider adding a badge to your model card

 ## Community Showcase

+Check out some of the projects and models that have been built using Axolotl! Have a model you'd like to add to our Community Showcase? Open a PR with your model.
+
 Open Access AI Collective
 - [Minotaur 13b](https://huggingface.co/openaccess-ai-collective/minotaur-13b)
 - [Manticore 13b](https://huggingface.co/openaccess-ai-collective/manticore-13b)
@@ -654,7 +710,9 @@ PocketDoc Labs

 ## Contributing 🤝

-Bugs? Please check for open issue else create a new [Issue](https://github.com/OpenAccess-AI-Collective/axolotl/issues/new).
+Please read the [contributing guide](./.github/CONTRIBUTING.md)
+
+Bugs? Please check the [open issues](https://github.com/OpenAccess-AI-Collective/axolotl/issues/bug) else create a new Issue.

 PRs are **greatly welcome**!

--- a/docker/Dockerfile-base
+++ b/docker/Dockerfile-base
@@ -40,7 +40,7 @@ ARG TORCH_CUDA_ARCH_LIST="7.0 7.5 8.0 8.6 9.0+PTX"

 RUN git clone https://github.com/Dao-AILab/flash-attention.git && \
    cd flash-attention && \
-    git checkout v2.0.1  && \
+    git checkout v2.0.4  && \
    python3 setup.py bdist_wheel && \
    cd csrc/fused_dense_lib && \
    python3 setup.py bdist_wheel && \
--- a/examples/llama-2/lora.yml
+++ b/examples/llama-2/lora.yml
@@ -2,6 +2,7 @@ base_model: meta-llama/Llama-2-7b-hf
 base_model_config: meta-llama/Llama-2-7b-hf
 model_type: LlamaForCausalLM
 tokenizer_type: LlamaTokenizer
+is_llama_derived_model: true

 load_in_8bit: true
 load_in_4bit: false
@@ -15,7 +16,7 @@ val_set_size: 0.01
 output_dir: ./lora-out

 sequence_len: 4096
-max_packed_sequence_len: 4096
+sample_packing: true

 adapter: lora
 lora_model_dir:
@@ -49,8 +50,8 @@ early_stopping_patience:
 resume_from_checkpoint:
 local_rank:
 logging_steps: 1
-xformers_attention: true
-flash_attention:
+xformers_attention:
+flash_attention: true

 warmup_steps: 10
 eval_steps: 20
@@ -64,4 +65,3 @@ special_tokens:
  bos_token: "<s>"
  eos_token: "</s>"
  unk_token: "<unk>"
-  pad_token: "<pad>"
--- a/examples/llama-2/qlora.yml
+++ b/examples/llama-2/qlora.yml
@@ -2,6 +2,7 @@ base_model: meta-llama/Llama-2-7b-hf
 base_model_config: meta-llama/Llama-2-7b-hf
 model_type: LlamaForCausalLM
 tokenizer_type: LlamaTokenizer
+is_llama_derived_model: true

 load_in_8bit: false
 load_in_4bit: true
@@ -18,7 +19,8 @@ adapter: qlora
 lora_model_dir:

 sequence_len: 4096
-max_packed_sequence_len: 4096
+sample_packing: true
+
 lora_r: 32
 lora_alpha: 16
 lora_dropout: 0.05
@@ -50,8 +52,8 @@ early_stopping_patience:
 resume_from_checkpoint:
 local_rank:
 logging_steps: 1
-xformers_attention: true
-flash_attention:
+xformers_attention:
+flash_attention: true

 warmup_steps: 10
 eval_steps: 20
@@ -65,4 +67,3 @@ special_tokens:
  bos_token: "<s>"
  eos_token: "</s>"
  unk_token: "<unk>"
-  pad_token: "<pad>"
--- a/scripts/finetune.py
+++ b/scripts/finetune.py
@@ -18,18 +18,13 @@ from optimum.bettertransformer import BetterTransformer
 from transformers import GenerationConfig, TextStreamer

 from axolotl.logging_config import configure_logging
-from axolotl.utils.bench import log_gpu_memory_usage
-from axolotl.utils.data import load_prepare_datasets, load_pretraining_dataset
+from axolotl.utils.config import normalize_config, validate_config
+from axolotl.utils.data import prepare_dataset
 from axolotl.utils.dict import DictDefault
-from axolotl.utils.distributed import barrier, is_main_process
+from axolotl.utils.distributed import is_main_process
 from axolotl.utils.models import load_model, load_tokenizer
 from axolotl.utils.tokenization import check_dataset_labels
-from axolotl.utils.trainer import (
-    calculate_total_num_steps,
-    process_datasets_for_packing,
-    setup_trainer,
-)
-from axolotl.utils.validation import validate_config
+from axolotl.utils.trainer import setup_trainer
 from axolotl.utils.wandb import setup_wandb_env_vars

 project_root = os.path.abspath(os.path.join(os.path.dirname(__file__), ".."))
@@ -39,30 +34,21 @@ sys.path.insert(0, src_dir)
 configure_logging()
 LOG = logging.getLogger("axolotl.scripts")

-
-DEFAULT_DATASET_PREPARED_PATH = "last_run_prepared"
 os.environ["HF_HUB_ENABLE_HF_TRANSFER"] = "1"


-def choose_device(cfg):
-    def get_device():
-        try:
-            if torch.cuda.is_available():
-                return f"cuda:{cfg.local_rank}"
+def print_axolotl_text_art():
+    ascii_art = """
+                           dP            dP   dP
+                           88            88   88
+.d8888b. dP.  .dP .d8888b. 88 .d8888b. d8888P 88
+88'  `88  `8bd8'  88'  `88 88 88'  `88   88   88
+88.  .88  .d88b.  88.  .88 88 88.  .88   88   88
+`88888P8 dP'  `dP `88888P' dP `88888P'   dP   dP
+"""

-            if torch.backends.mps.is_available():
-                return "mps"
-
-            raise SystemError("No CUDA/mps device found")
-        except Exception:  # pylint: disable=broad-exception-caught
-            return "cpu"
-
-    cfg.device = get_device()
-    if cfg.device_map != "auto":
-        if cfg.device.startswith("cuda"):
-            cfg.device_map = {"": cfg.local_rank}
-        else:
-            cfg.device_map = {"": cfg.device}
+    if is_main_process():
+        print(ascii_art)


 def get_multi_line_input() -> Optional[str]:
@@ -174,6 +160,7 @@ def train(
    prepare_ds_only: bool = False,
    **kwargs,
 ):
+    print_axolotl_text_art()
    if Path(config).is_dir():
        config = choose_config(config)

@@ -194,67 +181,18 @@ def train(

    validate_config(cfg)

-    # setup some derived config / hyperparams
-    cfg.gradient_accumulation_steps = cfg.gradient_accumulation_steps or (
-        cfg.batch_size // cfg.micro_batch_size
-    )
-    cfg.batch_size = (
-        cfg.batch_size or cfg.micro_batch_size * cfg.gradient_accumulation_steps
-    )
-    cfg.world_size = int(os.environ.get("WORLD_SIZE", 1))
-    cfg.local_rank = int(os.environ.get("LOCAL_RANK", 0))
-    choose_device(cfg)
-    cfg.ddp = cfg.ddp if cfg.ddp is not None else cfg.world_size != 1
-    if cfg.ddp:
-        cfg.device_map = {"": int(os.environ.get("LOCAL_RANK", 0))}
-        cfg.batch_size = cfg.batch_size * cfg.world_size
+    normalize_config(cfg)

    setup_wandb_env_vars(cfg)
-    if cfg.device == "mps":
-        cfg.load_in_8bit = False
-        cfg.tf32 = False
-        if cfg.bf16:
-            cfg.fp16 = True
-        cfg.bf16 = False
-
-    if cfg.tf32:
-        torch.backends.cuda.matmul.allow_tf32 = True

    # load the tokenizer first
-    tokenizer_config = cfg.tokenizer_config or cfg.base_model_config
-    LOG.info(f"loading tokenizer... {tokenizer_config}")
-    tokenizer = load_tokenizer(tokenizer_config, cfg.tokenizer_type, cfg)
+    LOG.info(f"loading tokenizer... {cfg.tokenizer_config or cfg.base_model_config}")
+    tokenizer = load_tokenizer(cfg)

    if (
        check_not_in(["shard", "merge_lora"], kwargs) and not cfg.inference
    ):  # don't need to load dataset for these
-        if not cfg.pretraining_dataset:
-            train_dataset, eval_dataset = load_prepare_datasets(
-                tokenizer, cfg, DEFAULT_DATASET_PREPARED_PATH
-            )
-        else:
-            train_dataset = load_pretraining_dataset(
-                cfg.pretraining_dataset,
-                tokenizer,
-                max_tokens=cfg.sequence_len,
-                seed=cfg.seed or 42,
-            )
-            # https://discuss.huggingface.co/t/how-to-use-huggingface-trainer-streaming-datasets-without-wrapping-it-with-torchdatas-iterablewrapper/25230
-            train_dataset = train_dataset.with_format("torch")
-            eval_dataset = None
-
-        if is_main_process():
-            # process on rank 0 first so it gets cached so other ranks load from cache
-            train_dataset, eval_dataset = process_datasets_for_packing(
-                cfg, train_dataset, eval_dataset
-            )
-        barrier()
-        if not is_main_process():
-            train_dataset, eval_dataset = process_datasets_for_packing(
-                cfg, train_dataset, eval_dataset
-            )
-        barrier()
-        total_num_steps = calculate_total_num_steps(cfg, train_dataset, tokenizer)
+        train_dataset, eval_dataset, total_num_steps = prepare_dataset(cfg, tokenizer)

    if cfg.debug or "debug" in kwargs:
        LOG.info("check_dataset_labels...")
@@ -269,8 +207,6 @@ def train(
        LOG.info("Finished preparing dataset. Exiting...")
        return

-    log_gpu_memory_usage(LOG, "baseline", cfg.device)
-
    # Load the model and tokenizer
    LOG.info("loading model and (optionally) peft_config...")
    model, peft_config = load_model(cfg, tokenizer)
@@ -354,6 +290,7 @@ def train(

    if not Path(cfg.output_dir).is_dir():
        os.makedirs(cfg.output_dir, exist_ok=True)
+    tokenizer.save_pretrained(cfg.output_dir)
    if cfg.flash_optimum:
        with torch.backends.cuda.sdp_kernel(
            enable_flash=True, enable_math=True, enable_mem_efficient=True
--- a/src/axolotl/monkeypatch/llama_attn_hijack_flash.py
+++ b/src/axolotl/monkeypatch/llama_attn_hijack_flash.py
@@ -2,142 +2,47 @@

 # copied from https://github.com/lm-sys/FastChat/blob/main/fastchat/train/llama_flash_attn_monkey_patch.py

-from typing import Optional, Tuple
+import warnings
+from typing import List, Optional, Tuple, Union

 import torch
+import torch.nn.functional as F
 import transformers
 from einops import rearrange
 from flash_attn.bert_padding import pad_input, unpad_input
+from transformers.modeling_outputs import BaseModelOutputWithPast
+from transformers.models.llama.modeling_llama import (
+    LlamaDecoderLayer as OriginalLlamaDecoderLayer,
+)
+from transformers.models.llama.modeling_llama import apply_rotary_pos_emb, repeat_kv
+
+from axolotl.monkeypatch.utils import get_cu_seqlens_from_pos_ids

 try:
-    from flash_attn.flash_attn_interface import flash_attn_varlen_qkvpacked_func
+    from flash_attn.flash_attn_interface import (  # pylint: disable=ungrouped-imports
+        flash_attn_kvpacked_func,
+        flash_attn_varlen_kvpacked_func,
+        flash_attn_varlen_qkvpacked_func,
+    )
 except ImportError:
+    from flash_attn.flash_attn_interface import (
+        flash_attn_unpadded_kvpacked_func as flash_attn_varlen_kvpacked_func,
+    )
    from flash_attn.flash_attn_interface import (
        flash_attn_unpadded_qkvpacked_func as flash_attn_varlen_qkvpacked_func,
    )

-from transformers.models.llama.modeling_llama import apply_rotary_pos_emb

-from axolotl.monkeypatch.utils import get_cu_seqlens_from_pos_ids
-
-
-def forward(
-    self,
-    hidden_states: torch.Tensor,
-    attention_mask: Optional[torch.Tensor] = None,
-    position_ids: Optional[torch.Tensor] = None,
-    past_key_value: Optional[Tuple[torch.Tensor]] = None,
-    output_attentions: bool = False,
-    use_cache: bool = False,
-) -> Tuple[torch.Tensor, Optional[torch.Tensor], Optional[Tuple[torch.Tensor]]]:
-    """Input shape: Batch x Time x Channel
-
-    attention_mask: [bsz, q_len]
-    """
-    # pylint: disable=duplicate-code
-    bsz, q_len, _ = hidden_states.size()
-
-    query_states = (
-        self.q_proj(hidden_states)
-        .view(bsz, q_len, self.num_heads, self.head_dim)
-        .transpose(1, 2)
+def replace_llama_attn_with_flash_attn(packed: Optional[bool] = False):
+    transformers.models.llama.modeling_llama.LlamaModel._prepare_decoder_attention_mask = (  # pylint: disable=protected-access
+        _prepare_decoder_attention_mask
    )
-    key_states = (
-        self.k_proj(hidden_states)
-        .view(bsz, q_len, self.num_heads, self.head_dim)
-        .transpose(1, 2)
-    )
-    value_states = (
-        self.v_proj(hidden_states)
-        .view(bsz, q_len, self.num_heads, self.head_dim)
-        .transpose(1, 2)
-    )
-    # [bsz, q_len, nh, hd]
-    # [bsz, nh, q_len, hd]
-
-    kv_seq_len = key_states.shape[-2]
-    assert past_key_value is None, "past_key_value is not supported"
-
-    cos, sin = self.rotary_emb(value_states, seq_len=kv_seq_len)
-    query_states, key_states = apply_rotary_pos_emb(
-        query_states, key_states, cos, sin, position_ids
-    )
-    # [bsz, nh, t, hd]
-    assert not output_attentions, "output_attentions is not supported"
-    assert not use_cache, "use_cache is not supported"
-
-    # Flash attention codes from
-    # https://github.com/HazyResearch/flash-attention/blob/main/flash_attn/flash_attention.py
-
-    # transform the data into the format required by flash attention
-    qkv = torch.stack(
-        [query_states, key_states, value_states], dim=2
-    )  # [bsz, nh, 3, q_len, hd]
-    qkv = qkv.transpose(1, 3)  # [bsz, q_len, 3, nh, hd]
-    # We have disabled _prepare_decoder_attention_mask in LlamaModel
-    # the attention_mask should be the same as the key_padding_mask
-    key_padding_mask = attention_mask
-
-    if key_padding_mask is None:
-        qkv = rearrange(qkv, "b s ... -> (b s) ...")
-        max_s = q_len
-        cu_q_lens = torch.arange(
-            0,
-            (bsz + 1) * q_len,
-            step=q_len,
-            dtype=torch.int32,
-            device=qkv.device,
+    transformers.models.llama.modeling_llama.LlamaAttention.forward = flashattn_forward
+    if packed:
+        transformers.models.llama.modeling_llama.LlamaDecoderLayer = LlamaDecoderLayer
+        transformers.models.llama.modeling_llama.LlamaModel.forward = (
+            llama_model_forward
        )
-        output = flash_attn_varlen_qkvpacked_func(
-            qkv, cu_q_lens, max_s, 0.0, softmax_scale=None, causal=True
-        )
-        output = rearrange(output, "(b s) ... -> b s ...", b=bsz)
-    elif position_ids.shape[0] == 1:
-        # special handling using sample packing
-        qkv = rearrange(qkv, "b s ... -> (b s) ...")
-        cu_q_lens, max_s = get_cu_seqlens_from_pos_ids(position_ids)
-        cu_q_lens = cu_q_lens.squeeze()
-
-        output = flash_attn_varlen_qkvpacked_func(
-            qkv, cu_q_lens, max_s, 0.0, softmax_scale=None, causal=True
-        )
-        output = rearrange(output, "(b s) ... -> b s ...", b=bsz)
-    else:
-        nheads = qkv.shape[-2]
-
-        # pylint: disable=invalid-name
-        x = rearrange(qkv, "b s three h d -> b s (three h d)")
-        x_unpad, indices, cu_q_lens, max_s = unpad_input(x, key_padding_mask)
-        x_unpad = rearrange(
-            x_unpad,
-            "nnz (three h d) -> nnz three h d",
-            three=3,
-            h=nheads,
-        )
-        output_unpad = flash_attn_varlen_qkvpacked_func(
-            x_unpad,
-            cu_q_lens,
-            max_s,
-            0.0,
-            softmax_scale=None,
-            causal=True,
-        )
-        output = rearrange(
-            pad_input(
-                rearrange(output_unpad, "nnz h d -> nnz (h d)"),
-                indices,
-                bsz,
-                q_len,
-            ),
-            "b s (h d) -> b s h d",
-            h=nheads,
-        )
-
-    return (
-        self.o_proj(rearrange(output, "b s h d -> b s (h d)")),
-        None,
-        None,
-    )


 # Disable the transformation of the attention mask in LlamaModel as the flash attention
@@ -153,8 +58,541 @@ def _prepare_decoder_attention_mask(
    return attention_mask


-def replace_llama_attn_with_flash_attn():
-    transformers.models.llama.modeling_llama.LlamaModel._prepare_decoder_attention_mask = (  # pylint: disable=protected-access
-        _prepare_decoder_attention_mask
+def flashattn_forward(
+    self,
+    hidden_states: torch.Tensor,
+    attention_mask: Optional[torch.Tensor] = None,
+    position_ids: Optional[torch.Tensor] = None,
+    past_key_value: Optional[Tuple[torch.Tensor]] = None,
+    output_attentions: bool = False,
+    use_cache: bool = False,
+    cu_seqlens: Optional[torch.Tensor] = None,
+    max_seqlen: Optional[torch.Tensor] = None,
+) -> Tuple[torch.Tensor, Optional[torch.Tensor], Optional[Tuple[torch.Tensor]]]:
+    """Input shape: Batch x Time x Channel
+
+    attention_mask: [bsz, q_len]
+    """
+    # pylint: disable=duplicate-code
+    bsz, q_len, _ = hidden_states.size()
+
+    if not hasattr(self, "pretraining_tp"):
+        self.pretraining_tp = 1
+
+    if self.pretraining_tp > 1:
+        key_value_slicing = (
+            self.num_key_value_heads * self.head_dim
+        ) // self.pretraining_tp
+        query_slices = self.q_proj.weight.split(
+            (self.num_heads * self.head_dim) // self.pretraining_tp, dim=0
+        )
+        key_slices = self.k_proj.weight.split(key_value_slicing, dim=0)
+        value_slices = self.v_proj.weight.split(key_value_slicing, dim=0)
+
+        query_states = [
+            F.linear(hidden_states, query_slices[i]) for i in range(self.pretraining_tp)
+        ]
+        query_states = torch.cat(query_states, dim=-1)
+
+        key_states = [
+            F.linear(hidden_states, key_slices[i]) for i in range(self.pretraining_tp)
+        ]
+        key_states = torch.cat(key_states, dim=-1)
+
+        value_states = [
+            F.linear(hidden_states, value_slices[i]) for i in range(self.pretraining_tp)
+        ]
+        value_states = torch.cat(value_states, dim=-1)
+
+    else:
+        query_states = self.q_proj(hidden_states)
+        key_states = self.k_proj(hidden_states)
+        value_states = self.v_proj(hidden_states)
+
+    query_states = query_states.view(
+        bsz, q_len, self.num_heads, self.head_dim
+    ).transpose(1, 2)
+    key_states = key_states.view(
+        bsz, q_len, self.num_key_value_heads, self.head_dim
+    ).transpose(1, 2)
+    value_states = value_states.view(
+        bsz, q_len, self.num_key_value_heads, self.head_dim
+    ).transpose(1, 2)
+    # [bsz, q_len, nh, hd]
+    # [bsz, nh, q_len, hd]
+
+    kv_seq_len = key_states.shape[-2]
+    if past_key_value is not None:
+        kv_seq_len += past_key_value[0].shape[-2]
+
+    cos, sin = self.rotary_emb(value_states, seq_len=kv_seq_len)
+    query_states, key_states = apply_rotary_pos_emb(
+        query_states, key_states, cos, sin, position_ids
    )
-    transformers.models.llama.modeling_llama.LlamaAttention.forward = forward
+    # [bsz, nh, t, hd]
+
+    if past_key_value is not None:
+        # reuse k, v, self_attention
+        key_states = torch.cat([past_key_value[0], key_states], dim=2)
+        value_states = torch.cat([past_key_value[1], value_states], dim=2)
+
+    past_key_value = (key_states, value_states) if use_cache else None
+
+    # repeat k/v heads if n_kv_heads < n_heads
+    key_states = repeat_kv(key_states, self.num_key_value_groups)
+    value_states = repeat_kv(value_states, self.num_key_value_groups)
+
+    if output_attentions:
+        warnings.warn(
+            "Output attentions is not supported for patched `LlamaAttention`, returning `None` instead."
+        )
+
+    #
+    # flash-attn v2 start
+    #
+
+    if self.training:
+        # during training q,k,v always have same seqlen
+        assert key_states.shape == query_states.shape
+        is_causal = True
+    else:
+        # turn off FA causal mask after first inference autoregressive iteration
+        # only on first autoregressive step q,k,v have same seqlen
+        is_causal = past_key_value is not None
+
+    if cu_seqlens is not None and max_seqlen is not None:
+        # special handling using sample packing
+        qkv = torch.stack(
+            [query_states, key_states, value_states], dim=2
+        )  # [bsz, nh, 3, q_len, hd]
+        qkv = qkv.transpose(1, 3)  # [bsz, q_len, 3, nh, hd]
+        qkv = rearrange(qkv, "b s ... -> (b s) ...")
+
+        output = flash_attn_varlen_qkvpacked_func(
+            qkv, cu_seqlens, max_seqlen, 0.0, softmax_scale=None, causal=is_causal
+        )
+        output = rearrange(output, "(b s) ... -> b s ...", b=bsz)
+    elif query_states.shape == key_states.shape:
+        query_states = query_states.transpose(1, 2)
+        key_states = key_states.transpose(1, 2)
+        value_states = value_states.transpose(1, 2)
+        qkv_unpad, cu_seqlens_q, max_seqlen_q, _, output_pad_fn = generate_qkv(
+            query_states,
+            key_states,
+            value_states,
+            qkvpacked=True,
+            # We have disabled _prepare_decoder_attention_mask in LlamaModel
+            # the attention_mask should be the same as the key_padding_mask
+            key_padding_mask=attention_mask,
+            query_padding_mask=attention_mask[:, -query_states.size(1) :]
+            if attention_mask is not None
+            else None,
+        )
+        output_unpad = flash_attn_varlen_qkvpacked_func(
+            qkv_unpad,
+            cu_seqlens_q,
+            max_seqlen_q,
+            0.0,
+            softmax_scale=None,
+            causal=is_causal,
+        )
+        output = output_pad_fn(output_unpad)
+    else:
+        query_states = query_states.transpose(1, 2)
+        key_states = key_states.transpose(1, 2)
+        value_states = value_states.transpose(1, 2)
+        if attention_mask is None or attention_mask.all().item():
+            output = flash_attn_kvpacked_func(
+                query_states,
+                torch.stack([key_states, value_states], 2),
+                causal=is_causal,
+            )
+        else:
+            (  # pylint: disable=unbalanced-tuple-unpacking
+                q_unpad,
+                kv_unpad,
+                cu_seqlens_q,
+                cu_seqlens_k,
+                max_seqlen_q,
+                max_seqlen_k,
+                _,
+                _,
+                output_pad_fn,
+            ) = generate_qkv(
+                query_states,
+                key_states,
+                value_states,
+                kvpacked=True,
+                key_padding_mask=attention_mask,
+                query_padding_mask=attention_mask[:, -query_states.size(1) :]
+                if attention_mask is not None
+                else None,
+            )
+            output_unpad = flash_attn_varlen_kvpacked_func(
+                q_unpad,
+                kv_unpad,
+                cu_seqlens_q,
+                cu_seqlens_k,
+                max_seqlen_q,
+                max_seqlen_k,
+                0.0,
+                softmax_scale=None,
+                causal=is_causal,
+            )
+            output = output_pad_fn(output_unpad)
+
+    attn_output = output
+    if attn_output.size() != (bsz, q_len, self.num_heads, self.head_dim):
+        raise ValueError(
+            f"`attn_output` should be of size {(bsz, q_len, self.num_heads, self.head_dim)}, but is"
+            f" {attn_output.size()}"
+        )
+    attn_output = rearrange(attn_output, "b s h d -> b s (h d)")
+
+    #
+    # flash-attn v2 end
+    #
+
+    if self.pretraining_tp > 1:
+        attn_output = attn_output.split(self.hidden_size // self.pretraining_tp, dim=2)
+        o_proj_slices = self.o_proj.weight.split(
+            self.hidden_size // self.pretraining_tp, dim=1
+        )
+        attn_output = sum(
+            F.linear(attn_output[i], o_proj_slices[i])
+            for i in range(self.pretraining_tp)
+        )
+    else:
+        attn_output = self.o_proj(attn_output)
+
+    return attn_output, None, past_key_value
+
+
+# based on https://github.com/Dao-AILab/flash-attention/blob/364a5b/tests/test_flash_attn.py#L38
+def generate_qkv(
+    q,
+    k,
+    v,
+    query_padding_mask=None,
+    key_padding_mask=None,
+    kvpacked=False,
+    qkvpacked=False,
+):  # pylint: disable=invalid-name,unnecessary-lambda-assignment
+    """
+    Arguments:
+        q: (batch_size, seqlen_q, nheads, d)
+        k: (batch_size, seqlen_k, nheads_k, d)
+        v: (batch_size, seqlen_k, nheads_k, d)
+        query_padding_mask: (batch_size, seqlen), bool
+        key_padding_mask: (batch_size, seqlen), bool
+    """
+    assert not (kvpacked and qkvpacked)
+    batch_size, seqlen_q, nheads, d = q.shape
+    _, seqlen_k, nheads_k, _ = k.shape
+    assert k.shape == (batch_size, seqlen_k, nheads_k, d)
+    assert v.shape == (batch_size, seqlen_k, nheads_k, d)
+
+    if query_padding_mask is not None:
+        q_unpad, indices_q, cu_seqlens_q, max_seqlen_q = unpad_input(
+            q, query_padding_mask
+        )
+
+        output_pad_fn = lambda output_unpad: pad_input(  # noqa: E731
+            output_unpad, indices_q, batch_size, seqlen_q
+        )
+
+    else:
+        q_unpad = rearrange(q, "b s h d -> (b s) h d")
+        cu_seqlens_q = torch.arange(
+            0,
+            (batch_size + 1) * seqlen_q,
+            step=seqlen_q,
+            dtype=torch.int32,
+            device=q_unpad.device,
+        )
+        max_seqlen_q = seqlen_q
+
+        output_pad_fn = lambda output_unpad: rearrange(  # noqa: E731
+            output_unpad, "(b s) h d -> b s h d", b=batch_size
+        )
+
+    if key_padding_mask is not None:
+        k_unpad, _, cu_seqlens_k, max_seqlen_k = unpad_input(k, key_padding_mask)
+        v_unpad, _, _, _ = unpad_input(v, key_padding_mask)
+    else:
+        k_unpad = rearrange(k, "b s h d -> (b s) h d")
+        v_unpad = rearrange(v, "b s h d -> (b s) h d")
+        cu_seqlens_k = torch.arange(
+            0,
+            (batch_size + 1) * seqlen_k,
+            step=seqlen_k,
+            dtype=torch.int32,
+            device=k_unpad.device,
+        )
+        max_seqlen_k = seqlen_k
+
+    if qkvpacked:
+        assert nheads == nheads_k
+        qkv_unpad = torch.stack([q_unpad, k_unpad, v_unpad], dim=1)
+        qkv = torch.stack([q, k, v], dim=2)
+        return (qkv_unpad, cu_seqlens_q, max_seqlen_q, qkv, output_pad_fn)
+
+    if kvpacked:
+        kv_unpad = torch.stack([k_unpad, v_unpad], dim=1)
+        kv = torch.stack([k, v], dim=2)
+        return (
+            q_unpad,
+            kv_unpad,
+            cu_seqlens_q,
+            cu_seqlens_k,
+            max_seqlen_q,
+            max_seqlen_k,
+            q,
+            kv,
+            output_pad_fn,
+        )
+
+    return (
+        q_unpad,
+        k_unpad,
+        v_unpad,
+        cu_seqlens_q,
+        cu_seqlens_k,
+        max_seqlen_q,
+        max_seqlen_k,
+        q,
+        k,
+        v,
+        output_pad_fn,
+    )
+
+
+def llama_model_forward(
+    self,
+    input_ids: torch.LongTensor = None,
+    attention_mask: Optional[torch.Tensor] = None,
+    position_ids: Optional[torch.LongTensor] = None,
+    past_key_values: Optional[List[torch.FloatTensor]] = None,
+    inputs_embeds: Optional[torch.FloatTensor] = None,
+    use_cache: Optional[bool] = None,
+    output_attentions: Optional[bool] = None,
+    output_hidden_states: Optional[bool] = None,
+    return_dict: Optional[bool] = None,
+) -> Union[Tuple, BaseModelOutputWithPast]:
+    output_attentions = (
+        output_attentions
+        if output_attentions is not None
+        else self.config.output_attentions
+    )
+    output_hidden_states = (
+        output_hidden_states
+        if output_hidden_states is not None
+        else self.config.output_hidden_states
+    )
+    use_cache = use_cache if use_cache is not None else self.config.use_cache
+
+    return_dict = (
+        return_dict if return_dict is not None else self.config.use_return_dict
+    )
+
+    # retrieve input_ids and inputs_embeds
+    if input_ids is not None and inputs_embeds is not None:
+        raise ValueError(
+            "You cannot specify both decoder_input_ids and decoder_inputs_embeds at the same time"
+        )
+    if input_ids is not None:
+        batch_size, seq_length = input_ids.shape
+    elif inputs_embeds is not None:
+        batch_size, seq_length, _ = inputs_embeds.shape
+    else:
+        raise ValueError(
+            "You have to specify either decoder_input_ids or decoder_inputs_embeds"
+        )
+
+    seq_length_with_past = seq_length
+    past_key_values_length = 0
+
+    if past_key_values is not None:
+        past_key_values_length = past_key_values[0][0].shape[2]
+        seq_length_with_past = seq_length_with_past + past_key_values_length
+
+    cu_seqlens = None
+    max_seqlen = None
+    if position_ids is None:
+        device = input_ids.device if input_ids is not None else inputs_embeds.device
+        position_ids = torch.arange(
+            past_key_values_length,
+            seq_length + past_key_values_length,
+            dtype=torch.long,
+            device=device,
+        )
+        position_ids = position_ids.unsqueeze(0).view(-1, seq_length)
+    else:
+        position_ids = position_ids.view(-1, seq_length).long()
+        cu_seqlens, max_seqlen = get_cu_seqlens_from_pos_ids(position_ids)
+        cu_seqlens = cu_seqlens.squeeze()
+
+    if inputs_embeds is None:
+        inputs_embeds = self.embed_tokens(input_ids)
+    # embed positions
+    if attention_mask is None:
+        attention_mask = torch.ones(
+            (batch_size, seq_length_with_past),
+            dtype=torch.bool,
+            device=inputs_embeds.device,
+        )
+    attention_mask = (
+        self._prepare_decoder_attention_mask(  # pylint: disable=protected-access
+            attention_mask,
+            (batch_size, seq_length),
+            inputs_embeds,
+            past_key_values_length,
+        )
+    )
+
+    hidden_states = inputs_embeds
+
+    if self.gradient_checkpointing and self.training:
+        if use_cache:
+            transformers.logger.warning_once(
+                "`use_cache=True` is incompatible with gradient checkpointing. Setting `use_cache=False`..."
+            )
+            use_cache = False
+
+    # decoder layers
+    all_hidden_states = () if output_hidden_states else None
+    all_self_attns = () if output_attentions else None
+    next_decoder_cache = () if use_cache else None
+
+    for idx, decoder_layer in enumerate(self.layers):
+        if output_hidden_states:
+            all_hidden_states += (hidden_states,)
+
+        past_key_value = past_key_values[idx] if past_key_values is not None else None
+
+        if self.gradient_checkpointing and self.training:
+
+            def create_custom_forward(module):
+                def custom_forward(*inputs):
+                    # None for past_key_value
+                    return module(*inputs)
+
+                return custom_forward
+
+            layer_outputs = torch.utils.checkpoint.checkpoint(
+                create_custom_forward(decoder_layer),
+                hidden_states,
+                attention_mask,
+                position_ids,
+                None,
+                output_attentions,
+                None,
+                cu_seqlens,
+                max_seqlen,
+            )
+        else:
+            layer_outputs = decoder_layer(
+                hidden_states,
+                attention_mask=attention_mask,
+                position_ids=position_ids,
+                past_key_value=past_key_value,
+                output_attentions=output_attentions,
+                use_cache=use_cache,
+                cu_seqlens=cu_seqlens,
+                max_seqlen=max_seqlen,
+            )
+
+        hidden_states = layer_outputs[0]
+
+        if use_cache:
+            next_decoder_cache += (layer_outputs[2 if output_attentions else 1],)
+
+        if output_attentions:
+            all_self_attns += (layer_outputs[1],)
+
+    hidden_states = self.norm(hidden_states)
+
+    # add hidden states from the last decoder layer
+    if output_hidden_states:
+        all_hidden_states += (hidden_states,)
+
+    next_cache = next_decoder_cache if use_cache else None
+    if not return_dict:
+        return tuple(
+            v
+            for v in [hidden_states, next_cache, all_hidden_states, all_self_attns]
+            if v is not None
+        )
+    return BaseModelOutputWithPast(
+        last_hidden_state=hidden_states,
+        past_key_values=next_cache,
+        hidden_states=all_hidden_states,
+        attentions=all_self_attns,
+    )
+
+
+class LlamaDecoderLayer(OriginalLlamaDecoderLayer):
+    """
+    patched version of LlamaDecoderLayer to pass through the precalculated cu_seqlens
+    """
+
+    def forward(
+        self,
+        hidden_states: torch.Tensor,
+        attention_mask: Optional[torch.Tensor] = None,
+        position_ids: Optional[torch.LongTensor] = None,
+        past_key_value: Optional[Tuple[torch.Tensor]] = None,
+        output_attentions: Optional[bool] = False,
+        use_cache: Optional[bool] = False,
+        cu_seqlens: Optional[torch.Tensor] = None,
+        max_seqlen: Optional[torch.Tensor] = None,
+    ) -> Tuple[
+        torch.FloatTensor, Optional[Tuple[torch.FloatTensor, torch.FloatTensor]]
+    ]:
+        """
+        Args:
+            hidden_states (`torch.FloatTensor`): input to the layer of shape `(batch, seq_len, embed_dim)`
+            attention_mask (`torch.FloatTensor`, *optional*): attention mask of size
+                `(batch, 1, tgt_len, src_len)` where padding elements are indicated by very large negative values.
+            output_attentions (`bool`, *optional*):
+                Whether or not to return the attentions tensors of all attention layers. See `attentions` under
+                returned tensors for more detail.
+            use_cache (`bool`, *optional*):
+                If set to `True`, `past_key_values` key value states are returned and can be used to speed up decoding
+                (see `past_key_values`).
+            past_key_value (`Tuple(torch.FloatTensor)`, *optional*): cached past key and value projection states
+            cu_seqlens (`torch.Tensor`, *optional*) cumulative sequence len when packing
+        """
+
+        residual = hidden_states
+
+        hidden_states = self.input_layernorm(hidden_states)
+
+        # Self Attention
+        hidden_states, self_attn_weights, present_key_value = self.self_attn(
+            hidden_states=hidden_states,
+            attention_mask=attention_mask,
+            position_ids=position_ids,
+            past_key_value=past_key_value,
+            output_attentions=output_attentions,
+            use_cache=use_cache,
+            cu_seqlens=cu_seqlens,
+            max_seqlen=max_seqlen,
+        )
+        hidden_states = residual + hidden_states
+
+        # Fully Connected
+        residual = hidden_states
+        hidden_states = self.post_attention_layernorm(hidden_states)
+        hidden_states = self.mlp(hidden_states)
+        hidden_states = residual + hidden_states
+
+        outputs = (hidden_states,)
+
+        if output_attentions:
+            outputs += (self_attn_weights,)
+
+        if use_cache:
+            outputs += (present_key_value,)
+
+        return outputs
--- a/src/axolotl/monkeypatch/llama_attn_hijack_sdp.py
+++ b/src/axolotl/monkeypatch/llama_attn_hijack_sdp.py
@@ -0,0 +1,140 @@
+"""
+Patched LlamaAttention to use torch.nn.functional.scaled_dot_product_attention
+"""
+
+import warnings
+from typing import Optional, Tuple
+
+import torch
+import torch.nn.functional as F
+import transformers.models.llama.modeling_llama
+from transformers.models.llama.modeling_llama import apply_rotary_pos_emb, repeat_kv
+
+
+def hijack_llama_sdp_attention():
+    transformers.models.llama.modeling_llama.LlamaAttention.forward = (
+        sdp_attention_forward
+    )
+
+
+def sdp_attention_forward(
+    self,
+    hidden_states: torch.Tensor,
+    attention_mask: Optional[torch.Tensor] = None,
+    position_ids: Optional[torch.LongTensor] = None,
+    past_key_value: Optional[Tuple[torch.Tensor]] = None,
+    output_attentions: bool = False,
+    use_cache: bool = False,
+) -> Tuple[torch.Tensor, Optional[torch.Tensor], Optional[Tuple[torch.Tensor]]]:
+    # pylint: disable=duplicate-code
+    bsz, q_len, _ = hidden_states.size()
+
+    if not hasattr(self, "pretraining_tp"):
+        self.pretraining_tp = 1
+
+    if self.pretraining_tp > 1:
+        key_value_slicing = (
+            self.num_key_value_heads * self.head_dim
+        ) // self.pretraining_tp
+        query_slices = self.q_proj.weight.split(
+            (self.num_heads * self.head_dim) // self.pretraining_tp, dim=0
+        )
+        key_slices = self.k_proj.weight.split(key_value_slicing, dim=0)
+        value_slices = self.v_proj.weight.split(key_value_slicing, dim=0)
+
+        query_states = [
+            F.linear(hidden_states, query_slices[i]) for i in range(self.pretraining_tp)
+        ]
+        query_states = torch.cat(query_states, dim=-1)
+
+        key_states = [
+            F.linear(hidden_states, key_slices[i]) for i in range(self.pretraining_tp)
+        ]
+        key_states = torch.cat(key_states, dim=-1)
+
+        value_states = [
+            F.linear(hidden_states, value_slices[i]) for i in range(self.pretraining_tp)
+        ]
+        value_states = torch.cat(value_states, dim=-1)
+
+    else:
+        query_states = self.q_proj(hidden_states)
+        key_states = self.k_proj(hidden_states)
+        value_states = self.v_proj(hidden_states)
+
+    query_states = query_states.view(
+        bsz, q_len, self.num_heads, self.head_dim
+    ).transpose(1, 2)
+    key_states = key_states.view(
+        bsz, q_len, self.num_key_value_heads, self.head_dim
+    ).transpose(1, 2)
+    value_states = value_states.view(
+        bsz, q_len, self.num_key_value_heads, self.head_dim
+    ).transpose(1, 2)
+    # [bsz, q_len, nh, hd]
+    # [bsz, nh, q_len, hd]
+
+    kv_seq_len = key_states.shape[-2]
+    if past_key_value is not None:
+        kv_seq_len += past_key_value[0].shape[-2]
+
+    cos, sin = self.rotary_emb(value_states, seq_len=kv_seq_len)
+    query_states, key_states = apply_rotary_pos_emb(
+        query_states, key_states, cos, sin, position_ids
+    )
+    # [bsz, nh, t, hd]
+
+    if past_key_value is not None:
+        # reuse k, v, self_attention
+        key_states = torch.cat([past_key_value[0], key_states], dim=2)
+        value_states = torch.cat([past_key_value[1], value_states], dim=2)
+
+    past_key_value = (key_states, value_states) if use_cache else None
+
+    # repeat k/v heads if n_kv_heads < n_heads
+    key_states = repeat_kv(key_states, self.num_key_value_groups)
+    value_states = repeat_kv(value_states, self.num_key_value_groups)
+
+    if output_attentions:
+        warnings.warn(
+            "Output attentions is not supported for patched `LlamaAttention`, returning `None` instead."
+        )
+
+    #
+    # sdp-attn start
+    #
+
+    with torch.backends.cuda.sdp_kernel():
+        attn_output = torch.nn.functional.scaled_dot_product_attention(
+            query_states,
+            key_states,
+            value_states,
+            attn_mask=attention_mask,
+            is_causal=False,
+        )
+
+    if attn_output.size() != (bsz, self.num_heads, q_len, self.head_dim):
+        raise ValueError(
+            f"`attn_output` should be of size {(bsz, self.num_heads, q_len, self.head_dim)}, but is"
+            f" {attn_output.size()}"
+        )
+    attn_output = attn_output.transpose(1, 2)
+    attn_output = attn_output.reshape(bsz, q_len, self.hidden_size)
+
+    #
+    # sdp-attn end
+    #
+
+    if self.pretraining_tp > 1:
+        attn_output = attn_output.split(self.hidden_size // self.pretraining_tp, dim=2)
+        o_proj_slices = self.o_proj.weight.split(
+            self.hidden_size // self.pretraining_tp, dim=1
+        )
+        attn_output = sum(
+            F.linear(attn_output[i], o_proj_slices[i])
+            for i in range(self.pretraining_tp)
+        )
+    else:
+        attn_output = self.o_proj(attn_output)
+
+    return attn_output, None, past_key_value
--- a/src/axolotl/monkeypatch/llama_attn_hijack_xformers.py
+++ b/src/axolotl/monkeypatch/llama_attn_hijack_xformers.py
@@ -3,13 +3,13 @@ Directly copied the code from https://raw.githubusercontent.com/oobabooga/text-g
 """

 import logging
-import math
+import warnings
 from typing import Optional, Tuple

 import torch
 import torch.nn.functional as F
 import transformers.models.llama.modeling_llama
-from torch import nn
+from transformers.models.llama.modeling_llama import apply_rotary_pos_emb, repeat_kv

 try:
    import xformers.ops
@@ -21,12 +21,6 @@ def hijack_llama_attention():
    transformers.models.llama.modeling_llama.LlamaAttention.forward = xformers_forward


-def hijack_llama_sdp_attention():
-    transformers.models.llama.modeling_llama.LlamaAttention.forward = (
-        sdp_attention_forward
-    )
-
-
 def xformers_forward(
    self,
    hidden_states: torch.Tensor,
@@ -81,15 +75,15 @@ def xformers_forward(
    value_states = value_states.view(
        bsz, q_len, self.num_key_value_heads, self.head_dim
    ).transpose(1, 2)
+    # [bsz, q_len, nh, hd]
+    # [bsz, nh, q_len, hd]

    kv_seq_len = key_states.shape[-2]
    if past_key_value is not None:
        kv_seq_len += past_key_value[0].shape[-2]
+
    cos, sin = self.rotary_emb(value_states, seq_len=kv_seq_len)
-    (
-        query_states,
-        key_states,
-    ) = transformers.models.llama.modeling_llama.apply_rotary_pos_emb(
+    query_states, key_states = apply_rotary_pos_emb(
        query_states, key_states, cos, sin, position_ids
    )
    # [bsz, nh, t, hd]
@@ -102,74 +96,50 @@ def xformers_forward(
    past_key_value = (key_states, value_states) if use_cache else None

    # repeat k/v heads if n_kv_heads < n_heads
-    key_states = transformers.models.llama.modeling_llama.repeat_kv(
-        key_states, self.num_key_value_groups
-    )
-    value_states = transformers.models.llama.modeling_llama.repeat_kv(
-        value_states, self.num_key_value_groups
-    )
+    key_states = repeat_kv(key_states, self.num_key_value_groups)
+    value_states = repeat_kv(value_states, self.num_key_value_groups)

-    # We only apply xformers optimizations if we don't need to output the whole attention matrix
-    if not output_attentions:
-        query_states = query_states.transpose(1, 2)
-        key_states = key_states.transpose(1, 2)
-        value_states = value_states.transpose(1, 2)
+    if output_attentions:
+        warnings.warn(
+            "Output attentions is not supported for patched `LlamaAttention`, returning `None` instead."
+        )

-        # This is a nasty hack. We know attention_mask in transformers is either LowerTriangular or all Zeros.
-        # We therefore check if one element in the upper triangular portion is zero. If it is, then the mask is all zeros.
-        if attention_mask is None or attention_mask[0, 0, 0, 1] == 0:
-            # input and output should be of form (bsz, q_len, num_heads, head_dim)
-            attn_output = xformers.ops.memory_efficient_attention(
-                query_states, key_states, value_states, attn_bias=None
-            )
-        else:
-            # input and output should be of form (bsz, q_len, num_heads, head_dim)
-            attn_output = xformers.ops.memory_efficient_attention(
-                query_states,
-                key_states,
-                value_states,
-                # attn_bias=attention_mask,
-                attn_bias=xformers.ops.LowerTriangularMask(),
-            )
-        attn_weights = None
+    #
+    # xformers-attn start
+    #
+
+    query_states = query_states.transpose(1, 2)
+    key_states = key_states.transpose(1, 2)
+    value_states = value_states.transpose(1, 2)
+
+    # This is a nasty hack. We know attention_mask in transformers is either LowerTriangular or all Zeros.
+    # We therefore check if one element in the upper triangular portion is zero. If it is, then the mask is all zeros.
+    if attention_mask is None or attention_mask[0, 0, 0, 1] == 0:
+        # input and output should be of form (bsz, q_len, num_heads, head_dim)
+        attn_output = xformers.ops.memory_efficient_attention(
+            query_states, key_states, value_states, attn_bias=None
+        )
    else:
-        attn_weights = torch.matmul(
-            query_states, key_states.transpose(2, 3)
-        ) / math.sqrt(self.head_dim)
-
-        if attn_weights.size() != (bsz, self.num_heads, q_len, kv_seq_len):
-            raise ValueError(
-                f"Attention weights should be of size {(bsz * self.num_heads, q_len, kv_seq_len)}, but is"
-                f" {attn_weights.size()}"
-            )
-
-        if attention_mask is not None:
-            if attention_mask.size() != (bsz, 1, q_len, kv_seq_len):
-                raise ValueError(
-                    f"Attention mask should be of size {(bsz, 1, q_len, kv_seq_len)}, but is {attention_mask.size()}"
-                )
-            attn_weights = attn_weights + attention_mask
-            attn_weights = torch.max(
-                attn_weights, torch.tensor(torch.finfo(attn_weights.dtype).min)
-            )
-
-        # upcast attention to fp32
-        attn_weights = nn.functional.softmax(
-            attn_weights, dim=-1, dtype=torch.float32
-        ).to(query_states.dtype)
-        attn_output = torch.matmul(attn_weights, value_states)
-
-        if attn_output.size() != (bsz, self.num_heads, q_len, self.head_dim):
-            raise ValueError(
-                f"`attn_output` should be of size {(bsz, self.num_heads, q_len, self.head_dim)}, but is"
-                f" {attn_output.size()}"
-            )
-
-        attn_output = attn_output.transpose(1, 2).contiguous()
-        # end x-formers vs. not x-formers if-else block
+        # input and output should be of form (bsz, q_len, num_heads, head_dim)
+        attn_output = xformers.ops.memory_efficient_attention(
+            query_states,
+            key_states,
+            value_states,
+            # attn_bias=attention_mask,
+            attn_bias=xformers.ops.LowerTriangularMask(),
+        )

+    if attn_output.size() != (bsz, q_len, self.num_heads, self.head_dim):
+        raise ValueError(
+            f"`attn_output` should be of size {(bsz, q_len, self.num_heads, self.head_dim)}, but is"
+            f" {attn_output.size()}"
+        )
    attn_output = attn_output.reshape(bsz, q_len, self.hidden_size)

+    #
+    # xformers-attn end
+    #
+
    if self.pretraining_tp > 1:
        attn_output = attn_output.split(self.hidden_size // self.pretraining_tp, dim=2)
        o_proj_slices = self.o_proj.weight.split(
@@ -182,103 +152,4 @@ def xformers_forward(
    else:
        attn_output = self.o_proj(attn_output)

-    return attn_output, attn_weights, past_key_value
-
-
-def sdp_attention_forward(
-    self,
-    hidden_states: torch.Tensor,
-    attention_mask: Optional[torch.Tensor] = None,
-    position_ids: Optional[torch.LongTensor] = None,
-    past_key_value: Optional[Tuple[torch.Tensor]] = None,
-    output_attentions: bool = False,
-    use_cache: bool = False,
-) -> Tuple[torch.Tensor, Optional[torch.Tensor], Optional[Tuple[torch.Tensor]]]:
-    # pylint: disable=duplicate-code
-    bsz, q_len, _ = hidden_states.size()
-
-    query_states = (
-        self.q_proj(hidden_states)
-        .view(bsz, q_len, self.num_heads, self.head_dim)
-        .transpose(1, 2)
-    )
-    key_states = (
-        self.k_proj(hidden_states)
-        .view(bsz, q_len, self.num_heads, self.head_dim)
-        .transpose(1, 2)
-    )
-    value_states = (
-        self.v_proj(hidden_states)
-        .view(bsz, q_len, self.num_heads, self.head_dim)
-        .transpose(1, 2)
-    )
-
-    kv_seq_len = key_states.shape[-2]
-    if past_key_value is not None:
-        kv_seq_len += past_key_value[0].shape[-2]
-    cos, sin = self.rotary_emb(value_states, seq_len=kv_seq_len)
-    (
-        query_states,
-        key_states,
-    ) = transformers.models.llama.modeling_llama.apply_rotary_pos_emb(
-        query_states, key_states, cos, sin, position_ids
-    )
-    # [bsz, nh, t, hd]
-
-    if past_key_value is not None:
-        # reuse k, v, self_attention
-        key_states = torch.cat([past_key_value[0], key_states], dim=2)
-        value_states = torch.cat([past_key_value[1], value_states], dim=2)
-
-    past_key_value = (key_states, value_states) if use_cache else None
-
-    # We only apply sdp attention if we don't need to output the whole attention matrix
-    if not output_attentions:
-        with torch.backends.cuda.sdp_kernel():
-            attn_output = torch.nn.functional.scaled_dot_product_attention(
-                query_states,
-                key_states,
-                value_states,
-                attn_mask=attention_mask,
-                is_causal=False,
-            )
-            attn_weights = None
-    else:
-        attn_weights = torch.matmul(
-            query_states, key_states.transpose(2, 3)
-        ) / math.sqrt(self.head_dim)
-
-        if attn_weights.size() != (bsz, self.num_heads, q_len, kv_seq_len):
-            raise ValueError(
-                f"Attention weights should be of size {(bsz * self.num_heads, q_len, kv_seq_len)}, but is"
-                f" {attn_weights.size()}"
-            )
-
-        if attention_mask is not None:
-            if attention_mask.size() != (bsz, 1, q_len, kv_seq_len):
-                raise ValueError(
-                    f"Attention mask should be of size {(bsz, 1, q_len, kv_seq_len)}, but is {attention_mask.size()}"
-                )
-            attn_weights = attn_weights + attention_mask
-            attn_weights = torch.max(
-                attn_weights, torch.tensor(torch.finfo(attn_weights.dtype).min)
-            )
-
-        # upcast attention to fp32
-        attn_weights = nn.functional.softmax(
-            attn_weights, dim=-1, dtype=torch.float32
-        ).to(query_states.dtype)
-        attn_output = torch.matmul(attn_weights, value_states)
-
-        if attn_output.size() != (bsz, self.num_heads, q_len, self.head_dim):
-            raise ValueError(
-                f"`attn_output` should be of size {(bsz, self.num_heads, q_len, self.head_dim)}, but is"
-                f" {attn_output.size()}"
-            )
-
-    attn_output = attn_output.transpose(1, 2)
-    attn_output = attn_output.reshape(bsz, q_len, self.hidden_size)
-
-    attn_output = self.o_proj(attn_output)
-
-    return attn_output, attn_weights, past_key_value
+    return attn_output, None, past_key_value
--- a/src/axolotl/prompt_strategies/alpaca_w_system.py
+++ b/src/axolotl/prompt_strategies/alpaca_w_system.py
@@ -92,8 +92,9 @@ class OpenOrcaSystemDataPrompter(SystemDataPrompter):
    def match_prompt_style(self):
        # pylint: disable=duplicate-code
        if self.prompt_style == PromptStyle.INSTRUCT.value:
-            self.turn_format = "### User:\n{instruction}\n\n### Additional Context:\n{input}\n\n### Assistant:\n"
-            self.turn_no_input_format = "### User:\n{instruction}\n\n### Assistant:\n"
+            self.turn_format = "### Human:\n{instruction}\n### Additional Context:\n{input}\n### Assistant:\n"
+            self.turn_no_input_format = "### Human:\n{instruction}\n### Assistant:\n"
+            self.system_format = "### System:\n{system}\n"
        if self.prompt_style == PromptStyle.CHAT.value:
            self.turn_format = "USER: {instruction}\n{input}\nASSISTANT:"
            self.turn_no_input_format = "USER: {instruction}\nASSISTANT:"
--- a/src/axolotl/prompt_tokenizers.py
+++ b/src/axolotl/prompt_tokenizers.py
@@ -74,8 +74,11 @@ class PromptTokenizingStrategy(abc.ABC):
            padding=False,
            return_tensors=None,
        )
+        if len(result["input_ids"]) == 0:
+            LOG.warning("Tokenizer result is empty. You may want to audit your dataset")
        if (
-            result["input_ids"][-1] != self.tokenizer.eos_token_id
+            len(result["input_ids"]) > 0
+            and result["input_ids"][-1] != self.tokenizer.eos_token_id
            and len(result["input_ids"]) < self.sequence_len
            and add_eos_token
        ):
--- a/src/axolotl/prompters.py
+++ b/src/axolotl/prompters.py
@@ -312,7 +312,9 @@ class ShareGPTPrompter:  # pylint: disable=too-few-public-methods
        if len(source) < 2:
            # If there isn't a back and forth conversation, ignore it
            # also happens on the data splitting leaving empty conversations
-            raise IndexError
+            raise IndexError(
+                f"A conversation entry has less than 2 messages :\n{source}"
+            )

        conv = self._conversation.copy()
        roles = {"human": conv.roles[0], "gpt": conv.roles[1]}
--- a/src/axolotl/utils/bench.py
+++ b/src/axolotl/utils/bench.py
@@ -4,13 +4,23 @@ import pynvml
 import torch


-def gpu_memory_usage(device):
+def gpu_memory_usage(device=0):
+    return torch.cuda.memory_allocated(device) / 1024.0**3
+
+
+def gpu_memory_usage_all(device=0):
+    usage = torch.cuda.memory_allocated(device) / 1024.0**3
+    reserved = torch.cuda.memory_reserved(device) / 1024.0**3
+    smi = gpu_memory_usage_smi(device)
+    return usage, reserved - usage, max(0, smi - reserved)
+
+
+def gpu_memory_usage_smi(device=0):
    if isinstance(device, torch.device):
        device = device.index
    if isinstance(device, str) and device.startswith("cuda:"):
        device = int(device[5:])

-    # NB torch.cuda.memory_usage returns zero so we use lower level api
    pynvml.nvmlInit()
    handle = pynvml.nvmlDeviceGetHandleByIndex(device)
    info = pynvml.nvmlDeviceGetMemoryInfo(handle)
@@ -18,6 +28,16 @@ def gpu_memory_usage(device):


 def log_gpu_memory_usage(log, msg, device):
+    if not torch.cuda.is_available():
+        return (0, 0, 0)
+
+    usage, cache, misc = gpu_memory_usage_all(device)
+    extras = []
+    if cache > 0:
+        extras.append(f"+{cache:.03f}GB cache")
+    if misc > 0:
+        extras.append(f"+{misc:.03f}GB misc")
    log.info(
-        f"GPU memory usage {msg}: {gpu_memory_usage(device):.03f} GB", stacklevel=2
+        f"GPU memory usage {msg}: {usage:.03f}GB ({', '.join(extras)})", stacklevel=2
    )
+    return usage, cache, misc
--- a/src/axolotl/utils/callbacks.py
+++ b/src/axolotl/utils/callbacks.py
@@ -74,10 +74,10 @@ class SaveBetterTransformerModelCallback(
        return control


-class PrintGPUStatsCallback(
+class GPUStatsCallback(
    TrainerCallback
 ):  # pylint: disable=too-few-public-methods disable=unused-argument
-    """Callback to print GPU utilization"""
+    """Callback to track GPU utilization"""

    def __init__(self, cfg):
        self.cfg = cfg
@@ -90,7 +90,7 @@ class PrintGPUStatsCallback(
        control: TrainerControl,
        **kwargs,
    ):
-        if not self.logged:
+        if not self.logged and state.global_step > 1:
            log_gpu_memory_usage(LOG, "while training", self.cfg.device)
            self.logged = True
        return control
--- a/src/axolotl/utils/validation.py
+++ b/src/axolotl/utils/validation.py
@@ -1,12 +1,70 @@
-"""Module for validating config files"""
+"""Module for working with config dicts"""

 import logging
+import os

 import torch

+from axolotl.utils.bench import log_gpu_memory_usage
+
 LOG = logging.getLogger("axolotl")


+def choose_device(cfg):
+    def get_device():
+        try:
+            if torch.cuda.is_available():
+                return f"cuda:{cfg.local_rank}"
+
+            if torch.backends.mps.is_available():
+                return "mps"
+
+            raise SystemError("No CUDA/mps device found")
+        except Exception:  # pylint: disable=broad-exception-caught
+            return "cpu"
+
+    cfg.device = get_device()
+    if cfg.device_map != "auto":
+        if cfg.device.startswith("cuda"):
+            cfg.device_map = {"": cfg.local_rank}
+        else:
+            cfg.device_map = {"": cfg.device}
+
+    # in `accelerate launch`, we need to not pass through any device map and let
+    # accelerate figure out which parts of the model to put on which gpu
+    accelerate_vars = [var for var in os.environ if var.startswith("ACCELERATE_USE_")]
+    if accelerate_vars:
+        cfg.device_map = None
+
+
+def normalize_config(cfg):
+    # setup some derived config / hyperparams
+    cfg.gradient_accumulation_steps = cfg.gradient_accumulation_steps or (
+        cfg.batch_size // cfg.micro_batch_size
+    )
+    cfg.batch_size = (
+        cfg.batch_size or cfg.micro_batch_size * cfg.gradient_accumulation_steps
+    )
+    cfg.world_size = int(os.environ.get("WORLD_SIZE", 1))
+    cfg.local_rank = int(os.environ.get("LOCAL_RANK", 0))
+    choose_device(cfg)
+    cfg.ddp = cfg.ddp if cfg.ddp is not None else cfg.world_size != 1
+    if cfg.ddp:
+        cfg.device_map = {"": int(os.environ.get("LOCAL_RANK", 0))}
+        cfg.batch_size = cfg.batch_size * cfg.world_size
+
+    if cfg.device == "mps":
+        cfg.load_in_8bit = False
+        cfg.tf32 = False
+        if cfg.bf16:
+            cfg.fp16 = True
+        cfg.bf16 = False
+    else:
+        torch.backends.cuda.matmul.allow_tf32 = cfg.tf32 or False
+
+    log_gpu_memory_usage(LOG, "baseline", cfg.device)
+
+
 def validate_config(cfg):
    if cfg.max_packed_sequence_len and cfg.sample_packing:
        raise ValueError(
@@ -89,7 +147,7 @@ def validate_config(cfg):
                "You should probably set bfloat16 or float16 to true to "
                "load the model in float16 for BetterTransformers"
            )
-        if int(torch.__version__.split(".")[0]) < 2:
+        if int(torch.__version__.split(".", maxsplit=1)[0]) < 2:
            LOG.warning("torch>=2.0.0 required")
            raise ValueError(
                f"flash_optimum for BetterTransformers may not be used with {torch.__version__}"
--- a/src/axolotl/utils/data.py
+++ b/src/axolotl/utils/data.py
@@ -41,9 +41,44 @@ from axolotl.prompters import (
    ShareGPTPrompter,
    SummarizeTLDRPrompter,
 )
-from axolotl.utils.distributed import barrier, is_main_process
+from axolotl.utils.distributed import is_main_process, zero_first
+from axolotl.utils.trainer import (
+    calculate_total_num_steps,
+    process_datasets_for_packing,
+)

 LOG = logging.getLogger("axolotl")
+DEFAULT_DATASET_PREPARED_PATH = "last_run_prepared"
+
+
+def prepare_dataset(cfg, tokenizer):
+    if not cfg.pretraining_dataset:
+        train_dataset, eval_dataset = load_prepare_datasets(
+            tokenizer, cfg, DEFAULT_DATASET_PREPARED_PATH
+        )
+    else:
+        train_dataset = load_pretraining_dataset(
+            cfg.pretraining_dataset,
+            tokenizer,
+            max_tokens=cfg.sequence_len,
+            seed=cfg.seed or 42,
+        )
+        # https://discuss.huggingface.co/t/how-to-use-huggingface-trainer-streaming-datasets-without-wrapping-it-with-torchdatas-iterablewrapper/25230
+        train_dataset = train_dataset.with_format("torch")
+        eval_dataset = None
+
+    with zero_first(is_main_process()):
+        train_dataset, eval_dataset = process_datasets_for_packing(
+            cfg, train_dataset, eval_dataset
+        )
+    if cfg.max_steps:
+        total_num_steps = min(
+            calculate_total_num_steps(cfg, train_dataset, tokenizer), cfg.max_steps
+        )
+        LOG.info(f"Maximum number of steps set at {total_num_steps}")
+    else:
+        total_num_steps = calculate_total_num_steps(cfg, train_dataset, tokenizer)
+    return train_dataset, eval_dataset, total_num_steps


 def load_tokenized_prepared_datasets(
@@ -440,7 +475,7 @@ def load_prepare_datasets(
            to_hash_test.encode(), usedforsecurity=False
        ).hexdigest()

-        if is_main_process():
+        with zero_first(is_main_process()):
            dataset = dataset.train_test_split(
                test_size=cfg.val_set_size,
                shuffle=False,
@@ -448,16 +483,6 @@ def load_prepare_datasets(
                train_new_fingerprint=train_fingerprint,
                test_new_fingerprint=test_fingerprint,
            )
-        barrier()
-        if not is_main_process():
-            dataset = dataset.train_test_split(
-                test_size=cfg.val_set_size,
-                shuffle=False,
-                seed=cfg.seed or 42,
-                train_new_fingerprint=train_fingerprint,
-                test_new_fingerprint=test_fingerprint,
-            )
-        barrier()

        train_dataset = dataset["train"]
        eval_dataset = dataset["test"]
--- a/src/axolotl/utils/dict.py
+++ b/src/axolotl/utils/dict.py
@@ -10,3 +10,6 @@ class DictDefault(Dict):

    def __missing__(self, key):
        return None
+
+    def __or__(self, other):
+        return DictDefault(super().__or__(other))
--- a/src/axolotl/utils/distributed.py
+++ b/src/axolotl/utils/distributed.py
@@ -1,6 +1,8 @@
 """
 utility helpers for distributed checks
 """
+from contextlib import contextmanager
+
 import torch.distributed as dist
 from accelerate import Accelerator

@@ -39,3 +41,15 @@ def is_main_process():
    if not is_distributed():
        return True
    return dist.get_rank() == 0
+
+
+@contextmanager
+def zero_first(is_main):
+    """
+    runs the wrapped context so that rank 0 runs first before other ranks
+    """
+    if not is_main:  # other ranks wait first
+        barrier()
+    yield
+    if is_main:  # then rank 0 waits after it has run the context
+        barrier()
--- a/src/axolotl/utils/models.py
+++ b/src/axolotl/utils/models.py
@@ -32,37 +32,27 @@ if TYPE_CHECKING:
    from axolotl.utils.dict import DictDefault  # noqa: F401


-def load_tokenizer(
-    tokenizer_config,
-    tokenizer_type,
-    cfg,
-):
+def load_tokenizer(cfg):
    tokenizer_kwargs = {}
    use_fast = True  # this is the default
+
    if cfg.tokenizer_use_fast is not None:
        use_fast = cfg.tokenizer_use_fast
    if cfg.tokenizer_legacy is not None:
        # True is the default w/ https://github.com/huggingface/transformers/pull/25224
        tokenizer_kwargs["legacy"] = cfg.tokenizer_legacy
-    if tokenizer_type:
-        tokenizer = getattr(transformers, tokenizer_type).from_pretrained(
-            tokenizer_config,
-            trust_remote_code=cfg.trust_remote_code or False,
-            use_fast=use_fast,
-            **tokenizer_kwargs,
-        )
-    else:
-        tokenizer = AutoTokenizer.from_pretrained(
-            tokenizer_config,
-            trust_remote_code=cfg.trust_remote_code or False,
-            use_fast=use_fast,
-            **tokenizer_kwargs,
-        )

-    LOG.debug(f"EOS: {tokenizer.eos_token_id} / {tokenizer.eos_token}")
-    LOG.debug(f"BOS: {tokenizer.bos_token_id} / {tokenizer.bos_token}")
-    LOG.debug(f"PAD: {tokenizer.pad_token_id} / {tokenizer.pad_token}")
-    LOG.debug(f"UNK: {tokenizer.unk_token_id} / {tokenizer.unk_token}")
+    tokenizer_cls = AutoTokenizer
+    if cfg.tokenizer_type:
+        tokenizer_cls = getattr(transformers, cfg.tokenizer_type)
+
+    tokenizer_config = cfg.tokenizer_config or cfg.base_model_config
+    tokenizer = tokenizer_cls.from_pretrained(
+        tokenizer_config,
+        trust_remote_code=cfg.trust_remote_code or False,
+        use_fast=use_fast,
+        **tokenizer_kwargs,
+    )

    if tokenizer.__class__.__name__ in [
        "LlamaTokenizer",
@@ -70,6 +60,11 @@ def load_tokenizer(
    ]:
        tokenizer.pad_token = LLAMA_DEFAULT_PAD_TOKEN

+    LOG.debug(f"EOS: {tokenizer.eos_token_id} / {tokenizer.eos_token}")
+    LOG.debug(f"BOS: {tokenizer.bos_token_id} / {tokenizer.bos_token}")
+    LOG.debug(f"PAD: {tokenizer.pad_token_id} / {tokenizer.pad_token}")
+    LOG.debug(f"UNK: {tokenizer.unk_token_id} / {tokenizer.unk_token}")
+
    if tokenizer.__class__.__name__ == "GPTNeoXTokenizerFast":
        tokenizer.add_special_tokens({"pad_token": "[PAD]"})
        os.environ["TOKENIZERS_PARALLELISM"] = "false"
@@ -92,7 +87,6 @@ def load_model(
    base_model = cfg.base_model
    base_model_config = cfg.base_model_config
    model_type = cfg.model_type
-    adapter = cfg.adapter

    # TODO refactor as a kwarg
    load_in_8bit = cfg.load_in_8bit
@@ -109,7 +103,7 @@ def load_model(
            )

            LOG.info("patching with flash attention")
-            replace_llama_attn_with_flash_attn()
+            replace_llama_attn_with_flash_attn(packed=cfg.sample_packing)
    elif cfg.is_llama_derived_model and cfg.xformers_attention:
        from axolotl.monkeypatch.llama_attn_hijack_xformers import (
            hijack_llama_attention,
@@ -118,9 +112,7 @@ def load_model(
        LOG.info("patching with xformers attention")
        hijack_llama_attention()
    elif cfg.is_llama_derived_model and cfg.sdp_attention:
-        from axolotl.monkeypatch.llama_attn_hijack_xformers import (
-            hijack_llama_sdp_attention,
-        )
+        from axolotl.monkeypatch.llama_attn_hijack_sdp import hijack_llama_sdp_attention

        LOG.info("patching with sdp attention")
        hijack_llama_sdp_attention()
@@ -144,8 +136,10 @@ def load_model(
        LOG.info("patching with xpos rope")
        replace_llama_rope_with_xpos_rope()

-    if cfg.is_llama_derived_model and (
-        cfg.max_packed_sequence_len or cfg.sample_packing
+    if (
+        cfg.is_llama_derived_model
+        and (cfg.max_packed_sequence_len or cfg.sample_packing)
+        and not cfg.inference
    ):
        from axolotl.monkeypatch.llama_expand_mask import hijack_expand_mask

@@ -235,12 +229,17 @@ def load_model(
        elif cfg.is_llama_derived_model and not cfg.trust_remote_code:
            from transformers import LlamaForCausalLM

+            config_kwargs = {}
+            if cfg.rope_scaling:
+                config_kwargs["rope_scaling"] = cfg.rope_scaling
            config = LlamaConfig.from_pretrained(
-                base_model_config, rope_scaling=cfg.rope_scaling
+                base_model_config,
+                **config_kwargs,
            )
            model = LlamaForCausalLM.from_pretrained(
                base_model,
                config=config,
+                device_map=cfg.device_map,
                load_in_8bit=cfg.load_in_8bit and cfg.adapter is not None,
                load_in_4bit=cfg.load_in_4bit and cfg.adapter is not None,
                torch_dtype=torch_dtype,
@@ -275,6 +274,7 @@ def load_model(
        elif model_type and not cfg.trust_remote_code:
            model = getattr(transformers, model_type).from_pretrained(
                base_model,
+                device_map=cfg.device_map,
                load_in_8bit=cfg.load_in_8bit and cfg.adapter is not None,
                load_in_4bit=cfg.load_in_4bit and cfg.adapter is not None,
                torch_dtype=torch_dtype,
@@ -305,6 +305,7 @@ def load_model(
            model = AutoModelForCausalLM.from_pretrained(
                base_model,
                config=config,
+                device_map=cfg.device_map,
                load_in_8bit=cfg.load_in_8bit and cfg.adapter is not None,
                load_in_4bit=cfg.load_in_4bit and cfg.adapter is not None,
                torch_dtype=torch_dtype,
@@ -318,6 +319,7 @@ def load_model(
        LOG.exception(err)
        model = AutoModelForCausalLM.from_pretrained(
            base_model,
+            device_map=cfg.device_map,
            load_in_8bit=cfg.load_in_8bit and cfg.adapter is not None,
            load_in_4bit=cfg.load_in_4bit and cfg.adapter is not None,
            torch_dtype=torch_dtype,
@@ -364,7 +366,7 @@ def load_model(
                    if hasattr(module, "weight"):
                        module.to(torch_dtype)

-    model, lora_config = load_adapter(model, cfg, adapter)
+    model, lora_config = load_adapter(model, cfg, cfg.adapter)

    if cfg.ddp and not load_in_8bit:
        model.to(f"cuda:{cfg.local_rank}")
@@ -381,9 +383,6 @@ def load_model(
                module.scales = module.scales.half()
                module.bias = module.bias.half()

-    if model.device.type == "cuda":
-        log_gpu_memory_usage(LOG, "after adapters", model.device)
-
    if (
        torch.cuda.device_count() > 1
        and int(os.getenv("WORLD_SIZE", "1")) > 1
@@ -406,6 +405,9 @@ def load_model(
    if cfg.flash_optimum:
        model = BetterTransformer.transform(model)

+    if cfg.adapter is not None:
+        log_gpu_memory_usage(LOG, "after adapters", model.device)
+
    # TODO resume_from_checkpoint handling
    return model, lora_config

--- a/src/axolotl/utils/trainer.py
+++ b/src/axolotl/utils/trainer.py
@@ -22,7 +22,7 @@ from transformers import EarlyStoppingCallback, Trainer, TrainingArguments
 from transformers.trainer_pt_utils import get_parameter_names

 from axolotl.utils.callbacks import (
-    PrintGPUStatsCallback,
+    GPUStatsCallback,
    SaveBetterTransformerModelCallback,
    SavePeftModelCallback,
 )
@@ -364,6 +364,9 @@ def setup_fsdp_envs(cfg):
 def setup_trainer(cfg, train_dataset, eval_dataset, model, tokenizer, total_num_steps):
    if cfg.fsdp:
        setup_fsdp_envs(cfg)
+    elif cfg.deepspeed:
+        os.environ["ACCELERATE_USE_DEEPSPEED"] = "true"
+
    warmup_steps = (
        cfg.warmup_steps
        if cfg.warmup_steps is not None
@@ -411,21 +414,13 @@ def setup_trainer(cfg, train_dataset, eval_dataset, model, tokenizer, total_num_
        if cfg.fsdp_config:
            training_arguments_kwargs["fsdp_config"] = dict(cfg.fsdp_config)

+    # deepspeed
+    if cfg.deepspeed:
+        training_arguments_kwargs["deepspeed"] = cfg.deepspeed
+
    if cfg.lr_quadratic_warmup is not None:
        training_arguments_kwargs["lr_quadratic_warmup"] = cfg.lr_quadratic_warmup

-    # deepspeed
-    if (
-        os.environ.get("ACCELERATE_USE_DEEPSPEED") == "true"
-        and torch.cuda.device_count() > 1
-    ):
-        if cfg.deepspeed:
-            training_arguments_kwargs["deepspeed"] = cfg.deepspeed
-        else:
-            # make a guess here
-            # TODO search Path("./") for one
-            training_arguments_kwargs["deepspeed"] = "./ds_config.json"
-
    if cfg.adam_beta1:
        training_arguments_kwargs["adam_beta1"] = cfg.adam_beta1
    if cfg.adam_beta2:
@@ -440,6 +435,9 @@ def setup_trainer(cfg, train_dataset, eval_dataset, model, tokenizer, total_num_
        training_arguments_kwargs["push_to_hub"] = True
        training_arguments_kwargs["hub_private_repo"] = True

+        if cfg.hub_strategy:
+            training_arguments_kwargs["hub_strategy"] = cfg.hub_strategy
+
    if cfg.save_safetensors:
        training_arguments_kwargs["save_safetensors"] = cfg.save_safetensors

@@ -448,8 +446,17 @@ def setup_trainer(cfg, train_dataset, eval_dataset, model, tokenizer, total_num_
            "sample_packing_efficiency"
        ] = cfg.sample_packing_eff_est

+    if cfg.val_set_size == 0:
+        training_arguments_kwargs["evaluation_strategy"] = "no"
+    elif cfg.eval_steps:
+        training_arguments_kwargs["evaluation_strategy"] = "steps"
+        training_arguments_kwargs["eval_steps"] = cfg.eval_steps
+    else:
+        # we have an eval set, but no steps defined, use epoch
+        training_arguments_kwargs["evaluation_strategy"] = "epoch"
+
    training_args = AxolotlTrainingArguments(  # pylint: disable=unexpected-keyword-arg
-        # max_steps=total_num_steps,  # this is helpful in case we don't actually know total # of steps
+        max_steps=total_num_steps if cfg.max_steps else -1,
        max_seq_length=cfg.sequence_len,
        per_device_train_batch_size=cfg.micro_batch_size,
        per_device_eval_batch_size=cfg.eval_batch_size
@@ -459,9 +466,7 @@ def setup_trainer(cfg, train_dataset, eval_dataset, model, tokenizer, total_num_
        eval_accumulation_steps=cfg.gradient_accumulation_steps,
        num_train_epochs=cfg.num_epochs,
        learning_rate=cfg.learning_rate,
-        evaluation_strategy="steps" if cfg.val_set_size > 0 else "no",
        save_strategy="steps" if cfg.save_steps else "epoch",
-        eval_steps=cfg.eval_steps if cfg.val_set_size > 0 else None,
        save_steps=cfg.save_steps,
        output_dir=cfg.output_dir,
        save_total_limit=cfg.save_total_limit if cfg.save_total_limit else 4,
@@ -555,7 +560,7 @@ def setup_trainer(cfg, train_dataset, eval_dataset, model, tokenizer, total_num_
        trainer_kwargs["optimizers"] = (optimizer, lr_scheduler)

    callbacks = []
-    callbacks.append(PrintGPUStatsCallback(cfg))
+    callbacks.append(GPUStatsCallback(cfg))
    # TODO on_save callback to sync checkpoints to GCP/AWS in background
    if cfg.early_stopping_patience:
        early_stop_cb = EarlyStoppingCallback(
--- a/tests/fixtures/conversation.tokenized_llama2chat.json
+++ b/tests/fixtures/conversation.tokenized_llama2chat.json
--- a/tests/test_dict.py
+++ b/tests/test_dict.py
@@ -72,6 +72,13 @@ class DictDefaultTest(unittest.TestCase):

        assert cfg.random_key is None, "DictDefault should return None for missing keys"

+    def test_dict_or(self):
+        cfg = DictDefault({}) | DictDefault({})
+
+        assert (
+            cfg.random_key is None
+        ), "DictDefault should return None for missing keys after | operation"
+
    def test_dict_nested_missingparentkey(self):
        """
        Due to subclassing Dict, DictDefault will error if we try to access a nested key whose parent key does not exist.
--- a/tests/test_tokenizers.py
+++ b/tests/test_tokenizers.py
@@ -13,17 +13,22 @@ class TestTokenizers(unittest.TestCase):
    """

    def test_default_use_fast(self):
-        cfg = DictDefault({})
-        tokenizer = load_tokenizer("huggyllama/llama-7b", None, cfg)
+        cfg = DictDefault(
+            {
+                "tokenizer_config": "huggyllama/llama-7b",
+            }
+        )
+        tokenizer = load_tokenizer(cfg)
        assert "Fast" in tokenizer.__class__.__name__

    def test_dont_use_fast(self):
        cfg = DictDefault(
            {
+                "tokenizer_config": "huggyllama/llama-7b",
                "tokenizer_use_fast": False,
            }
        )
-        tokenizer = load_tokenizer("huggyllama/llama-7b", None, cfg)
+        tokenizer = load_tokenizer(cfg)
        assert "Fast" not in tokenizer.__class__.__name__


--- a/tests/test_validation.py
+++ b/tests/test_validation.py
@@ -6,8 +6,8 @@ from typing import Optional

 import pytest

+from axolotl.utils.config import validate_config
 from axolotl.utils.dict import DictDefault
-from axolotl.utils.validation import validate_config


 class ValidationTest(unittest.TestCase):
Author	SHA1	Message	Date
Wing Lian	587dbbfc02	fix the patch to work properly and work with FSDP Some checks failed pre-commit / pre-commit (push) Has been cancelled Details PyTest / test (3.10) (push) Has been cancelled Details PyTest / test (3.9) (push) Has been cancelled Details	2023-08-17 20:33:56 -04:00
Wing Lian	6c306d9186	enable eval dataloader using multipack again	2023-08-17 18:17:31 -04:00
Wing Lian	7565fb9d63	update forwards so we only calculate cu_seqlens once	2023-08-17 18:17:31 -04:00
Wing Lian	a6b737d5ff	copy LlamaModel.forward and LlamaDecoderLayer.forward into monkeypatch	2023-08-17 18:17:30 -04:00
Wing Lian	cf95b57c0a	fix patch to check position ids and don't use multipack for evals	2023-08-17 18:17:30 -04:00
Aman Karmani	13f7efaf74	speed up flash-attn inference	2023-08-17 18:17:30 -04:00
Aman Karmani	d773384f74	update flash-attn patch for 70B/GQA and inference using helper from flash-attn tests	2023-08-17 18:17:30 -04:00
Aman Karmani	985dcbc051	sync xformers patch to follow shared format and be diffable	2023-08-17 18:17:30 -04:00
Aman Karmani	5d0b27e5a1	split sdp attn into its own patch	2023-08-17 18:17:30 -04:00
Wing Lian	8cace80175	fix fixture for new tokenizer handling in transformers (#428 )	2023-08-17 17:01:52 -04:00
Wing Lian	1b7e8604bb	fix orca prompts (#422 )	2023-08-16 11:21:03 -04:00
NanoCode012	3d1f203b62	Fix(docs): Remove gptq+lora and fix xformer compat list (#423 )	2023-08-16 13:56:48 +09:00
Wing Lian	d3d6fd6ae6	just resort to tags ans use main-latest (#424 )	2023-08-16 00:39:57 -04:00
NanoCode012	b7449a997f	Fix(template): Inform to place stack trace to Issue (#417 ) * Fix(template): Inform to place stack trace to Issue * Update following suggestions Co-authored-by: Wing Lian <wing.lian@gmail.com> --------- Co-authored-by: Wing Lian <wing.lian@gmail.com>	2023-08-16 11:55:48 +09:00
Wing Lian	5f80b3560b	use inputs for image rather than outputs for docker metadata (#420 )	2023-08-15 18:26:59 -04:00
Wing Lian	24959091d7	hopefully improve the README (#419 ) * hopefully improve the README * exitcode -9 help * table of contents * formatting	2023-08-15 15:30:53 -04:00
Wing Lian	7af816699e	tag with latest as well for axolotl-runpod (#418 ) * tag with latest as well for axolotl-runpod * no dev branch for now	2023-08-15 15:30:41 -04:00
mhenrichsen	f806e86a6e	Merge pull request #413 from mhenrichsen/chore/update-deepseed-config update path to align with fsdp example	2023-08-15 20:08:23 +02:00
NanoCode012	2b990eb628	Feat(doc): Add lr_quadratic_warmup to readme (#412 )	2023-08-16 02:55:48 +09:00
mhenrichsen	bd8cab49c9	update path to align with fsdp example	2023-08-15 19:51:58 +02:00
NanoCode012	c01015f33f	Fix(config): Update handling of deepspeed config (#404 ) * Fix(config): Update handling of deepspeed config * feat: auto set deepspeed env if deepspeed passed * fix: update new deepspeed instructions	2023-08-16 01:22:43 +09:00
NanoCode012	72fe3f8e3d	Fix(docs): Update flash attn requirements (#409 )	2023-08-15 22:40:52 +09:00
Wing Lian	47961fdb8b	update docs for tokenizer_legacy (#401 ) * update docs for tokenizer_legacy * add default info	2023-08-15 22:34:42 +09:00
NanoCode012	7ad37cb6d7	Fix(template): Remove iPhone/android from Issue template (#407 )	2023-08-15 22:32:51 +09:00
Wing Lian	29241cf1e4	Ax art (#405 ) * axolotl text art :D * only print art on rank0 * lint and pr feedback	2023-08-15 08:34:30 -04:00
lightningRalf	31db0ecce4	add templates, CoC and contributing guide (#126 ) * add templates, CoC and contributing guide * Update .github/SECURITY.md correct responsible person Co-authored-by: Wing Lian <wing.lian@gmail.com> * Update bug-report.yaml axolotl version switch with axolotl branch-commit * update CONTRIBUTING doc * update reporting link * linter fixes * chore: fix linter --------- Co-authored-by: Wing Lian <wing.lian@gmail.com> Co-authored-by: NanoCode012 <kevinvong@rocketmail.com>	2023-08-15 07:41:05 -04:00
Wing Lian	da10af03e9	fix eval steps and strategy (#403 )	2023-08-15 07:28:50 -04:00
Wing Lian	85cf4f8e2c	better handling of empty input ids when tokenizing (#395 ) * better handling of empty input ids when tokenizing * Add warning if tokenizer resulted in empty result * fix len comparison for linter	2023-08-15 01:09:59 -04:00
Aman Karmani	2e22404d2d	add utils.data.prepare_dataset	2023-08-14 21:28:29 -07:00
NanoCode012	be294fd605	Feat(doc): Add how to save by epochs (#396 )	2023-08-15 13:24:25 +09:00
Wing Lian	fc2d6be96d	use context manager to run things on rank0 before others (#397 )	2023-08-15 00:10:47 -04:00
Wing Lian	1687be6a35	don't use mask expansion for inference (#392 )	2023-08-14 20:52:54 -04:00
NanoCode012	41ecb451c2	Feat(doc): Add max_steps to readme (#389 )	2023-08-15 00:34:22 +09:00
Gabriel Puliatti	3c2ad00d07	Feat(config): add max steps (#387 )	2023-08-14 11:19:29 -04:00
florian peyron	5d48a10548	Added "epoch" evaluation_strategy (#388 )	2023-08-14 10:59:23 -04:00
NanoCode012	73a0b6ead5	Feat(config): Add hub_strategy (#386 )	2023-08-14 07:12:55 -04:00
florian peyron	63fdb5a7fb	Error msg for sharegpt if conv has less than 2 msg (#379 )	2023-08-14 17:40:40 +09:00
mhenrichsen	fdffef5940	new llama-2 default settings (#370 ) * new default settings * fix whitespace * rm max packed sequence length --------- Co-authored-by: Mads Henrichsen <mads@BrbartiendeMads.lan>	2023-08-14 17:39:09 +09:00
Wing Lian	919246fbc1	don't pass rope_scaling kwarg if it's None (#383 )	2023-08-13 18:57:38 -04:00
Wing Lian	ffac902c1b	bump flash-attn to 2.0.4 for the base docker image (#382 )	2023-08-13 17:55:04 -04:00
Charles Goddard	15f6e57eaa	Fix crash when running without CUDA	2023-08-13 13:36:40 -07:00
NanoCode012	729c299256	Feat(doc): Improve sharegpt doc (#378 ) * Feat(doc): Improve sharegpt doc * Fix typo	2023-08-14 00:36:00 +09:00
Wing Lian	86a91e260b	save tokenizer before training starts (#380 )	2023-08-13 11:28:58 -04:00
Aman Gupta Karmani	094fc2c6e6	try to detect accelerate and only use device_map=None in that case (#373 )	2023-08-13 00:32:07 -04:00
Wing Lian	2dafa730ef	Create FUNDING.yml	2023-08-13 00:30:34 -04:00
Wing Lian	343ac84e5a	fix check for flash attn branching (#377 )	2023-08-12 22:48:08 -04:00
Aman Karmani	0c967279ce	remove unnecessary local variable	2023-08-13 01:58:39 +00:00
Aman Karmani	efb3b2c95e	simplify `load_tokenizer`	2023-08-12 18:55:06 -07:00
Aman Karmani	7b55fe6419	improve GPU logging to break out pytorch cache and system mem	2023-08-12 18:52:57 -07:00
Aman Karmani	e029ab34ea	quiet noise from llama tokenizer by setting pad token earlier	2023-08-12 18:31:40 -07:00
Aman Karmani	8cec513447	extract module for working with cfg	2023-08-12 18:25:27 -07:00
Aman Karmani	a13e45d548	fix DefaultDict.__or__	2023-08-13 01:15:50 +00:00