Initial model upload

LICENSE

@@ -0,0 +1,201 @@
+                                 Apache License
+                           Version 2.0, January 2004
+                        http://www.apache.org/licenses/
+
+   TERMS AND CONDITIONS FOR USE, REPRODUCTION, AND DISTRIBUTION
+
+   1. Definitions.
+
+      "License" shall mean the terms and conditions for use, reproduction,
+      and distribution as defined by Sections 1 through 9 of this document.
+
+      "Licensor" shall mean the copyright owner or entity authorized by
+      the copyright owner that is granting the License.
+
+      "Legal Entity" shall mean the union of the acting entity and all
+      other entities that control, are controlled by, or are under common
+      control with that entity. For the purposes of this definition,
+      "control" means (i) the power, direct or indirect, to cause the
+      direction or management of such entity, whether by contract or
+      otherwise, or (ii) ownership of fifty percent (50%) or more of the
+      outstanding shares, or (iii) beneficial ownership of such entity.
+
+      "You" (or "Your") shall mean an individual or Legal Entity
+      exercising permissions granted by this License.
+
+      "Source" form shall mean the preferred form for making modifications,
+      including but not limited to software source code, documentation
+      source, and configuration files.
+
+      "Object" form shall mean any form resulting from mechanical
+      transformation or translation of a Source form, including but
+      not limited to compiled object code, generated documentation,
+      and conversions to other media types.
+
+      "Work" shall mean the work of authorship, whether in Source or
+      Object form, made available under the License, as indicated by a
+      copyright notice that is included in or attached to the work
+      (an example is provided in the Appendix below).
+
+      "Derivative Works" shall mean any work, whether in Source or Object
+      form, that is based on (or derived from) the Work and for which the
+      editorial revisions, annotations, elaborations, or other modifications
+      represent, as a whole, an original work of authorship. For the purposes
+      of this License, Derivative Works shall not include works that remain
+      separable from, or merely link (or bind by name) to the interfaces of,
+      the Work and Derivative Works thereof.
+
+      "Contribution" shall mean any work of authorship, including
+      the original version of the Work and any modifications or additions
+      to that Work or Derivative Works thereof, that is intentionally
+      submitted to Licensor for inclusion in the Work by the copyright owner
+      or by an individual or Legal Entity authorized to submit on behalf of
+      the copyright owner. For the purposes of this definition, "submitted"
+      means any form of electronic, verbal, or written communication sent
+      to the Licensor or its representatives, including but not limited to
+      communication on electronic mailing lists, source code control systems,
+      and issue tracking systems that are managed by, or on behalf of, the
+      Licensor for the purpose of discussing and improving the Work, but
+      excluding communication that is conspicuously marked or otherwise
+      designated in writing by the copyright owner as "Not a Contribution."
+
+      "Contributor" shall mean Licensor and any individual or Legal Entity
+      on behalf of whom a Contribution has been received by Licensor and
+      subsequently incorporated within the Work.
+
+   2. Grant of Copyright License. Subject to the terms and conditions of
+      this License, each Contributor hereby grants to You a perpetual,
+      worldwide, non-exclusive, no-charge, royalty-free, irrevocable
+      copyright license to reproduce, prepare Derivative Works of,
+      publicly display, publicly perform, sublicense, and distribute the
+      Work and such Derivative Works in Source or Object form.
+
+   3. Grant of Patent License. Subject to the terms and conditions of
+      this License, each Contributor hereby grants to You a perpetual,
+      worldwide, non-exclusive, no-charge, royalty-free, irrevocable
+      (except as stated in this section) patent license to make, have made,
+      use, offer to sell, sell, import, and otherwise transfer the Work,
+      where such license applies only to those patent claims licensable
+      by such Contributor that are necessarily infringed by their
+      Contribution(s) alone or by combination of their Contribution(s)
+      with the Work to which such Contribution(s) was submitted. If You
+      institute patent litigation against any entity (including a
+      cross-claim or counterclaim in a lawsuit) alleging that the Work
+      or a Contribution incorporated within the Work constitutes direct
+      or contributory patent infringement, then any patent licenses
+      granted to You under this License for that Work shall terminate
+      as of the date such litigation is filed.
+
+   4. Redistribution. You may reproduce and distribute copies of the
+      Work or Derivative Works thereof in any medium, with or without
+      modifications, and in Source or Object form, provided that You
+      meet the following conditions:
+
+      (a) You must give any other recipients of the Work or
+          Derivative Works a copy of this License; and
+
+      (b) You must cause any modified files to carry prominent notices
+          stating that You changed the files; and
+
+      (c) You must retain, in the Source form of any Derivative Works
+          that You distribute, all copyright, patent, trademark, and
+          attribution notices from the Source form of the Work,
+          excluding those notices that do not pertain to any part of
+          the Derivative Works; and
+
+      (d) If the Work includes a "NOTICE" text file as part of its
+          distribution, then any Derivative Works that You distribute must
+          include a readable copy of the attribution notices contained
+          within such NOTICE file, excluding those notices that do not
+          pertain to any part of the Derivative Works, in at least one
+          of the following places: within a NOTICE text file distributed
+          as part of the Derivative Works; within the Source form or
+          documentation, if provided along with the Derivative Works; or,
+          within a display generated by the Derivative Works, if and
+          wherever such third-party notices normally appear. The contents
+          of the NOTICE file are for informational purposes only and
+          do not modify the License. You may add Your own attribution
+          notices within Derivative Works that You distribute, alongside
+          or as an addendum to the NOTICE text from the Work, provided
+          that such additional attribution notices cannot be construed
+          as modifying the License.
+
+      You may add Your own copyright statement to Your modifications and
+      may provide additional or different license terms and conditions
+      for use, reproduction, or distribution of Your modifications, or
+      for any such Derivative Works as a whole, provided Your use,
+      reproduction, and distribution of the Work otherwise complies with
+      the conditions stated in this License.
+
+   5. Submission of Contributions. Unless You explicitly state otherwise,
+      any Contribution intentionally submitted for inclusion in the Work
+      by You to the Licensor shall be under the terms and conditions of
+      this License, without any additional terms or conditions.
+      Notwithstanding the above, nothing herein shall supersede or modify
+      the terms of any separate license agreement you may have executed
+      with Licensor regarding such Contributions.
+
+   6. Trademarks. This License does not grant permission to use the trade
+      names, trademarks, service marks, or product names of the Licensor,
+      except as required for reasonable and customary use in describing the
+      origin of the Work and reproducing the content of the NOTICE file.
+
+   7. Disclaimer of Warranty. Unless required by applicable law or
+      agreed to in writing, Licensor provides the Work (and each
+      Contributor provides its Contributions) on an "AS IS" BASIS,
+      WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or
+      implied, including, without limitation, any warranties or conditions
+      of TITLE, NON-INFRINGEMENT, MERCHANTABILITY, or FITNESS FOR A
+      PARTICULAR PURPOSE. You are solely responsible for determining the
+      appropriateness of using or redistributing the Work and assume any
+      risks associated with Your exercise of permissions under this License.
+
+   8. Limitation of Liability. In no event and under no legal theory,
+      whether in tort (including negligence), contract, or otherwise,
+      unless required by applicable law (such as deliberate and grossly
+      negligent acts) or agreed to in writing, shall any Contributor be
+      liable to You for damages, including any direct, indirect, special,
+      incidental, or consequential damages of any character arising as a
+      result of this License or out of the use or inability to use the
+      Work (including but not limited to damages for loss of goodwill,
+      work stoppage, computer failure or malfunction, or any and all
+      other commercial damages or losses), even if such Contributor
+      has been advised of the possibility of such damages.
+
+   9. Accepting Warranty or Additional Liability. While redistributing
+      the Work or Derivative Works thereof, You may choose to offer,
+      and charge a fee for, acceptance of support, warranty, indemnity,
+      or other liability obligations and/or rights consistent with this
+      License. However, in accepting such obligations, You may act only
+      on Your own behalf and on Your sole responsibility, not on behalf
+      of any other Contributor, and only if You agree to indemnify,
+      defend, and hold each Contributor harmless for any liability
+      incurred by, or claims asserted against, such Contributor by reason
+      of your accepting any such warranty or additional liability.
+
+   END OF TERMS AND CONDITIONS
+
+   APPENDIX: How to apply the Apache License to your work.
+
+      To apply the Apache License to your work, attach the following
+      boilerplate notice, with the fields enclosed by brackets "[]"
+      replaced with your own identifying information. (Don't include
+      the brackets!)  The text should be enclosed in the appropriate
+      comment syntax for the file format. We also recommend that a
+      file or class name and description of purpose be included on the
+      same "printed page" as the copyright notice for easier
+      identification within third-party archives.
+
+   Copyright [yyyy] [name of copyright owner]
+
+   Licensed under the Apache License, Version 2.0 (the "License");
+   you may not use this file except in compliance with the License.
+   You may obtain a copy of the License at
+
+       http://www.apache.org/licenses/LICENSE-2.0
+
+   Unless required by applicable law or agreed to in writing, software
+   distributed under the License is distributed on an "AS IS" BASIS,
+   WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
+   See the License for the specific language governing permissions and
+   limitations under the License.

README.md

@@ -1,3 +1,497 @@
----
-license: gpl-3.0
----
+# Mistral-finetune
+
+<a target="_blank" href="https://colab.research.google.com/github/mistralai/mistral-finetune/blob/main/tutorials/mistral_finetune_7b.ipynb">
+  <img src="https://colab.research.google.com/assets/colab-badge.svg" alt="Open In Colab"/>
+</a>
+
+
+`mistral-finetune` is a light-weight codebase that enables memory-efficient and performant finetuning of Mistral's models.
+It is based on [LoRA](https://arxiv.org/abs/2106.09685), a training paradigm where most weights are frozen and only 1-2% additional weights in the form of low-rank matrix perturbations are trained. 
+
+For maximum efficiency it is recommended to use a A100 or H100 GPU. The codebase is optimized 
+for multi-GPU-single-node training setups, but for smaller models, such as the 7B a single GPU suffices.
+
+> **Note**
+> 
+> - The goal of this repository is to provide a simple, guided entrypoint to finetune Mistral models.
+> As such, it is fairly opinionated (especially around data formatting) and does not aim at being exhaustive
+> across multiple model architecture or hardware types.
+> For more generic approaches, you can check out some other great projects like 
+> [torchtune](https://pytorch.org/torchtune/stable/overview.html).
+
+## Installation
+
+To get started with Mistral LoRA fine-tuning, follow these steps:
+
+1. Clone this repository:
+```
+cd $HOME && git clone https://github.com/mistralai/mistral-finetune.git
+```
+
+2. Install all required dependencies:
+```
+cd mistral-finetune
+pip install -r requirements.txt
+```
+
+## Model download
+
+We recommend fine-tuning one of the official Mistral models which you can download here:
+
+| Model          | Link                                                                                                    | Checksum                          |
+|----------------|---------------------------------------------------------------------------------------------------------|-----------------------------------|
+| 7B Base V3       | [7B Base](https://models.mistralcdn.com/mistral-7b-v0-3/mistral-7B-v0.3.tar)                            | `0663b293810d7571dad25dae2f2a5806`|
+| 7B Instruct v3 | [7B Instruct v3](https://models.mistralcdn.com/mistral-7b-v0-3/mistral-7B-Instruct-v0.3.tar)             | `80b71fcb6416085bcb4efad86dfb4d52`|
+| 8x7B Base V1   | [8x7B Base](https://huggingface.co/mistralai/Mixtral-8x7B-v0.1)                                                                        | (HF link)                                |
+| 8x7B Instruct V1 | [8x7B Instruct](https://models.mistralcdn.com/mixtral-8x7b-v0-1/Mixtral-8x7B-v0.1-Instruct.tar) | `8e2d3930145dc43d3084396f49d38a3f` |
+| 8x22 Instruct V3 | [8x22 Instruct](https://models.mistralcdn.com/mixtral-8x22b-v0-3/mixtral-8x22B-Instruct-v0.3.tar)        | `471a02a6902706a2f1e44a693813855b`|
+| 8x22B Base V3  | [8x22B Base](https://models.mistralcdn.com/mixtral-8x22b-v0-3/mixtral-8x22B-v0.3.tar)                        | `a2fa75117174f87d1197e3a4eb50371a`|
+
+**Important Notice**: For 8x7B Base V1 and 8x7B Instruct V1, it is necessary to use our v3 tokenizer and extend the vocabulary size to 32768 prior to fine-tuning. For detailed instructions on this process, please refer to the ["Model extension"](https://github.com/mistralai/mistral-finetune?tab=readme-ov-file#model-extension) section. 
+
+E.g., to download the 7B-base model you can run the following command:
+```sh
+mkdir -p ~/${HOME}/mistral_models
+cd ${HOME} && wget https://models.mistralcdn.com/mistral-7b-v0-3/mistral-7B-v0.3.tar
+tar -xf mistral-7B-v0.3.tar -C mistral_models
+```
+
+Make sure to modify your training script and add the path to the downloaded 
+folder as `model_id_or_path`.
+
+E.g., modify [example/7B.yaml](https://github.com/mistralai/mistral-finetune/blob/main/example/7B.yaml) to include the absolute path to `$HOME/mistral_models/7B`:
+
+```
+model_id_or_path: "/Users/johndoe/mistral_models/7B"
+```
+
+## Prepare dataset 
+
+To ensure effective training, `mistral-finetune` has strict 
+requirements for how the training data has to be formatted.
+
+All data files must be stored in jsonl format files.
+
+You can build two types of data files:
+
+### _Pretrain_:
+
+Pretrain data corresponds to plain text data stored in the `"text"` key. E.g:
+
+```jsonl
+{"text": "Text contained in document n°1"}
+{"text": "Text contained in document n°2"}
+```
+
+### _Instruct_:
+
+Currently two different types of instruction following data are supported:
+
+- _Instruct_: conversational data stored in the `"messages"` key in the form of a list. Each list item is a dictionary containing the `"content"` and `"role"` keys. `"role"` is a string being one of "user", "assistant" or "system_prompt". The loss will only be computed if "role" == "assistant". E.g.:
+
+```jsonl
+{
+  "messages": [
+    {
+      "role": "user",
+      "content": "User interaction n°1 contained in document n°1"
+    },
+    {
+      "role": "assistant",
+      "content": "Bot interaction n°1 contained in document n°1"
+    },
+    {
+      "role": "user",
+      "content": "User interaction n°2 contained in document n°1"
+    },
+    {
+      "role": "assistant",
+      "content": "Bot interaction n°2 contained in document n°1"
+    }
+  ]
+}
+{
+  "messages": [
+    {
+      "role": "user",
+      "content": "User interaction n°1 contained in document n°2"
+    },
+    {
+      "role": "assistant",
+      "content": "Bot interaction n°1 contained in document n°2"
+    },
+    {
+      "role": "user",
+      "content": "User interaction n°2 contained in document n°2"
+    },
+    {
+      "role": "assistant",
+      "content": "Bot interaction n°2 contained in document n°2",
+      "weight": 0,  # don't train on n°2
+    },
+    {
+      "role": "user",
+      "content": "User interaction n°3 contained in document n°2"
+    },
+    {
+      "role": "assistant",
+      "content": "Bot interaction n°3 contained in document n°2"
+    }
+  ]
+}
+```
+
+- _Function calling_: conversational data stored in the `"messages"` key in the form of a list. Each list item is a dictionary containing the `"role"` and `"content"` or `"tool_calls"` keys. `"role"` is a string being one of "user", "assistant", "system_prompt", or "tool". The loss will only be computed if "role" == "assistant".
+
+**Note**: In function calling the `"id"` of `"tool_calls"` and the `"tool_call_id"` are randomly generated strings of exactly 9 chars. We recommend to generate this automatically 
+in a data preparation script as is done [here](https://github.com/mistralai/mistral-finetune/blob/208b25c0f7299bb78d06cea25b82adee03834319/utils/reformat_data_glaive.py#L74).
+
+E.g.:
+
+```jsonl
+{
+  "messages": [
+    {
+      "role": "system",
+      "content": "You are an helpful assistant who has access to the following functions to help the user, you can use the functions if needed"
+    },
+    {
+      "role": "user",
+      "content": "Can you help me generate an anagram of the word \"listen\"?"
+    },
+    {
+      "role": "assistant",
+      "tool_calls": [
+        {
+          "id": "TX92Jm8Zi",
+          "type": "function",
+          "function": {
+            "name": "generate_anagram",
+            "arguments": "{\"word\": \"listen\"}"
+          }
+        }
+      ]
+    },
+    {
+      "role": "tool",
+      "content": "{\"anagram\": \"silent\"}",
+      "tool_call_id": "TX92Jm8Zi"
+    },
+    {
+      "role": "assistant",
+      "content": "The anagram of the word \"listen\" is \"silent\"."
+    },
+    {
+      "role": "user",
+      "content": "That's amazing! Can you generate an anagram for the word \"race\"?"
+    },
+    {
+      "role": "assistant",
+      "tool_calls": [
+        {
+          "id": "3XhQnxLsT",
+          "type": "function",
+          "function": {
+            "name": "generate_anagram",
+            "arguments": "{\"word\": \"race\"}"
+          }
+        }
+      ]
+    }
+  ],
+  "tools": [
+    {
+      "type": "function",
+      "function": {
+        "name": "generate_anagram",
+        "description": "Generate an anagram of a given word",
+        "parameters": {
+          "type": "object",
+          "properties": {
+            "word": {
+              "type": "string",
+              "description": "The word to generate an anagram of"
+            }
+          },
+          "required": [
+            "word"
+          ]
+        }
+      }
+    }
+  ]
+}
+```
+
+## Verify dataset
+
+Before starting a training run you should verify that your dataset is correctly formatted and get an 
+estimation of the training time. You can do so by using the [./utils/validate_data](https://github.com/mistralai/mistral-finetune/blob/main/utils/validate_data.py) script.
+
+Note that this step is crucial to ensure that the data is correctly formatted.
+
+### Instruction following
+
+Let's go over a simple example to train a model in instruction following:
+
+- 1. **Load a chunk of [Ultachat_200k](https://huggingface.co/datasets/HuggingFaceH4/ultrachat_200k)**
+
+Create the data folder and navigate to the folder.
+```sh
+cd $HOME && mkdir -p data && cd $HOME/data
+```
+
+Load the data into a Pandas Dataframe. 
+
+**Note**: Make sure to have pandas and pyarrow installed (`pip install pandas pyarrow`).
+
+```py
+import pandas as pd
+
+df = pd.read_parquet('https://huggingface.co/datasets/HuggingFaceH4/ultrachat_200k/resolve/main/data/test_gen-00000-of-00001-3d4cd8309148a71f.parquet')
+```
+- 2. Split into train and eval
+
+```py
+df_train=df.sample(frac=0.95,random_state=200)
+df_eval=df.drop(df_train.index)
+```
+
+- 3. Save data to jsonl
+
+```py
+df_train.to_json("ultrachat_chunk_train.jsonl", orient="records", lines=True)
+df_eval.to_json("ultrachat_chunk_eval.jsonl", orient="records", lines=True)
+```
+
+- 4. Modify your training yaml to include the ultrachat dataset and verify the yaml
+
+Modify [example/7B.yaml](https://github.com/mistralai/mistral-finetune/blob/main/example/7B.yaml) to include the absolute path to `$HOME/data/ultrachat_chunk_train.jsonl` as well as a dataset mixing weight for training and `$HOME/data/ultrachat_chunk_eval.jsonl` for eval, *e.g.*
+
+```
+data:
+  instruct_data: "/Users/johndoe/data/ultrachat_chunk_train.jsonl"
+  eval_instruct_data: "/Users/johndoe/data/ultrachat_chunk_eval.jsonl"
+```
+
+Now you can verify your training yaml to make sure the data is correctly formatted and to get an estimate of your training time.
+
+```
+cd $HOME/mistral-finetune
+python -m utils.validate_data --train_yaml example/7B.yaml
+```
+
+Upon completion you should see an error report with many of the following errors:
+
+```
+The data in line 1412 of dataset /Users/johndoe/data/ultrachat_chunk_eval.jsonl is incorrectly formated.Expected last role to be one of: [assistant] but got user
+The data in line 1413 of dataset /Users/johndoe/data/ultrachat_chunk_eval.jsonl is incorrectly formated.Expected last role to be one of: [assistant] but got user
+The data in line 1414 of dataset /Users/johndoe/data/ultrachat_chunk_eval.jsonl is incorrectly formated.Expected last role to be one of: [assistant] but got user
+The data in line 1415 of dataset /Users/johndoe/data/ultrachat_chunk_eval.jsonl is incorrectly formated.Expected last role to be one of: [assistant] but got user
+```
+
+Many conversations seem to end with the 'user' role which is unnecessary as we only train on 'assistant' messages and thus would unnecessarily process data.
+
+You can make use of [./utils/reformat_data.py](https://github.com/mistralai/mistral-finetune/blob/main/utils/reformat_data.py) to correct the data:
+
+```
+cd $HOME/mistral-finetune
+python -m utils.reformat_data $HOME/data/ultrachat_chunk_train.jsonl
+python -m utils.reformat_data $HOME/data/ultrachat_chunk_eval.jsonl
+```
+
+You should see that a couple of samples will be skipped.
+
+- 5. Potentially change number of training steps
+
+Upon correction of the dataset, run the script again
+
+```
+cd $HOME/mistral-finetune
+python -m utils.validate_data --train_yaml example/7B.yaml
+```
+
+You should get a summary of the data input and training parameters:
+
+```
+Train States
+ --------------------
+{
+   "expected": {
+       "eta": "00:52:44",
+       "data_tokens": 25169147,
+       "train_tokens": 131072000,
+       "epochs": "5.21",
+       "max_steps": 500,
+       "data_tokens_per_dataset": {
+           "/Users/johndoe/data/ultrachat_chunk_train.jsonl": "25169147.0"
+       },
+       "train_tokens_per_dataset": {
+           "/Users/johndoe/data/ultrachat_chunk_train.jsonl": "131072000.0"
+       },
+       "epochs_per_dataset": {
+           "/Users/johndoe/data/ultrachat_chunk_train.jsonl": "5.2"
+       }
+   },
+}
+```
+
+Having `max_steps` set to 500 would lead to iterating through the dataset roughly 5 times which is reasonable, but might 
+be a bit too much. A recommended setting is shown below which would only take 30min on a 8xH100 cluster.
+
+### Function calling
+
+Next let's go over a more advanced use case to fine-tune a model on function calling.
+Function calling requires the data to be in the format as [explained above](#instruct). Let's go over an example.
+
+- 1. **Load a chat-formatted version of the [Glaive function calling dataset](https://huggingface.co/datasets/Locutusque/function-calling-chatml)**
+
+Create the data folder and navigate to the folder.
+```sh
+cd $HOME && mkdir -p data && cd $HOME/data
+```
+
+Load the data into a Pandas Dataframe.
+
+**Note**: Make sure to have pandas and pyarrow installed (`pip install pandas pyarrow`).
+
+```py
+import pandas as pd
+
+df = pd.read_parquet('https://huggingface.co/datasets/Locutusque/function-calling-chatml/resolve/main/data/train-00000-of-00001-f0b56c6983b4a78f.parquet')
+```
+- 2. Split into train and eval
+
+```py
+df_train=df.sample(frac=0.95,random_state=200)
+df_eval=df.drop(df_train.index)
+```
+
+- 3. Save data to jsonl
+
+```py
+df_train.to_json("glaive_train.jsonl", orient="records", lines=True)
+df_eval.to_json("glaive_eval.jsonl", orient="records", lines=True)
+```
+
+- 4. Reformat dataset
+
+As one can see the dataset does not follow the required function calling format, so it will need to be reformatted. Among other things `"from"` should be renamed to `"user"` and superfluous `"\n"` characters should be removed.
+For this dataset you can make use of [`./utils/reformat_data_glaive.py`](https://github.com/mistralai/mistral-finetune/blob/main/utils/reformat_data_glaive.py):
+
+```
+cd $HOME/mistral-finetune
+python -m utils.reformat_data_glaive $HOME/data/glaive_train.jsonl
+python -m utils.reformat_data_glaive $HOME/data/glaive_eval.jsonl
+```
+
+Running this command will make sure that most samples are in the correct format.
+
+**Note**: It is impossible to write reformatting scripts that work for all kinds of datasets. 
+If you have datasets that don't yet follow the required format above, you will most probably have to 
+create a reformatting script yourself (mistral-chat or chat-gpt is your best friend here!).
+
+- 5. Validate dataset
+
+You can now validate the dataset by setting `data.instruct_data` and `data.eval_instruct_data` to
+`$HOME/data/glaive_train.jsonl` and `$HOME/data/glaive_eval.jsonl` in `example/7B.yaml` respectively.
+
+The reformatted datasets still has some errors which can be removed with `--create_corrected`. For this, make sure to add
+`--create_corrected` as follows:
+
+```
+cd $HOME/mistral-finetune
+python -m utils.validate_data --train_yaml example/7B.yaml --create_corrected
+```
+
+Running this command will show a couple of errors and save two new datasets `$HOME/data/glaive_train.jsonl.corrected` and `$HOME/data/glaive_eval.jsonl.corrected`. Make sure to use these two dataset in `example/7B.yaml` and run the command again. Now the dataset should be correctly formatted!
+
+
+## Start training
+
+Having followed the [dataset verification section](#verify-dataset), we can now start training.
+For faster training, we recommend setting max_steps to only 300. Make sure to define `run_dir` to your experiment folder and optionally set `wandb_project` to a Weights & Biases project for logging`, *e.g.*:
+```
+max_steps: 300
+run_dir: "/Users/johndoe/ultra_chat_test"
+wandb.project: ultra_chat
+```
+
+Optionally you can also set `wandb`
+
+Save the training configuration and start training! Make sure to set `--nproc-per-node` to the number of available GPUs.
+
+```
+cd $HOME/mistral-finetune
+torchrun --nproc-per-node 8 --master_port $RANDOM -m train example/7B.yaml
+```
+
+Training on ultra-chat should take around 30min on a 8xH100 node and the resulting weights should give an MT Bench score around 6.3.
+
+Training on glaive should take around 1h on a 8xH100 node and the resulting weights should work nicely for function calling.
+
+## Customizing training configuration
+
+The example `mistral-finetune/examples/7B` defines reasonable parameters for learning rate, weight decay, etc... but you are advised to 
+customize these settings for your use case.
+
+Generally, a training configuration should fill the following parameters:
+
+- `model_id_or_path` defines the model to start training from. This can be a path to a pre-trained model or a local model directory.
+- `run_dir` defines the directory where training checkpoints and metrics are stored.
+- `seq_len` defines the sequence length for training. This is the maximum length of input sequences the model will process. Samples are packed to reach a length of `seq_len` for maximum training efficiency.
+- `batch_size` defines the number of training examples used per GPU. **Note**: The overall effective batch_size (in tokens) across all GPUs equals `num_gpus` x `batch_size` x `seq_len`.
+- `max_steps` defines the maximum number of training steps. This is the total number of iterations the training process will run. It can be adjusted based on the specific needs of your training scenario. Total number of tokens seen during training is `max_steps` x `num_gpus` x `batch_size` x `seq_len`.
+- `optim.lr` defines the learning rate. This is the initial learning rate for the optimizer.
+- `optim.weight_decay` defines weight decay. Weight decay is a regularization technique used to prevent overfitting by penalizing large weights. We recommend leaving it at 0.1.
+- `optim.pct_start` defines the percentage of the total training steps used for the learning rate warm-up phase before it starts to decrease. It corresponds to pct_start of PyTorch's OneCycleLR.
+- `lora.rank` defines the size of the LoRA (Low-Rank Adaptation) adapters. We recommend 64 or less, which adjusts the rank of the low-rank decomposition used in LoRA.
+- `seed` defines the random seed for initialization and data shuffling/sampling. Setting a seed ensures reproducibility of results.
+- `log_freq` defines the logging frequency. This specifies how often (in steps) to log training metrics.
+- `data.instruct_data` is the path to the instruction data used for training. This field has to be filled with one or multiple data sources in the format as explained above. Each data source should either be a path to jsonl file of a path to a directory containing jsonl files followed by a weighting to define the importance of this dataset: `<path/to/data_source>:<weight>`. E.g.: `data.instruct_data: "/path/to/data1.jsonl:5.,/path/to/data2.jsonl:1.,/path/to/dir_of_jsonls:1."`
+- `data.data` is an optional path to additional pretraining data in the format as explained above. Note that this field can be left blank.
+- `data.eval_instruct_data` is an optional path to evaluation instruction data to run cross-validation at every `eval_freq` steps. Cross-validation metrics are displayed as `loss` and `perplexity`.
+- `eval_freq` defines how often (in steps) to evaluate the model. This specifies the interval at which the model is evaluated on the validation set.
+- `no_eval` is a flag to enable or disable intermediate evaluation. Setting it to False enables periodic evaluation during training.
+- `ckpt_freq` defines how often (in steps) to save checkpoints. This specifies the interval at which the model's state is saved.
+- `ckpt_only_lora` defines whether to only save the trained LoRA checkpoints or whether the trained LoRA should directly be merged into the base model and saved. **Note**: When setting `ckpt_only_lora=False` make sure that you have enough CPU and GPU memory to save the full model on a single process (this is usually only possible for the 7B model).
+- `wandb.key` is used to pass your Weights & Biases (wandb) API key for logging. This allows you to log training metrics to the wandb dashboard.
+- `wandb.project` defines the wandb project name. This is where the training run will be logged in the wandb interface.
+
+## Inference
+
+Once your model is trained, you should try it out in inference. We recommend using [mistral-inference](https://github.com/mistralai/mistral-inference). 
+
+Make sure to have `mistral_inference` correctly installed:
+```
+pip install mistral_inference
+```
+
+Assuming your `lora.safetensors` is saved under `$HOME/ultra_chat_test/checkpoints/checkpoint_000300/consolidated/lora.safetensors`, you can chat with the model using `mistral_inference`, *e.g.*:
+
+```sh
+mistral-chat /mnt/slow/runs/patrick/mistral-finetune/7B/ --max_tokens 256 --temperature 1.0 --instruct --lora_path $HOME/ultra_chat_test/checkpoints/checkpoint_000300/consolidated/lora.safetensors
+```
+
+## Model extension
+
+**Important**: Note that one can only fine-tune mistral models that are compatible with the v3 tokenizer which entails that the models have a vocabulary size of 32768 - not 32000. One can however easily extend older version of vocabulary size 32000 to have a vocabulary size of 32768 by using:
+```
+python -m utils.extend_model_vocab --original_model_ckpt /folder/to/old/model --extended_model_ckpt /folder/to/extended/model
+```
+
+Once the extension has worked, one can fine-tune using the newly created model checkpoint in `/folder/to/extended/model`.
+
+## FAQ:
+
+> - What's the best practice of fine-tuning MoEs?
+
+We see a higher degree of performance variance in when fine-tuning MoE models. It's not unusual to find that fine-tuning MoEs models with different seeds can lead to a high variance in performance. We did not observe such a high variance with dense models. Therefore, we suggest running multiple instances of the same fine-tuning process on MoEs models and selecting the one that performs best.
+
+> - How can I determine the number of tokens used during the model training process?
+  
+You can use the following script to find out: https://github.com/mistralai/mistral-finetune/blob/main/utils/validate_data.py. This script accepts a .yaml training file as input and returns the number of tokens the model is being trained on.
+
+> - What should I do if I encounter a CUDA out-of-memory error?
+  
+One possible solution is to reduce the batch size per GPU. The batch size is equal to `seq_len` x `batch_size`. Try setting `batch_size` to 1 and reduce `seq_len`. You can define the `batch_size` and `seq_len` in the .yaml file.

example/7B.yaml

@@ -0,0 +1,37 @@
+# data
+data:
+  instruct_data: "/root/data/mol_instructions_train.jsonl"  # Fill this with the path to your training data
+  data: ""  # Optionally fill with pretraining data
+  eval_instruct_data: ""  # Optionally fill with evaluation data
+
+# model
+model_id_or_path: "/root/mistral_models/7B-v0.3"  # Path to downloaded model
+lora:
+  rank: 64
+
+# optim
+seq_len: 32768
+batch_size: 2
+ #TODO try other values
+max_steps: 500
+optim:
+  lr: 5.e-5
+  weight_decay: 0.05
+  pct_start: 0.05
+
+# other
+seed: 99
+log_freq: 1
+eval_freq: 100
+no_eval: True
+ckpt_freq: 100
+
+ckpt_only_lora: False  # Save only trained LoRA adapters. Set to `False` to merge LoRA adapter into the base model and save full fine-tuned model
+
+run_dir: "/root/mistral-finetune/runseed99"
+
+wandb:
+  project: "CHEMISTral7b-ft"
+  offline: False  # Set to True if you want to use wandb in offline mode
+  key: "aaf77f83a4e316f6a8b47fa975ab6b5e73c7c8df"  # Optionally set your WandB API key
+  run_name: "runseed99"  # Optionally name your WandB run

finetune/args.py

@@ -0,0 +1,116 @@
+import logging
+import os
+from dataclasses import dataclass, field
+from pathlib import Path
+from typing import Optional
+
+from simple_parsing.helpers import Serializable
+
+from model.args import LoraArgs
+
+from .data.args import DataArgs
+
+
+@dataclass
+class OptimArgs(Serializable):
+    lr: float = 3e-4
+    weight_decay: float = 0.1
+    pct_start: float = 0.3
+
+
+@dataclass
+class WandbArgs(Serializable):
+    project: Optional[str] = None  # Fill this argument to use wandb.
+    offline: bool = False
+    key: Optional[str] = None
+    run_name: Optional[str] = None
+
+    def __post_init__(self) -> None:
+        if self.project is not None:
+            try:
+                import wandb  # noqa: F401
+            except ImportError:
+                raise ImportError("`wandb` not installed. Either make sure `wandb` is installed or set `wandb:project` to None.")
+
+            if len(self.project) == 0:
+                raise ValueError("`wandb.project` must not be an empty string.")
+
+@dataclass
+class MLFlowArgs(Serializable):
+    tracking_uri: Optional[str] = None
+    experiment_name: Optional[str] = None
+
+    def __post_init__(self) -> None:
+        if self.tracking_uri is not None:
+            try:
+                import mlflow  # noqa: F401
+            except ImportError:
+                raise ImportError("`mlflow` not installed. Either make sure `mlflow` is installed or set `mlflow.tracking_uri` to None.")
+
+            if self.experiment_name is None:
+                raise ValueError("If `mlflow.tracking_uri` is set, `mlflow.experiment_name` must be set as well.")
+
+
+
+@dataclass
+class TrainArgs(Serializable):
+    data: DataArgs
+
+    # if specified, instruct_tokenizer and model will be loaded
+    model_id_or_path: str  # Path to the directory containing the initial model or model id: "mistral-small"
+
+    run_dir: str  # Path to the directory where everything will be saved. It needs to be empty.
+    # Name of the wandb run, if None it will be set to the name of the run_dir.
+
+    optim: OptimArgs = field(default_factory=OptimArgs)
+    seed: int = 0
+    # Number of steps to accumulate gradients before calling doing an optimizer step.
+    num_microbatches: int = 1
+
+    seq_len: int = 2048  # Number of tokens per batch per device.
+    batch_size: int = 1
+    max_norm: float = 1.0  # Gradient clipping.
+    max_steps: int = 100  # Number of training steps.
+    log_freq: int = 1  # Number of steps between each logging.
+
+    # Number of steps between each checkpoint saving. If inferior to 1, only the last checkpoint will be saved.
+    ckpt_freq: int = 0
+    ckpt_only_lora: bool = True
+    # If True, no checkpoint will be saved. This is useful for development.
+    no_ckpt: bool = False
+    num_ckpt_keep: Optional[int] = 3
+    eval_freq: int = 0
+    no_eval: bool = True
+
+    # Efficiency
+    # Determines whether gradient checkpointing should be utilized or not during the training process. Gradient checkpointing can be beneficial in reducing memory usage at the cost of slightly longer training times.
+    checkpoint: bool = True
+
+    world_size: Optional[int] = field(init=False, default=None)
+
+    # logging
+    wandb: WandbArgs = field(default_factory=WandbArgs)
+    mlflow: MLFlowArgs = field(default_factory=MLFlowArgs)
+
+    # LoRA
+    lora: Optional[LoraArgs] = field(default_factory=LoraArgs)
+
+    def __post_init__(self) -> None:
+        assert getattr(self, "world_size", None) is None
+        self.world_size = int(os.environ.get("WORLD_SIZE", -1))
+
+        if self.wandb.offline:
+            command = f"cd {self.run_dir}; wandb sync --sync-all"
+            logging.info(f"to sync wandb offline, run: {command}")
+
+        assert self.num_microbatches >= 1
+
+        assert self.num_ckpt_keep is None or self.num_ckpt_keep >= 1
+
+        if self.model_id_or_path is not None:
+            Path(self.model_id_or_path).exists()
+
+        if not self.ckpt_only_lora:
+            logging.warning(
+                "You are have disabled `ckpt_only_lora` and are thus merging the trained LoRA checkpoint into the base model upon checkpointing. This might lead to OOM erros - make sure you have enough CPU and GPU memory."
+            )

finetune/checkpointing.py

@@ -0,0 +1,246 @@
+import json
+import logging
+import shutil
+from pathlib import Path
+from typing import Dict, List, Optional, Union
+
+import safetensors.torch
+import torch
+from mistral_common.tokens.tokenizers.sentencepiece import InstructTokenizerBase
+from torch.distributed import barrier
+from torch.distributed.fsdp.fully_sharded_data_parallel import FullyShardedDataParallel
+
+from model.transformer import LoRALinear
+
+from .distributed import get_rank, get_world_size
+from .utils import TrainState
+
+logger = logging.getLogger("checkpointing")
+
+
+def main_logger_info(message: str) -> None:
+    if get_rank() == 0:
+        logger.info(message)
+
+
+class Checkpointer:
+    """A class to save PyTorch model and optimizer states"""
+
+    def __init__(
+        self,
+        model: FullyShardedDataParallel,
+        state: TrainState,
+        run_dir: Union[Path, str],
+        optimizer: Optional[torch.optim.Optimizer] = None,
+        num_ckpt_keep: Optional[int] = None,
+    ):
+        self.model = model
+        self.optimizer = optimizer
+        self.state = state
+        self.run_dir = Path(run_dir)
+        self.rank = get_rank()
+        self.num_ckpt_keep = num_ckpt_keep
+
+    @property
+    def ckpt_dir(self) -> Path:
+        return self.run_dir / "checkpoints"
+
+    @property
+    def dst_dir(self) -> Path:
+        return self.ckpt_dir / f"checkpoint_{self.state.step:06d}" / "consolidated"
+
+    @staticmethod
+    def consolidated_path(
+        ckpt_dir: Path, use_safetensors: bool, save_only_lora: Optional[bool] = False
+    ) -> Path:
+        suffix = "safetensors" if use_safetensors else "00.pth"
+        prefix = "lora" if save_only_lora else "consolidated"
+
+        return ckpt_dir / f"{prefix}.{suffix}"
+
+    @staticmethod
+    def _tmp(ckpt_dir: Path) -> Path:
+        return ckpt_dir.with_name(f"tmp.{ckpt_dir.name}")
+
+    def write_params_info(self, tmp_dst: Path):
+        params_path = tmp_dst / "params.json"
+        with open(params_path, "w") as f:
+            model_args = self.model.args.to_dict()
+
+            f.write(json.dumps(model_args, indent=4))
+
+    def delete_old_ckpts(self) -> List[Path]:
+        all_saved_ckpts = [d for d in self.ckpt_dir.iterdir() if d.is_dir()]
+
+        # Sort directories by creation time (oldest to newest)
+        all_saved_ckpts.sort(key=lambda x: x.stat().st_ctime, reverse=True)
+
+        ckpts_to_delete = all_saved_ckpts[self.num_ckpt_keep :]
+
+        for ckpt_to_delete in ckpts_to_delete:
+            try:
+                shutil.rmtree(ckpt_to_delete)
+                main_logger_info(f"Deleted ckpt: {ckpt_to_delete}")
+            except OSError as e:
+                main_logger_info(f"Error deleting directory {ckpt_to_delete}: {e}")
+
+        return ckpts_to_delete
+
+    @staticmethod
+    def get_lora_states(state_dict: Dict[str, torch.Tensor]) -> Dict[str, torch.Tensor]:
+        return {k: v for k, v in state_dict.items() if "lora" in k}
+
+    @staticmethod
+    def get_non_lora_states(
+        state_dict: Dict[str, torch.Tensor]
+    ) -> Dict[str, torch.Tensor]:
+        return {
+            k: v
+            for k, v in state_dict.items()
+            if not any(l_key in k for l_key in ["lora", "frozen"])
+        }
+
+    @torch.no_grad()
+    def retrieve_save_states(
+        self, save_only_lora: bool, save_dtype: torch.dtype
+    ) -> Dict[str, torch.Tensor]:
+        if save_only_lora:
+            assert (
+                self.model.args.lora.enable
+            ), "Cannot save LoRA checkpoint as LoRA training is not enabled."
+
+        # remove all potential hooks
+        for module in self.model.modules():
+            if isinstance(module, LoRALinear) and hasattr(module, "_merge_lora_handle"):
+                module._merge_lora_handle.remove()  # type: ignore
+
+        # merge weights if we don't just save LoRA
+        if not save_only_lora:
+
+            def merge_lora(
+                m: torch.nn.Module,
+                destination: Dict[str, torch.Tensor],
+                prefix: str,
+                *args,
+            ):
+                weight = m.merge_weight()  # type: ignore
+                destination[prefix + "weight"] = weight
+
+            for module in self.model.modules():
+                if isinstance(module, LoRALinear):
+                    module._merge_lora_handle = module._register_state_dict_hook(
+                        merge_lora
+                    )
+
+        offload_to_cpu = get_world_size() > 1
+        if save_only_lora:
+
+            def is_trainable_fsdp(
+                module: Union[torch.nn.Module, FullyShardedDataParallel]
+            ):
+                is_fsdp = isinstance(module, FullyShardedDataParallel)
+                all_params_have_grads = is_fsdp and all(
+                    p.requires_grad is True for p in module.parameters()
+                )
+
+                # need to make sure only lowest fsdp wrap is used
+                is_leaf_node = is_fsdp and len(list(module.module.children())) == 0  # type: ignore
+
+                return is_fsdp and all_params_have_grads and is_leaf_node
+
+            # extract all modules with only trainable weights
+            modules = {
+                k: m for k, m in self.model.named_modules() if is_trainable_fsdp(m)
+            }
+
+            states = {}
+            for key, module in modules.items():
+                assert isinstance(
+                    module, FullyShardedDataParallel
+                ), "`module` should be an instance of `FullyShardedDataParallel`"
+                parent_prefix = key.replace("_fsdp_wrapped_module.", "").replace(
+                    "_checkpoint_wrapped_module.", ""
+                )
+                with module.summon_full_params(
+                    module, writeback=True, offload_to_cpu=offload_to_cpu
+                ):
+                    states.update(
+                        {
+                            f"{parent_prefix}.{k}": v.to(dtype=save_dtype)
+                            for k, v in module.state_dict().items()
+                        }
+                    )
+        else:
+            # make sure you have enough CPU RAM available to save the full model
+            assert isinstance(
+                self.model, FullyShardedDataParallel
+            ), "`self.model` should be an instance of `FullyShardedDataParallel`"
+            with self.model.summon_full_params(
+                self.model, writeback=True, offload_to_cpu=offload_to_cpu
+            ):
+                states = self.get_non_lora_states(self.model.state_dict())
+                states = {k: v.to(dtype=save_dtype) for k, v in states.items()}
+
+        states = dict(sorted(states.items()))
+        return states
+
+    @staticmethod
+    def save_tokenizer(instruct_tokenizer: InstructTokenizerBase, tmp_dst: Path):
+        serialized_spm = instruct_tokenizer.tokenizer._model.serialized_model_proto()  # type: ignore
+
+        tokenizer_path = tmp_dst / "tokenizer.model.v3"
+
+        with open(tokenizer_path, "wb") as f:
+            f.write(serialized_spm)
+
+    @torch.no_grad()
+    def save_checkpoint(
+        self,
+        save_only_lora: bool,
+        dtype: torch.dtype = torch.float16,
+        instruct_tokenizer: Optional[InstructTokenizerBase] = None,
+    ):
+        tmp_dst = self._tmp(self.dst_dir)
+        main_logger_info(
+            f"Dumping checkpoint in {self.dst_dir} using tmp name: {tmp_dst.name}"
+        )
+
+        assert not self.dst_dir.exists(), f"dst exists {self.dst_dir}"
+        tmp_dst.mkdir(parents=True, exist_ok=True)
+
+        states: Dict[str, torch.Tensor] = self.retrieve_save_states(
+            save_only_lora, dtype
+        )
+
+        barrier()
+
+        if self.rank == 0:
+            # save checkpoint in tmp path
+            safetensors.torch.save_file(
+                states,
+                self.consolidated_path(
+                    tmp_dst, use_safetensors=True, save_only_lora=save_only_lora
+                ),  # always use safetensors for checkpointing
+            )
+
+            self.write_params_info(tmp_dst)
+
+            # save tokenizer
+            if instruct_tokenizer is not None:
+                self.save_tokenizer(instruct_tokenizer, tmp_dst)
+
+            assert not self.dst_dir.exists(), f"should not happen! {self.dst_dir}"
+            tmp_dst.rename(self.dst_dir)
+
+            logger.info(
+                f"Done dumping checkpoint in {self.dst_dir} for step: {self.state.step}"
+            )
+
+            # delete last n checkpoints
+            if self.num_ckpt_keep is not None:
+                ckpts_to_delete = self.delete_old_ckpts()
+                logger.info(
+                    f"Done deleting checkpoints {', '.join([str(c) for c in ckpts_to_delete])}"
+                )
+
+        main_logger_info("Done!")

finetune/data/args.py

@@ -0,0 +1,61 @@
+import logging
+from dataclasses import dataclass, field
+
+from simple_parsing.helpers import Serializable
+
+logger = logging.getLogger("data")
+
+
+@dataclass()
+class InstructArgs(Serializable):
+    shuffle: bool = True
+
+    # For function calling training examples only the last tool call
+    # of the assistant message can be used for training. Therefore,
+    # we chunk longer function calling conversations into multiple
+    # training samples to not loose any data. E.g.:
+    # [[
+    #   UserMessage_1, AssisantToolCallMessage_1, ToolMessage_1, AssisantMessage_1
+    #   UserMessage_2, AssisantToolCallMessage_2, ToolMessage_2, AssisantMessage_2
+    # ]]
+    # => is chunked into two training samples:
+    # [[
+    #   UserMessage_1, AssisantToolCallMessage_1, ToolMessage_1, AssisantMessage_1
+    # ],
+    # [
+    #   UserMessage_1, AssisantToolCallMessage_1, ToolMessage_1, AssisantMessage_1
+    #   UserMessage_2, AssisantToolCallMessage_2, ToolMessage_2, AssisantMessage_2
+    # ]]
+    # NOTE: Only if your data is already pre-chunked should this argument be set to False
+    dynamic_chunk_fn_call: bool = True
+
+
+@dataclass()
+class DataArgs(Serializable):
+    # The data arguments `data` and `instruct_data` are a string in the format
+    # "data_source_dir_1:weight_1,data_source_dir_2:weight_2,...". The weight
+    # will be used to sample the data sources. If the sum of the weights is
+    # not 1 when concatenating the two arguments `data` and `instruct_data`,
+    # it will be normalized. The data sources folders must contain jsonl files.
+    # If the value is an empty string, no data will be used for the corresponding
+    # data type.
+    data: str = (
+        ""  # Each line in the jsonl files inside the data source directories must be a dictionary with a "text" key. See Readme for more details. Can be left empty.
+    )
+    shuffle: bool = False
+    instruct_data: str = (
+        ""  # Each line in the jsonl files inside the data source directories must be a dictionary with a "interactions" key. See Readme for more details. Can be left empty.
+    )
+    eval_instruct_data: str = (
+        ""  # Each line in the jsonl files inside the data source directories must be a dictionary with a "interactions" key. See Readme for more details. Can be left empty.
+    )
+    instruct: InstructArgs = field(default_factory=InstructArgs)
+
+    def __post_init__(self) -> None:
+        if (
+            self.instruct.shuffle is False
+            and self.instruct.dynamic_chunk_fn_call is True
+        ):
+            raise ValueError(
+                "Make sure to either enable `data.instruct.shuffle=True` or `data.instruct.dynamic_chunk_fn_call=False`. Dynamic chunking is only possible if data is loaded and shuffled before training."
+            )

finetune/data/data_loader.py

@@ -0,0 +1,126 @@
+import dataclasses
+from typing import Any, Iterator, List, Optional
+
+import numpy as np
+from mistral_common.tokens.tokenizers.sentencepiece import InstructTokenizerBase
+
+from .args import DataArgs
+from .dataset import build_dataset
+
+
+@dataclasses.dataclass
+class Batch:
+    x: np.ndarray
+    y: np.ndarray
+    sizes: List[int]
+    y_mask: Optional[np.ndarray] = None
+    is_pad_only: bool = False
+
+    def __post_init__(self):
+        assert self.x.ndim == 1
+        assert self.x.shape == self.y.shape
+        assert self.x.dtype == np.int64
+        assert self.y.dtype == np.int64
+        assert isinstance(self.sizes, list)
+        assert sum(self.sizes) == self.x.size == self.y.size
+
+        if self.y_mask is not None:
+            assert self.y_mask.size == self.y.size, (self.y_mask.shape, self.y.shape)
+            assert self.y_mask.dtype == bool
+            assert sum(self.sizes) == self.y_mask.size
+            assert not self.y_mask.all()
+            assert self.y_mask.any()
+
+        if self.is_pad_only:
+            assert np.sum(np.abs(self.y)) == 0
+            assert np.sum(np.abs(self.x)) == 0
+            assert self.y_mask is None
+            # create all 0's mask for pad samples
+            self.y_mask = np.zeros_like(self.x)
+
+
+
+
+@dataclasses.dataclass
+class BatchList:
+    x: List[List[int]] = dataclasses.field(default_factory=list)
+    y: List[List[int]] = dataclasses.field(default_factory=list)
+    sizes: List[List[int]] = dataclasses.field(default_factory=list)
+    y_mask: List[Optional[List[int]]] = dataclasses.field(default_factory=list)
+
+    def __post_init__(self):
+        assert self.x == [], "`BatchList` has to be empty at init."
+        assert self.y == [], "`BatchList` has to be empty at init."
+        assert self.sizes == [], "`BatchList` has to be empty at init."
+        assert self.y_mask == [], "`BatchList` has to be empty at init."
+
+    def __len__(self) -> int:
+        return len(self.x)
+
+    def add(self, x: List[int], y: List[int], sizes: List[int], y_mask: Optional[List[int]] = None):
+        self.x.append(x)
+        self.y.append(y)
+        self.sizes.append(sizes)
+        self.y_mask.append(y_mask)
+
+    def empty(self):
+        self.x = []
+        self.y = []
+        self.sizes = []
+        self.y_mask = []
+
+    @staticmethod
+    def flatten_to_numpy(list_of_lists: List[List[Any]], dtype: np.dtype) -> np.array:
+        return np.array([el for sublist in list_of_lists for el in sublist], dtype=dtype)
+
+    def create_batch(self) -> Batch:
+        x_np: np.array = self.flatten_to_numpy(self.x, dtype=np.int64)
+        y_np: np.array = self.flatten_to_numpy(self.y, dtype=np.int64)
+        sizes = sum(self.sizes, [])  # noqa
+
+        y_mask_np: Optional[np.array] = self.flatten_to_numpy(self.y_mask, dtype=bool)
+        y_mask_np = None if y_mask_np.all() else y_mask_np
+
+        return Batch(x_np, y_np, sizes, y_mask_np)
+
+
+
+
+def build_data_loader(
+    instruct_tokenizer: InstructTokenizerBase,
+    args: DataArgs,
+    batch_size: int,
+    seq_len: int,
+    seed: Optional[int],
+    rank: int,
+    world_size: int,
+    is_eval: bool,
+) -> Iterator[Batch]:
+    pretrain_data = args.data if not is_eval else ""
+    instruct_data = args.instruct_data if not is_eval else args.eval_instruct_data
+
+    dataset = build_dataset(
+        pretrain_data=pretrain_data,
+        instruct_data=instruct_data,
+        instruct_args=args.instruct,
+        instruct_tokenizer=instruct_tokenizer,
+        seq_len=seq_len,
+        seed=seed,
+        rank=rank,
+        world_size=world_size,
+        is_eval=is_eval,
+        shuffle_pretrain=args.shuffle,
+    )
+
+    batch_list = BatchList()
+    for sample in dataset:
+        assert all(s >= 0 for s in sample.sizes)
+
+        batch_list.add(sample.x, sample.y, sample.sizes, sample.mask)
+
+        if len(batch_list) == batch_size:
+            batch: Batch = batch_list.create_batch()
+            yield batch
+
+            batch_list.empty()
+

finetune/data/dataset.py

@@ -0,0 +1,475 @@
+import dataclasses
+import itertools
+import json
+import logging
+from dataclasses import dataclass
+from pathlib import Path
+from typing import Any, Dict, Iterator, List, Optional, Set, Tuple, Union
+
+import numpy as np
+import torch.distributed as dist
+from mistral_common.protocol.instruct.messages import (
+    FinetuningAssistantMessage,
+    SystemMessage,
+)
+from mistral_common.tokens.tokenizers.sentencepiece import InstructTokenizerBase
+
+from finetune.distributed import get_rank
+
+from .args import InstructArgs
+from .tokenize import (
+    Mask,
+    SampleType,
+    TokenSample,
+    TrainingInstructSample,
+    build_instruct_sample,
+    encode,
+)
+
+logger = logging.getLogger("dataset")
+
+
+_LOADED_DATASETS: Dict[Path, List[str]] = {}
+
+
+def main_logger_info(message: str) -> None:
+    if dist.is_initialized() and get_rank() == 0:
+        logger.info(message)
+
+
+def load_file(path: Path, world_size: int, rank: int) -> List[str]:
+    lines = []
+    with path.open() as f:
+        for idx, line in enumerate(f):
+            if not idx % world_size == rank:
+                continue
+            lines.append(line)
+    return lines
+
+
+def maybe_load_local_dataset(
+    path: Path, chunk: bool, rank: int, world_size: int, instruct_tokenizer: InstructTokenizerBase, sample_type: SampleType
+) -> List[TokenSample]:
+    global _LOADED_DATASETS
+
+    if path in _LOADED_DATASETS:
+        return _LOADED_DATASETS[path]
+
+    main_logger_info(f"Loading {path} ...")
+    lines: List[str] = load_file(path, rank=rank, world_size=world_size)
+
+    if chunk:
+        lines += maybe_chunk_lines(lines)
+
+    tokens_list: List[TokenSample] = []
+    for line in lines:
+        data = json.loads(line)
+
+        token_sample: TokenSample = encode(
+            data,
+            instruct_tokenizer=instruct_tokenizer,
+            as_type=sample_type,
+        )
+        tokens_list.append(token_sample)
+
+    main_logger_info(f"{path} loaded and tokenized.")
+    _LOADED_DATASETS[path] = tokens_list
+
+    return _LOADED_DATASETS[path]
+
+
+@dataclass
+class DataDir:
+    path: Path
+    sample_type: SampleType
+
+    @property
+    def jsonl_files(self):
+        assert self.path.exists(), f"Make sure that {self.path} exists"
+        jsonl_files = list(self.path.rglob("*jsonl"))
+        assert (
+            len(jsonl_files) > 0
+        ), f"{self.path} does not seem to have any files ending with '.jsonl'"
+        return jsonl_files
+
+
+@dataclass
+class DataFile:
+    path: Path
+    sample_type: SampleType
+
+    @property
+    def jsonl_files(self):
+        assert self.path.exists(), f"Make sure that {self.path} exists"
+        return [self.path]
+
+
+def parse_data_sources(
+    pretrain_data: str,
+    instruct_data: str,
+) -> Tuple[List[Union[DataDir, DataFile]], List[float]]:
+    seen: Set[str] = set()
+    sources: List[Union[DataDir, DataFile]] = []
+    weights: List[float] = []
+    for sample_sources, sample_type in [
+        (pretrain_data, SampleType.PRETRAIN),
+        (instruct_data, SampleType.INSTRUCT),
+    ]:
+        for source in sample_sources.strip().split(","):
+            if not source:
+                continue
+
+            source_items = source.strip().split(":")
+            if len(source_items) == 1:
+                path_ = source_items[0]
+                weight = 1.0
+            elif len(source_items) == 2:
+                path_, weight_ = source_items
+                weight = float(weight_)
+            else:
+                raise ValueError(
+                    f"{source} is not correctly formatted. Make sure to format each data source as <path/to/data>:<weight> or just <path/to/data>"
+                )
+
+            assert (
+                path_ not in seen
+            ), f"{path_} seems to be duplicated. Make sure to only add it once."
+            assert (
+                weight > 0
+            ), f"Make sure to define strictly positive data sampling weights, not {weight}"
+
+            data: Union[DataDir, DataFile]
+            if Path(path_).is_dir():
+                data = DataDir(path=Path(path_), sample_type=sample_type)
+            elif Path(path_).is_file():
+                data = DataFile(path=Path(path_), sample_type=sample_type)
+            else:
+                raise FileNotFoundError(
+                    f"The path {path_} does not exist. Make sure {path_} is either a file or directory that contains training data."
+                )
+
+            sources.append(data)
+            weights.append(weight)
+
+            seen.add(path_)
+
+    sum_weights = sum(weights)
+    n_weights = [weight / sum_weights for weight in weights]
+
+    assert min(n_weights) > 0
+    assert (
+        abs(1 - sum(n_weights)) < 1e-8
+    ), f"Defined data sampling weights {weights} must sum to 1."
+    return sources, n_weights
+
+
+@dataclasses.dataclass()
+class SequenceMaskAndSizes:
+    """
+    Concatenation of samples to reach a given size
+    """
+
+    x: List[int]
+    y: List[int]
+    mask: Mask
+    sizes: List[int]
+
+    def __post_init__(self):
+        assert sum(self.sizes) == len(self.x) == len(self.y) == len(self.mask)
+
+
+def sequence_iterator(
+    ds_it: Iterator[TokenSample],
+    seq_len: int,
+    is_finite: bool,
+) -> Iterator[SequenceMaskAndSizes]:
+    """
+    Creates sequences of length `seq_len` from the dataset iterator by concatenating samples.
+    """
+    x_buffer: List[int] = []
+    y_buffer: List[int] = []
+    mask_buffer: Mask = []
+
+    sizes: List[int] = []
+    n_missing = seq_len
+    for sample in ds_it:
+        assert 0 <= len(x_buffer) < seq_len, len(x_buffer)
+        assert n_missing == seq_len - len(
+            x_buffer
+        ), f"n_missing: {n_missing} | seq_len - len(x_buffer) {seq_len - len(x_buffer)}"
+
+        tokens, mask = sample.tokens, sample.masks[1:]
+        x, y = tokens[:-1], tokens[1:]
+        cur_pos = 0
+
+        while cur_pos < len(x):
+            size = len(x[cur_pos : cur_pos + n_missing])
+
+            curr_mask = mask[cur_pos : cur_pos + n_missing]
+            if not any(curr_mask):
+                cur_pos += size
+                # we have a sequence with a mask filled with False
+                continue
+
+            x_buffer.extend(x[cur_pos : cur_pos + n_missing])
+            y_buffer.extend(y[cur_pos : cur_pos + n_missing])
+            mask_buffer.extend(curr_mask)
+            n_missing -= size
+            sizes.append(size)
+
+            cur_pos += size
+
+            if n_missing == 0:
+                assert len(mask_buffer) == len(x_buffer) == seq_len == len(y_buffer)
+                assert sum(sizes) == seq_len
+                # we don't want to yield sequences with a mask filled with False
+                if any(mask_buffer):
+                    yield SequenceMaskAndSizes(
+                        x=x_buffer,
+                        y=y_buffer,
+                        mask=mask_buffer,
+                        sizes=sizes,
+                    )
+                x_buffer, y_buffer = [], []
+                mask_buffer = []
+                sizes = []
+                n_missing = seq_len
+
+    if is_finite:
+        # if dataloader is in eval, pad to seq length
+        if any(mask_buffer):
+            mask_buffer.extend(n_missing * [False])
+            x_buffer.extend(n_missing * [0])
+            y_buffer.extend(n_missing * [0])
+            sizes.append(n_missing)
+
+            yield SequenceMaskAndSizes(
+                x=x_buffer,
+                y=y_buffer,
+                mask=mask_buffer,
+                sizes=sizes,
+            )
+
+
+def build_dataset(
+    pretrain_data: str,
+    instruct_data: str,
+    instruct_args: InstructArgs,
+    instruct_tokenizer: InstructTokenizerBase,
+    seq_len: int,
+    seed: Optional[int],
+    rank: int,
+    world_size: int,
+    is_eval: bool,
+    shuffle_pretrain: bool = False,
+) -> Iterator[SequenceMaskAndSizes]:
+    sources, probabilities = parse_data_sources(
+        pretrain_data=pretrain_data, instruct_data=instruct_data
+    )
+
+    def do_shuffle(source: Union[DataDir, DataFile]) -> bool:
+        shuffle = {
+            SampleType.PRETRAIN: shuffle_pretrain,
+            SampleType.INSTRUCT: instruct_args.shuffle,
+        }[source.sample_type]
+
+        return not is_eval and shuffle
+
+    dataset_iterators = [
+        get_dataset_iterator(
+            source,
+            instruct_args=instruct_args,
+            instruct_tokenizer=instruct_tokenizer,
+            rank=rank,
+            world_size=world_size,
+            is_finite=is_eval,
+            seed=seed,
+            shuffle_at_epoch=do_shuffle(source),
+        )
+        for source in sources
+    ]
+
+    sequence_iterators = [
+        sequence_iterator(
+            ds_it=it,
+            seq_len=seq_len,
+            is_finite=is_eval,
+        )
+        for it in dataset_iterators
+    ]
+
+    if is_eval:
+        combined_iterator = itertools.chain.from_iterable(sequence_iterators)
+    else:
+        # make sure random_seed is different per rank and original seed
+        random_seed = np.array((seed, rank))
+        rng = np.random.RandomState(seed=random_seed)
+        combined_iterator = interleave_iterators(
+            sequence_iterators, probabilities=probabilities, rng=rng
+        )
+
+    return combined_iterator
+
+
+def get_rng(seed: int, rank: int) -> np.random.RandomState:
+    random_seed = np.array((seed, rank))
+    rng = np.random.RandomState(seed=random_seed)
+    return rng
+
+
+def get_dataset_iterator(
+    source: Union[DataDir, DataFile],
+    instruct_args: InstructArgs,
+    instruct_tokenizer: InstructTokenizerBase,
+    rank: int,
+    world_size: int,
+    is_finite: bool,
+    seed: Optional[int],
+    shuffle_at_epoch: bool,
+) -> Iterator[TokenSample]:
+    jsonl_files = source.jsonl_files
+    rng: Optional[np.random.RandomState] = (
+        get_rng(seed, rank) if seed is not None else None
+    )
+
+    chunk_dataset = (
+        instruct_args.dynamic_chunk_fn_call
+        and source.sample_type == SampleType.INSTRUCT
+    )
+
+    if not is_finite:
+        # train mode
+        while True:
+            for jsonl_file in jsonl_files:
+                if shuffle_at_epoch:
+                    assert rng is not None, "`seed` has to be passed when shuffling"
+                    # will preload all data into RAM, shuffle and yield
+                    yield from preload_and_yield(
+                        jsonl_file,
+                        chunk_dataset=chunk_dataset,
+                        rank=rank,
+                        world_size=world_size,
+                        rng=rng,
+                        instruct_tokenizer=instruct_tokenizer,
+                        sample_type=source.sample_type,
+                    )
+                else:
+                    # will read data on-the-fly and yield
+                    main_logger_info(f"Lazily loading {jsonl_file} ...")
+                    yield from lazy_load_and_yield(
+                        jsonl_file,
+                        rank=rank,
+                        world_size=world_size,
+                        instruct_tokenizer=instruct_tokenizer,
+                        sample_type=source.sample_type,
+                    )
+    else:
+        # eval mode
+        for jsonl_file in jsonl_files:
+            # No need to shuffle for eval
+            yield from lazy_load_and_yield(
+                jsonl_file,
+                rank=rank,
+                world_size=world_size,
+                instruct_tokenizer=instruct_tokenizer,
+                sample_type=source.sample_type,
+            )
+
+
+def preload_and_yield(
+    jsonl_file: Path,
+    chunk_dataset: bool,
+    rank: int,
+    world_size: int,
+    rng: np.random.RandomState,
+    instruct_tokenizer: InstructTokenizerBase,
+    sample_type: SampleType,
+) -> Iterator[TokenSample]:
+    # only instruct data has to be chunked
+    # load dataset if not already loaded. Make sure to only load 1/world_size dataset
+    tokens_list = maybe_load_local_dataset(
+        jsonl_file, chunk=chunk_dataset, rank=rank, world_size=world_size, instruct_tokenizer=instruct_tokenizer, sample_type=sample_type
+    )
+
+    if sample_type == SampleType.PRETRAIN:
+        assert chunk_dataset is False, "Pretrain data should not have chunking enabled."
+
+    main_logger_info(f"Shuffling {jsonl_file} ...")
+    rng.shuffle(tokens_list)
+
+    for token_sample in tokens_list:
+        yield token_sample
+
+def lazy_load_and_yield(
+    jsonl_file: Path,
+    rank: int,
+    world_size: int,
+    instruct_tokenizer: InstructTokenizerBase,
+    sample_type: SampleType,
+):
+    with jsonl_file.open() as file_handle:
+        for idx, line in enumerate(file_handle):
+            if not idx % world_size == rank:
+                continue
+
+            data = json.loads(line)
+            yield encode(
+                data,
+                instruct_tokenizer=instruct_tokenizer,
+                as_type=sample_type,
+            )
+
+
+def maybe_chunk_lines(lines: List[str]) -> List[str]:
+    extra_lines: List[str] = []
+    for line in lines:
+        data = json.loads(line)
+        # mult-turn fn call data will be chunked and shorder conversations are added additionally
+        maybe_chunked_lines = maybe_chunk_data(data)
+        extra_lines.extend([json.dumps(line) for line in maybe_chunked_lines])
+
+    return extra_lines
+
+
+def maybe_chunk_data(data: Dict[str, Any]) -> List[Dict[str, Any]]:
+    # think about always allowing both open-ai and non-open-ai data
+    sample = build_instruct_sample(data)
+
+    def num_assistant_messages(sample: TrainingInstructSample) -> int:
+        return len(
+            [m for m in sample.messages if isinstance(m, FinetuningAssistantMessage)]
+        )
+
+    chunk_data = []
+    while sample.only_last is True and num_assistant_messages(sample) > 1:
+        assert sample == build_instruct_sample(sample.dict())
+        last_message = sample.messages.pop()
+
+        # 1. First pop until and including last assistant message
+        system_message = None
+        while not isinstance(last_message, FinetuningAssistantMessage):
+            last_message = sample.messages.pop()
+            if isinstance(last_message, SystemMessage):
+                system_message = last_message
+
+        # 2. Second pop until and excluding last assistant message
+        prev_last_message = sample.messages[-1]
+        while not isinstance(prev_last_message, FinetuningAssistantMessage):
+            last_message = sample.messages.pop()
+            if isinstance(last_message, SystemMessage):
+                system_message = last_message
+
+            prev_last_message = sample.messages[-1]
+
+        # if system_message is not None, append again
+        if system_message is not None:
+            sample.messages.append(system_message)
+        chunk_data.append(sample.dict())
+
+    return chunk_data
+
+
+def interleave_iterators(iterators: List[Iterator], probabilities, rng):
+    while True:
+        it_id = rng.choice(range(len(iterators)), p=probabilities)
+        yield next(iterators[it_id])

finetune/data/exceptions.py

@@ -0,0 +1,56 @@
+class MessageFormatError(Exception):
+    def __init__(self, message, data):
+        self._message = message
+        self._begin_data = data[:20]
+        super().__init__()
+
+    def __str__(self):
+        return f"A message starting with {self._begin_data} is incorrectly formated." + self._message
+
+
+class ToolCallFormatError(Exception):
+    def __init__(self, message, data):
+        self._message = message
+        self._begin_data = data[:20]
+        super().__init__()
+
+    def __str__(self):
+        return f"A tool call assistant message starting with {self._begin_data} of the conversation is incorrectly formated. " + self._message
+
+
+class FunctionFormatError(Exception):
+    def __init__(self, message, data):
+        self._message = message
+        self._begin_data = data[:20]
+        super().__init__()
+
+    def __str__(self):
+        return (
+            f"A function of the conversation starting with {self._begin_data} is incorrectly formated. "
+            + self._message
+        )
+
+
+class ConversationFormatError(Exception):
+    def __init__(self, message, data):
+        self._message = message
+        self._begin_data = data[:20]
+        super().__init__()
+
+    def __str__(self):
+        return (
+            f"A conversation starting with {self._begin_data} is incorrectly formated. " + self._message
+        )
+
+
+class UnrecognizedRoleError(Exception):
+    def __init__(self, role, allowed_roles):
+        self._role = role
+        self._allowed_roles = allowed_roles
+        super().__init__()
+
+    def __str__(self):
+        return (
+            f"The following role: {self._role} is not recognized in line: {self.line} of the dataset {self.dataset}. Make sure that each role is one of {self._allowed_roles}"
+            + self._message
+        )

finetune/data/tokenize.py

@@ -0,0 +1,355 @@
+import logging
+from dataclasses import dataclass
+from enum import Enum
+from typing import Any, Dict, List, Optional, Union
+
+from mistral_common.protocol.instruct.messages import (
+    FinetuningAssistantMessage,
+    Roles,
+    SystemMessage,
+    ToolMessage,
+    UserMessage,
+)
+from mistral_common.protocol.instruct.tool_calls import (
+    Function,
+    FunctionCall,
+    Tool,
+    ToolCall,
+)
+from mistral_common.protocol.instruct.validator import (
+    MistralRequestValidatorV3,
+    ValidationMode,
+)
+from mistral_common.tokens.instruct.request import InstructRequest
+from mistral_common.tokens.tokenizers.base import Tokenizer
+from mistral_common.tokens.tokenizers.sentencepiece import InstructTokenizerBase
+
+from .exceptions import (
+    ConversationFormatError,
+    FunctionFormatError,
+    MessageFormatError,
+    ToolCallFormatError,
+    UnrecognizedRoleError,
+)
+
+logger = logging.getLogger("tokenize")
+
+Sequence = List[int]
+Mask = List[bool]
+
+
+class TrainingInstructSample(InstructRequest):
+    available_tools: Optional[List[Tool]] = None
+    only_last: bool = False
+
+
+@dataclass()
+class TokenSample:
+    tokens: Sequence
+    masks: Mask
+
+
+class SampleType(str, Enum):
+    PRETRAIN = "pretrain"
+    INSTRUCT = "instruct"
+
+
+def encode(
+    data: Dict[str, Any],
+    instruct_tokenizer: InstructTokenizerBase,
+    as_type: SampleType,
+) -> TokenSample:
+    sample: Union[str, TrainingInstructSample]
+    if as_type == SampleType.PRETRAIN:
+        sample = get_pretrain_sample(data)
+    elif as_type == SampleType.INSTRUCT:
+        sample = build_instruct_sample(data)
+
+    return tokenize(sample=sample, instruct_tokenizer=instruct_tokenizer)
+
+
+def get_pretrain_sample(data: Dict[str, Any]) -> str:
+    content_keys = ["text", "content"]
+    assert not all(
+        k in data for k in content_keys
+    ), "Make sure to have either 'text' or 'content' in your data. Not both."
+    assert any(
+        data.get(k) is not None for k in content_keys
+    ), f"Must have one of 'text' or 'content' in your data. Only have {data.keys()}"
+
+    # get first non-None value
+    sample = None
+    for key in content_keys:
+        sample = data[key] if key in data else sample
+
+    assert isinstance(sample, str), sample
+
+    return sample
+
+
+def build_instruct_sample(data: Dict[str, Any]) -> TrainingInstructSample:
+    messages: List[
+        SystemMessage | UserMessage | FinetuningAssistantMessage | ToolMessage
+    ] = []
+    # optional data fields that might be set
+    available_tools: Optional[List[Tool]] = data.get("available_tools")
+    system_prompt = data.get("system_prompt")
+
+    messages_keys = ["messages", "interactions"]
+    content_keys = ["content", "text"]  # both are accepted
+    allowed_roles = [role.value for role in Roles]
+
+    if not any(messages_key in data for messages_key in messages_keys):
+        err = f"The conversation does not contain one of '{', '.join(messages_keys)}' key, but only {', '.join(data.keys())}. Make sure that the conversation includes one of '{', '.join(messages_keys)}'."
+        raise ConversationFormatError(err, str(data))
+
+    if all(messages_key in data for messages_key in messages_keys):
+        err = f"The conversation cannot contain both of '{', '.join(messages_keys)}' key, but only one of the two."
+        raise ConversationFormatError(err, str(data))
+
+    # get first non-None value
+    data_messages: Optional[List[Dict[str, Any]]] = None
+    for key in messages_keys:
+        data_messages = data[key] if key in data else data_messages
+
+    assert data_messages is not None, "data_messages can't be None"
+
+    if "available_tools" in data and "tools" in data:
+        err = "The conversation contains both an `available_tools` and `tools` key. You can only have one."
+        raise ConversationFormatError(err, str(data))
+
+    if data.get("tools", None) is not None and len(data["tools"]) > 0:
+        available_tools = _parse_available_tools(data["tools"])
+    elif (
+        data.get("available_tools", None) is not None
+        and len(data["available_tools"]) > 0
+    ):
+        available_tools = _parse_available_tools(data["available_tools"])
+
+    for data_message in data_messages:
+        is_tool_call = data_message.get("tool_calls") is not None
+
+        if "role" not in data_message:
+            err = f"A message does not contain a 'role' key, but only {', '.join(data_message.keys())}. Make sure that the message includes the key 'role'."
+            raise MessageFormatError(err, str(data))
+
+        role = data_message["role"]
+
+        if all(key in data_message for key in content_keys):
+            err = f"A {role} message contains both a 'text' and 'content' key. Make sure that there is only one of the two."
+            raise MessageFormatError(err, str(data))
+
+        content: Optional[str] = None
+        for key in content_keys:
+            content = content if content is not None else data_message.get(key)
+
+        # non-function call message must have content
+        if not is_tool_call and content is None:
+            err = f"A {role} message does not contain one of '{content_keys}' key, but only {', '.join(data_message.keys())}. Make sure that the message includes one of '{content_keys}' keys."
+            raise MessageFormatError(err, str(data))
+
+        if role not in allowed_roles:
+            raise UnrecognizedRoleError(role, allowed_roles)
+
+        if data_message["role"] == "user":
+            assert content is not None
+            messages.append(UserMessage(content=content))
+        elif data_message["role"] == "assistant":
+            tool_calls: Optional[List[ToolCall]] = None
+
+            if is_tool_call:
+                tool_calls = _parse_tool_calls(data_message["tool_calls"])
+
+            weight = data_message.get("weight")
+            messages.append(
+                FinetuningAssistantMessage(
+                    content=content, tool_calls=tool_calls, weight=weight
+                )
+            )
+        elif data_message["role"] == "system":
+            if system_prompt is not None:
+                err = "Multiple messages with role 'system' encountered. Only one is allowed."
+                raise MessageFormatError(err, str(data))
+
+            system_prompt = content
+        elif data_message["role"] == "tool":
+            assert content is not None
+            tool_message = _parse_tool_message(content, data_message)
+            messages.append(tool_message)
+
+    # validate created messages
+    validator = MistralRequestValidatorV3(ValidationMode.finetuning)
+    validator.validate_messages(messages)
+    validator._validate_tools(available_tools or [])
+
+    # whether to train only on last assistant message
+    only_last = data.get("only_last", False) or available_tools is not None
+
+    return TrainingInstructSample(
+        messages=messages,
+        system_prompt=system_prompt,
+        available_tools=available_tools,
+        only_last=only_last,
+    )
+
+
+def _parse_available_tools(tools: List[Dict[str, Any]]) -> List[Tool]:
+    available_tools = []
+    for tool in tools:
+        if "function" not in tool:
+            raise FunctionFormatError(
+                "A tool dict does not have a 'function' key.", str(tool)
+            )
+
+        func_data = tool["function"]
+
+        for key in ["name", "description", "parameters"]:
+            if key not in func_data:
+                raise FunctionFormatError(
+                    f"A function dict does not have a {key} key.", str(func_data)
+                )
+
+        if not isinstance(func_data["parameters"], dict):
+            raise FunctionFormatError(
+                f"A function 'parameters' key has to be of type dict, but is {type(func_data['parameters'])}. If the function has no parameters pass an empyt dict ", str(func_data)
+            )
+
+        description = func_data["description"]
+        function = Function(
+            name=func_data["name"],
+            description=description,
+            parameters=func_data["parameters"],
+        )
+
+    available_tools.append(Tool(function=function))
+    return available_tools
+
+
+def _parse_tool_calls(calls: List[Dict[str, Any]]) -> List[ToolCall]:
+    for key in ["id", "function"]:
+        if not all(key in call for call in calls):
+            err = f"A tool call of an assistant message does not have a {key} key"
+            raise ToolCallFormatError(err, str(calls))
+
+    for key in ["name", "arguments"]:
+        if not all(key in call["function"] for call in calls):
+            err = (
+                f"A tool call function of an assistant message does not have a {key} key"
+            )
+            raise ToolCallFormatError(err, str(calls))
+
+    if not all(isinstance(call["function"]["arguments"], str) for call in calls):
+        err = "A tool call function of an assistant message does not have a 'arguments' key of type str"
+        raise ToolCallFormatError(err, str(calls))
+
+    tool_calls = [
+        ToolCall(
+            id=call["id"],
+            function=FunctionCall(
+                name=call["function"]["name"],
+                arguments=call["function"]["arguments"],
+            ),
+        )
+        for call in calls
+    ]
+    return tool_calls
+
+
+def _parse_tool_message(content: str, data_message: Dict[str, Any]) -> ToolMessage:
+    if "tool_call_id" not in data_message:
+        err = f"A tool message does not contain a 'tool_call_id' key, but only {', '.join(data_message.keys())}. Make sure that the message includes the key 'tool_call_id'."
+        raise MessageFormatError(err, str(data_message))
+
+    call_id = data_message["tool_call_id"]
+    # name is deprecated in v3, but we'll add it nevertheless for now
+    name = data_message.get("name")
+
+    return ToolMessage(content=content, tool_call_id=call_id, name=name)
+
+
+def tokenize(
+    sample: Union[str, TrainingInstructSample],
+    instruct_tokenizer: InstructTokenizerBase,
+) -> TokenSample:
+    if isinstance(sample, str):
+        tokenizer: Tokenizer = instruct_tokenizer.tokenizer
+        return tokenize_pretrain(sample, tokenizer)
+    elif isinstance(sample, TrainingInstructSample):
+        return tokenize_instruct(sample, instruct_tokenizer)
+
+    raise ValueError(
+        f"`sample` has to be either of type `str` or `TrainingInstructSample`, not {type(sample)}."
+    )
+
+
+def tokenize_pretrain(sample: str, tokenizer: Tokenizer) -> TokenSample:
+    tokens = tokenizer.encode(sample, bos=True, eos=True)
+    masks = [True] * len(tokens)
+    return TokenSample(tokens, masks)
+
+
+def tokenize_instruct(
+    sample: TrainingInstructSample,
+    instruct_tokenizer: InstructTokenizerBase,
+) -> TokenSample:
+    tokens: List[int] = instruct_tokenizer.start()
+    masks: List[bool] = [False]
+
+    mask_all_but_last = sample.only_last
+
+    # find first and last user message
+    user_messages = [
+        i for i, msg in enumerate(sample.messages) if isinstance(msg, UserMessage)
+    ]
+    first_user_idx = user_messages[0] if user_messages else -1
+    last_user_idx = user_messages[-1] if user_messages else -1
+
+    for msg_idx, message in enumerate(sample.messages):
+        if isinstance(message, UserMessage):
+            curr_tokens = instruct_tokenizer.encode_user_message(
+                message,
+                available_tools=sample.available_tools,
+                is_last=msg_idx == last_user_idx,
+                is_first=msg_idx == first_user_idx,
+                system_prompt=sample.system_prompt,
+            )
+            curr_masks = [False] * len(curr_tokens)  # only predict bot answers
+        elif isinstance(message, ToolMessage):
+            curr_tokens = instruct_tokenizer.encode_tool_message(
+                message, is_before_last_user_message=msg_idx < last_user_idx
+            )
+            curr_masks = [False] * len(curr_tokens)  # only predict bot answers
+        elif isinstance(message, FinetuningAssistantMessage):
+            is_last_message = msg_idx == (len(sample.messages) - 1)
+
+            # we don't want to predict a random call id
+            message = maybe_remove_call_id(message, is_last_message=is_last_message)
+
+            curr_tokens = instruct_tokenizer.encode_assistant_message(
+                message, is_before_last_user_message=False
+            )
+
+            is_weighted = message.weight is None or message.weight == 1
+            is_relevant = (not mask_all_but_last) or is_last_message
+            if is_weighted and is_relevant:
+                curr_masks = [True] * len(curr_tokens)  # only predict bot answers
+            else:
+                # in function calling we only backprop through last message
+                curr_masks = [False] * len(curr_tokens)
+
+        tokens.extend(curr_tokens)
+        masks.extend(curr_masks)
+
+    return TokenSample(tokens, masks)
+
+
+def maybe_remove_call_id(message: FinetuningAssistantMessage, is_last_message: bool):
+    if message.tool_calls is None or not is_last_message:
+        return message
+
+    # remove call id
+    message.tool_calls = [
+        ToolCall(function=call.function) for call in message.tool_calls
+    ]
+
+    return message

finetune/distributed.py

@@ -0,0 +1,59 @@
+import logging
+import os
+from functools import lru_cache
+from typing import List, Union
+
+import torch
+import torch.distributed as dist
+
+logger = logging.getLogger("distributed")
+
+BACKEND = "nccl"
+
+
+@lru_cache()
+def get_rank() -> int:
+    return dist.get_rank()
+
+
+@lru_cache()
+def get_world_size() -> int:
+    return dist.get_world_size()
+
+
+def visible_devices() -> List[int]:
+    return [int(d) for d in os.environ["CUDA_VISIBLE_DEVICES"].split(",")]
+
+
+def set_device():
+    logger.info(f"torch.cuda.device_count: {torch.cuda.device_count()}")
+    logger.info(f"CUDA_VISIBLE_DEVICES: {os.environ['CUDA_VISIBLE_DEVICES']}")
+    logger.info(f"local rank: {int(os.environ['LOCAL_RANK'])}")
+
+    assert torch.cuda.is_available()
+
+    assert len(visible_devices()) == torch.cuda.device_count()
+
+    if torch.cuda.device_count() == 1:
+        # gpus-per-task set to 1
+        torch.cuda.set_device(0)
+        return
+
+    local_rank = int(os.environ["LOCAL_RANK"])
+    logger.info(f"Set cuda device to {local_rank}")
+
+    assert 0 <= local_rank < torch.cuda.device_count(), (
+        local_rank,
+        torch.cuda.device_count(),
+    )
+    torch.cuda.set_device(local_rank)
+
+
+def avg_aggregate(metric: Union[float, int]) -> Union[float, int]:
+    buffer = torch.tensor([metric], dtype=torch.float32, device="cuda")
+    dist.all_reduce(buffer, op=dist.ReduceOp.SUM)
+    return buffer[0].item() / get_world_size()
+
+
+def is_torchrun() -> bool:
+    return "TORCHELASTIC_RESTART_COUNT" in os.environ

finetune/eval.py

@@ -0,0 +1,77 @@
+import logging
+from typing import List
+
+import numpy as np
+import torch.cuda
+import torch.distributed as dist
+from torch.distributed.fsdp.fully_sharded_data_parallel import FullyShardedDataParallel
+
+from .data.data_loader import Batch
+from .distributed import get_rank, get_world_size
+from .loss import compute_loss_with_mask
+from .utils import TrainState
+
+logger = logging.getLogger("eval")
+
+
+def main_logger_info(message: str) -> None:
+    if get_rank() == 0:
+        logger.info(message)
+
+
+def evaluate(
+    model: FullyShardedDataParallel,
+    batches: List[Batch],
+    state: TrainState,
+):
+    # Create fake samples to make FSDP happy for unbalanced data
+    num_samples = torch.tensor([len(batches)], device="cuda", dtype=torch.long)
+    all_num_samples = [torch.zeros_like(num_samples) for _ in range(get_world_size())]
+
+    torch.distributed.all_gather(all_num_samples, num_samples)
+
+    total_num_samples = int(torch.tensor(all_num_samples).sum().item())
+    max_num_samples = int(torch.tensor(all_num_samples).max().item())
+
+    for _ in range(max_num_samples - int(num_samples.item())):
+        pad_x = np.zeros_like(batches[-1].x)
+        pad_y = np.zeros_like(batches[-1].y)
+        pad_sizes = batches[-1].sizes.copy()
+
+        pad_batch = Batch(pad_x, pad_y, pad_sizes, is_pad_only=True)
+        batches.append(pad_batch)
+
+    # eval mode!
+    model.eval()
+
+    eval_loss = torch.tensor(0.0).cuda()
+    main_logger_info("Start eval...")
+    for batch in batches:
+        x = torch.from_numpy(batch.x).cuda()
+        y = torch.from_numpy(batch.y).cuda()
+        y_mask = (
+            torch.from_numpy(batch.y_mask).cuda() if batch.y_mask is not None else None
+        )
+
+        with torch.no_grad():
+            output = model(
+                input_ids=x,
+                seqlens=batch.sizes,
+            )
+
+            if y_mask.sum() > 0:
+                eval_loss += compute_loss_with_mask(output, y, y_mask)
+
+            assert batch.is_pad_only or y.abs().sum() != 0, "Pad sample is used to compute loss."
+
+    # sum loss
+    main_logger_info("Eval finished!")
+
+    dist.all_reduce(eval_loss, op=dist.ReduceOp.SUM)
+    eval_loss /= total_num_samples
+
+    state.this_eval_loss = eval_loss.item()
+    state.this_eval_perplexity = (2**eval_loss).item()
+
+    # train mode!
+    model.train()

finetune/loss.py

@@ -0,0 +1,16 @@
+from typing import Optional
+
+import torch
+from torch.nn import functional as F
+
+
+def compute_loss_with_mask(
+    logits: torch.Tensor, target: torch.Tensor, target_mask: Optional[torch.Tensor]
+):
+    if target_mask is None:
+        return F.cross_entropy(logits, target, reduction="mean")
+
+    mb_loss = F.cross_entropy(logits, target, reduction="none")
+    mb_loss = torch.sum(mb_loss * target_mask) / torch.sum(target_mask)
+
+    return mb_loss

finetune/mixed_precision.py

@@ -0,0 +1,47 @@
+from typing import Iterable
+
+import torch
+
+
+def prepare_mixed_precision(
+    params: Iterable[torch.nn.Parameter],
+    param_dtype: torch.dtype,
+    optim_dtype: torch.dtype,
+):
+    """Appends a freshly allocated fp32 tensor copy of all params to parameters that can be updated."""
+    with torch.no_grad():
+        for p in params:
+            if p.requires_grad:
+                # Mixed precision: let's save a fp32 param tensor to each params that require a grad
+                p._mp_param = torch.empty_like(p, dtype=optim_dtype)  # type: ignore
+                p._mp_param.copy_(p.to(optim_dtype))  # type: ignore
+
+            p.data = p.data.to(param_dtype)
+
+
+def upcast_mixed_precision(
+    params: Iterable[torch.nn.Parameter], optim_dtype: torch.dtype
+):
+    """Make sure to run this function BEFORE optimizer.step() so that all weights and optimizer states are updated in fp32 in .step()"""
+    with torch.no_grad():
+        for p in params:
+            if p.requires_grad and p.grad is not None:
+                # store original tensor in p._temp
+                p._temp = p.data  # type: ignore
+                # upcast data for the optimizer step
+                p.data = p._mp_param  # type: ignore
+                p.grad = p.grad.to(optim_dtype)
+
+
+def downcast_mixed_precision(
+    params: Iterable[torch.nn.Parameter], param_dtype: torch.dtype
+):
+    """Make sure to run this function AFTER optimizer.step() as optimizer.step() will update data underlying p.data and p._mp_param pointers"""
+    with torch.no_grad():
+        for p in params:
+            if p.requires_grad and p.grad is not None:
+                # copy fp32 weights into bfloat16 tensor
+                p._temp.copy_(p.data)  # type: ignore
+                # set _temp again to the data tensor
+                p.data = p._temp  # type: ignore
+                p.grad = p.grad.to(param_dtype)

finetune/monitoring/metrics_logger.py

@@ -0,0 +1,226 @@
+import json
+import logging
+import os
+from datetime import datetime, timedelta
+from pathlib import Path
+from typing import Any, Dict, Optional, Union
+
+from torch.utils.tensorboard import SummaryWriter
+
+from finetune.args import MLFlowArgs, TrainArgs, WandbArgs
+from finetune.utils import TrainState
+
+logger = logging.getLogger("metrics_logger")
+
+GB = 1024**3
+
+
+def get_train_logs(
+    state: TrainState,
+    loss: float,
+    lr: float,
+    peak_allocated_mem: float,
+    allocated_mem: float,
+    train_args: TrainArgs,
+) -> Dict[str, Union[float, int]]:
+    metrics = {
+        "lr": lr,
+        "step": state.step,
+        "loss": loss,
+        "percent_done": 100 * state.step / train_args.max_steps,
+        "peak_allocated_mem": peak_allocated_mem / GB,
+        "allocated_mem": allocated_mem / GB,
+        "wps": state.wps,
+        "avg_wps": state.avg_wps,
+        "eta_in_seconds": state.eta,
+    }
+
+    return metrics
+
+
+def get_eval_logs(
+    step: int,
+    train_loss: float,
+    perplexity: Optional[float],
+    eval_loss: Optional[float],
+) -> Dict[str, Union[float, int]]:
+    eval_dict = {"step": step, "train_loss": train_loss}
+
+    if perplexity is not None:
+        eval_dict["perplexity"] = perplexity
+
+    if eval_loss is not None:
+        eval_dict["eval_loss"] = eval_loss
+    return eval_dict
+
+
+def train_log_msg(
+    state: TrainState, logs: Dict[str, Union[float, int]], loss: float
+) -> str:
+    metrics: Dict[str, Union[float, int, datetime]] = dict(logs)  # shallow copy
+    metrics.pop("eta_in_seconds")
+
+    metrics["eta"] = datetime.now() + timedelta(seconds=state.eta)
+    metrics["step"] = state.step
+    metrics["loss"] = loss
+
+    parts = []
+    for key, fmt, new_name in [
+        ("step", "06", None),
+        ("percent_done", "03.1f", "done (%)"),
+        ("loss", ".3f", None),
+        ("lr", ".1e", None),
+        ("peak_allocated_mem", ".1f", "peak_alloc_mem (GB)"),
+        ("allocated_mem", ".1f", "alloc_mem (GB)"),
+        ("wps", ".1f", "words_per_second"),
+        ("avg_wps", ".1f", "avg_words_per_second"),
+        ("eta", "%Y-%m-%d %H:%M:%S", "ETA"),
+    ]:
+        name = key if new_name is None else new_name
+        try:
+            parts.append(f"{name}: {metrics[key]:>{fmt}}")
+        except KeyError:
+            logger.error(f"{key} not found in {sorted(metrics.keys())}")
+            raise
+
+    return " - ".join(parts)
+
+
+def eval_log_msg(logs: Dict[str, Union[float, int]]) -> str:
+    parts = []
+    for key, fmt, new_name in [
+        ("step", "06", None),
+        ("perplexity", ".3f", "eval_perplexity"),
+        ("eval_loss", ".3f", None),
+        ("train_loss", ".3f", None),
+    ]:
+        name = key if new_name is None else new_name
+        if key in logs:
+            parts.append(f"{name}: {logs[key]:>{fmt}}")
+
+    return " - ".join(parts)
+
+
+class MetricsLogger:
+    def __init__(
+        self,
+        dst_dir: Path,
+        tag: str,
+        is_master: bool,
+        wandb_args: WandbArgs,
+        mlflow_args: MLFlowArgs,
+        config: Optional[Dict[str, Any]] = None,
+    ):
+        self.dst_dir = dst_dir
+        self.tag = tag
+        self.is_master = is_master
+        self.jsonl_path = dst_dir / f"metrics.{tag}.jsonl"
+        self.tb_dir = dst_dir / "tb"
+        self.summary_writer: Optional[SummaryWriter] = None
+
+        if not self.is_master:
+            return
+
+        filename_suffix = f".{tag}"
+        self.tb_dir.mkdir(exist_ok=True)
+        self.summary_writer = SummaryWriter(
+            log_dir=str(self.tb_dir),
+            max_queue=1000,
+            filename_suffix=filename_suffix,
+        )
+        self.is_wandb = wandb_args.project is not None
+        self.is_mlflow = mlflow_args.tracking_uri is not None
+
+        if self.is_wandb:
+            import wandb
+
+            if wandb_args.key is not None:
+                wandb.login(key=wandb_args.key)  # LLM
+            if wandb_args.offline:
+                os.environ["WANDB_MODE"] = "offline"
+            if wandb.run is None:
+                logger.info("initializing wandb")
+                wandb.init(
+                    config=config,
+                    dir=dst_dir,
+                    project=wandb_args.project,
+                    job_type="training",
+                    name=wandb_args.run_name or dst_dir.name,
+                    resume=False,
+                )
+
+            self.wandb_log = wandb.log
+
+        if self.is_mlflow:
+            import mlflow
+
+            mlflow.set_tracking_uri(mlflow_args.tracking_uri)
+            mlflow.set_experiment(mlflow_args.experiment_name or dst_dir.name)
+
+            if tag == "train":
+                mlflow.start_run()
+
+            self.mlflow_log = mlflow.log_metric
+
+    def log(self, metrics: Dict[str, Union[float, int]], step: int):
+        if not self.is_master:
+            return
+
+        metrics_to_ignore = {"step"}
+        assert self.summary_writer is not None
+        for key, value in metrics.items():
+            if key in metrics_to_ignore:
+                continue
+            assert isinstance(value, (int, float)), (key, value)
+            self.summary_writer.add_scalar(
+                tag=f"{self.tag}.{key}", scalar_value=value, global_step=step
+            )
+
+            if self.is_mlflow:
+                self.mlflow_log(f"{self.tag}.{key}", value, step=step)
+
+        if self.is_wandb:
+            # grouping in wandb is done with /
+            self.wandb_log(
+                {
+                    f"{self.tag}/{key}": value
+                    for key, value in metrics.items()
+                    if key not in metrics_to_ignore
+                },
+                step=step,
+            )
+
+        metrics_: Dict[str, Any] = dict(metrics)  # shallow copy
+        if "step" in metrics_:
+            assert step == metrics_["step"]
+        else:
+            metrics_["step"] = step
+        metrics_["at"] = datetime.utcnow().isoformat()
+        with self.jsonl_path.open("a") as fp:
+            fp.write(f"{json.dumps(metrics_)}\n")
+
+    def close(self):
+        if not self.is_master:
+            return
+
+        if self.summary_writer is not None:
+            self.summary_writer.close()
+            self.summary_writer = None
+
+        if self.is_wandb:
+            import wandb
+
+            # to be sure we are not hanging while finishing
+            wandb.finish()
+
+        if self.is_mlflow:
+            import mlflow
+
+            mlflow.end_run()
+
+    def __del__(self):
+        if self.summary_writer is not None:
+            raise RuntimeError(
+                "MetricsLogger not closed properly! You should "
+                "make sure the close() method is called!"
+            )

finetune/monitoring/utils.py

@@ -0,0 +1,34 @@
+import datetime
+import logging
+import sys
+import time
+
+
+class DeltaTimeFormatter(logging.Formatter):
+    def format(self, record):
+        delta = datetime.timedelta(
+            seconds=int(record.relativeCreated / 1000)
+        )  # no milliseconds
+        record.delta = delta
+        return super().format(record)
+
+
+def set_logger(level: int = logging.INFO):
+    root = logging.getLogger()
+    root.handlers.clear()
+    root.setLevel(level)
+    tz, *_ = time.tzname
+
+    LOGFORMAT = "%(asctime)s - %(delta)s - %(name)s - %(levelname)s - %(message)s"
+    TIMEFORMAT = f"%Y-%m-%d %H:%M:%S ({tz})"
+    formatter = DeltaTimeFormatter(LOGFORMAT, TIMEFORMAT)
+
+    handler = logging.StreamHandler(sys.stdout)
+    handler.setLevel(level)
+    handler.setFormatter(formatter)
+    root.addHandler(handler)
+
+    handler = logging.StreamHandler(sys.stderr)
+    handler.setLevel(logging.WARNING)
+    handler.setFormatter(formatter)
+    root.addHandler(handler)

finetune/utils.py

@@ -0,0 +1,83 @@
+import contextlib
+import dataclasses
+import datetime
+import logging
+import time
+from typing import Optional, Protocol
+
+import torch
+
+logger = logging.getLogger("utils")
+
+
+@dataclasses.dataclass
+class TrainState:
+    max_steps: int
+    step: int = 0
+    elapsed_time: float = 0.0
+    n_seen_tokens: int = 0
+    this_step_time: float = 0.0
+    begin_step_time: float = 0.0
+    this_eval_perplexity: Optional[float] = None
+    this_eval_loss: Optional[float] = None
+
+    def start_step(self):
+        self.step += 1
+        self.begin_step_time = time.time()
+
+    def end_step(self, n_batch_tokens: int):
+        self.this_step_time = time.time() - self.begin_step_time
+        self.this_step_tokens = n_batch_tokens
+
+        self.elapsed_time += self.this_step_time
+        self.n_seen_tokens += self.this_step_tokens
+
+        self.begin_step_time = time.time()
+
+    @property
+    def wps(self):
+        return self.this_step_tokens / self.this_step_time
+
+    @property
+    def avg_wps(self):
+        return self.n_seen_tokens / self.elapsed_time
+
+    @property
+    def eta(self):
+        steps_left = self.max_steps - self.step
+        avg_time_per_step = self.elapsed_time / self.step
+
+        return steps_left * avg_time_per_step
+
+
+def set_random_seed(seed: int) -> None:
+    """Set random seed for reproducibility."""
+    torch.manual_seed(seed)
+    torch.cuda.manual_seed(seed)
+
+
+class Closable(Protocol):
+    def close(self):
+        pass
+
+
+@contextlib.contextmanager
+def logged_closing(thing: Closable, name: str):
+    """
+    Logging the closing to be sure something is not hanging at exit time
+    """
+    try:
+        setattr(thing, "wrapped_by_closing", True)
+        yield
+    finally:
+        logger.info(f"Closing: {name}")
+        try:
+            thing.close()
+        except Exception:
+            logger.error(f"Error while closing {name}!")
+            raise
+        logger.info(f"Closed: {name}")
+
+
+def now_as_str() -> str:
+    return datetime.datetime.now().strftime("%Y-%m-%d_%H-%M-%S")

finetune/wrapped_model.py

@@ -0,0 +1,227 @@
+import functools
+import json
+import logging
+import math
+from pathlib import Path
+from typing import Callable, Union
+
+import safetensors
+import torch
+import torch.distributed.fsdp.wrap as torch_wrap
+from torch.distributed.fsdp import BackwardPrefetch
+from torch.distributed.fsdp.api import ShardingStrategy
+from torch.distributed.fsdp.fully_sharded_data_parallel import FullyShardedDataParallel
+
+from model.args import ModelArgs, MoeArgs
+from model.transformer import Transformer, TransformerBlock
+
+from .args import LoraArgs
+from .checkpointing import Checkpointer
+from .distributed import (
+    get_rank,
+    get_world_size,
+)
+
+logger = logging.getLogger(__name__)
+
+
+def main_logger_info(message: str) -> None:
+    if get_rank() == 0:
+        logger.info(message)
+
+
+def get_fsdp_policy(is_lora: bool) -> Callable[[torch.nn.Module], bool]:
+    """
+    This function instantiates the FSDP wrap policy.
+    - Each Transformers block becomes it's own FSDP group so that only a single Transformer block is sharded at a time
+    - If LoRA is enabled, we additionally create seperate FSDP sub-groups for every trainable and non-trainable parameter group
+      since this is a requirement for mixed requires_grad=True/False training. See: https://pytorch.org/docs/stable/fsdp.html
+    """
+
+    # Each transformer block becomes a FSDP group, each being sharded seperately
+    transformer_block_wrap_policy = functools.partial(
+        torch_wrap.transformer_auto_wrap_policy,
+        transformer_layer_cls=(TransformerBlock,),
+    )
+
+    if not is_lora:
+        return transformer_block_wrap_policy
+
+    def fsdp_lora_policy_fn(module):
+        return all(p.requires_grad for p in module.parameters())
+
+    # For LoRA training, trainable and non-trainable parameters need to be put into
+    # different FSDP groups
+    fsdp_lora_policy = functools.partial(
+        torch_wrap.lambda_auto_wrap_policy, lambda_fn=fsdp_lora_policy_fn
+    )
+
+    policies = [fsdp_lora_policy, transformer_block_wrap_policy]
+
+    return functools.partial(torch_wrap._or_policy, policies=policies)
+
+
+def log_train_params(model: Union[torch.nn.Module, FullyShardedDataParallel]):
+    world_size = get_world_size()
+
+    num_params = world_size * sum(p.numel() for p in model.parameters())
+    num_train_params = world_size * sum(
+        p.numel() for p in model.parameters() if p.requires_grad
+    )
+
+    main_logger_info(
+        f"{num_train_params:,.0f} out of {num_params:,.0f} parameter are finetuned ({num_train_params / num_params * 100:.2f}%)."
+    )
+
+
+def initialize_lora_parameters(model: torch.nn.Module, param_dtype: torch.dtype):
+    """
+        Initialize LoRA layers with Kaiming uniform and zeros.
+        See original paper for more info: https://arxiv.org/abs/2106.09685 and
+        original github repo: https://github.com/microsoft/LoRA/blob/a0a92e0f26c067cf94747bdbf1ce73793fa44d19/loralib/layers.py#L122
+    """
+    for m_name, module in model.named_modules():
+        if all(p.is_meta for p in module.parameters()):
+            for p_name, param in module.named_parameters():
+                module._parameters[p_name] = torch.nn.Parameter(
+                    torch.empty_like(param, device="cpu", dtype=param_dtype)
+                )
+                param = module._parameters[p_name]
+
+                if m_name.split(".")[-1] == "lora_A":
+                    torch.nn.init.kaiming_uniform_(param, a=math.sqrt(5))
+                elif m_name.split(".")[-1] == "lora_B":
+                    torch.nn.init.zeros_(param)
+                else:
+                    raise ValueError(
+                        "Only Lora layers should be randomely initialized."
+                    )
+
+
+def load_model(
+    folder: Path,
+    lora: LoraArgs,
+    checkpoint: bool,
+    param_dtype: torch.dtype,
+) -> FullyShardedDataParallel:
+    with open(folder / "params.json", "r") as f:
+        args = json.loads(f.read())
+
+    model_args = ModelArgs(
+        lora=lora,
+        dim=args["dim"],
+        n_layers=args["n_layers"],
+        head_dim=args["head_dim"],
+        hidden_dim=args["hidden_dim"],
+        n_heads=args["n_heads"],
+        n_kv_heads=args["n_kv_heads"],
+        norm_eps=args["norm_eps"],
+        vocab_size=args["vocab_size"],
+    )
+
+    if model_args.vocab_size == 32000:
+        raise ValueError(
+            f"Fine-tuning is not supported for older model versions with vocab_size 32000. Make sure to extend your model to vocab_size=32768 using `python -m utils.extend_model_vocab --original_model_ckpt {folder} --extended_model_ckpt {folder}_extended`."
+        )
+
+    assert (
+        model_args.vocab_size >= 32768
+    ), "Make sure to use a model with a vocab size of at least 32768"
+
+    if args.get("rope_theta") is not None:
+        model_args.rope_theta = args["rope_theta"]
+
+    if args.get("moe") is not None:
+        model_args.moe = MoeArgs(**args["moe"])
+
+    with torch.device("meta"):
+        model = Transformer(args=model_args, checkpoint=checkpoint)
+
+    if get_rank() == 0:
+        state_dict = load_state_dict(folder, dtype=param_dtype)
+
+        model.load_state_dict(state_dict, assign=True)  # type: ignore
+        logger.info("Loaded model on cpu!")
+
+        if lora.enable:
+            logger.info("Initializing lora layers ...")
+            # initialize LoRA layers
+            initialize_lora_parameters(model, param_dtype)
+
+        assert not any(
+            p.is_meta for p in model.parameters()
+        ), "All parameters should be intialized by now"
+        assert all(
+            p.dtype == param_dtype for p in model.parameters()
+        ), f"All parameters should be on {param_dtype}"
+
+        logger.info("Finished initialization!")
+        param_init_fn = None
+    else:
+
+        def param_init_fn(m):
+            m.to_empty(device=torch.cuda.current_device(), recurse=False)
+            m.to(param_dtype)
+
+        assert all(
+            p.is_meta for p in model.parameters()
+        ), "All parameters should be on meta"
+
+    torch.distributed.barrier()
+
+    # only finetune LoRA parameters and freeze before wrapping
+    if lora.enable:
+        for name, param in model.named_parameters():
+            if "lora" in name:
+                param.requires_grad = True
+            else:
+                param.requires_grad = False
+
+    auto_wrap_policy = get_fsdp_policy(model_args.lora.enable)
+
+    main_logger_info(f"Sharding model over {get_world_size()} GPUs ...")
+
+    wrapped_model = FullyShardedDataParallel(
+        model,
+        sharding_strategy=ShardingStrategy.FULL_SHARD,
+        auto_wrap_policy=auto_wrap_policy,
+        backward_prefetch=BackwardPrefetch.BACKWARD_PRE,
+        limit_all_gathers=True,
+        device_id=torch.cuda.current_device(),
+        sync_module_states=True,
+        param_init_fn=param_init_fn,
+    )
+    main_logger_info("Model sharded!")
+
+    log_train_params(wrapped_model)
+
+    return wrapped_model
+
+
+@torch.no_grad()
+def load_state_dict(path: Path, dtype: torch.dtype):
+    assert path.is_dir(), path
+
+    this_safetensors_path = Checkpointer.consolidated_path(path, use_safetensors=True)
+    this_torch_path = Checkpointer.consolidated_path(path, use_safetensors=False)
+
+    assert (
+        this_safetensors_path.exists() or this_torch_path.exists()
+    ), f"Either {this_safetensors_path} or {this_torch_path} must exist."
+    assert not (
+        this_safetensors_path.exists() and this_torch_path.exists()
+    ), f"Only one of {this_safetensors_path} or {this_torch_path} should exist."
+
+    if this_safetensors_path.exists():
+        logger.info(f"Reloading model from {this_safetensors_path} ...")
+        model_state_dict = safetensors.torch.load_file(this_safetensors_path)
+    else:
+        logger.info(f"Reloading model from {this_torch_path} ...")
+        model_state_dict = torch.load(this_torch_path)
+
+    logger.info(f"Converting model to dtype {dtype} ...")
+
+    for k, v in model_state_dict.items():
+        model_state_dict[k] = v.to(dtype)
+
+    return model_state_dict

huggingface.ipynb

@@ -0,0 +1,40 @@
+{
+ "cells": [
+  {
+   "cell_type": "code",
+   "execution_count": null,
+   "metadata": {},
+   "outputs": [],
+   "source": [
+    "from huggingface_hub import HfApi, HfFolder, Repository\n",
+    "\n",
+    "repo_id = \"your_username/your_model_name\"\n",
+    "repo_local_path = \"./path_to_your_model\"\n",
+    "\n",
+    "# Create the repository object and clone the repo\n",
+    "repo = Repository(local_dir=repo_local_path, clone_from=repo_id)\n",
+    "\n",
+    "# Copy your model files to the repository\n",
+    "model_files = [\"config.json\", \"pytorch_model.bin\", \"tokenizer_config.json\", \"vocab.json\"]\n",
+    "for file in model_files:\n",
+    "    shutil.copy(file, repo_local_path)\n",
+    "\n",
+    "# Push the model files to the repository\n",
+    "repo.push_to_hub(commit_message=\"Initial model upload\")\n"
+   ]
+  }
+ ],
+ "metadata": {
+  "kernelspec": {
+   "display_name": "chemistralpy310",
+   "language": "python",
+   "name": "python3"
+  },
+  "language_info": {
+   "name": "python",
+   "version": "3.10.14"
+  }
+ },
+ "nbformat": 4,
+ "nbformat_minor": 2
+}