Fix Google Colab notebook 2024-05 (#1662) [skip ci]

* include mlflow installation in the colab notebook

Without explicitly installing mlflow the `accelerate launch` command fails.

* update the colab noteboko to use the latest tinyllama config

examples/colab-notebooks/colab-axolotl-example.ipynb

@@ -46,7 +46,7 @@
     "!pip install torch==\"2.1.2\"\n",
     "!pip install -e git+https://github.com/OpenAccess-AI-Collective/axolotl#egg=axolotl\n",
     "!pip install flash-attn==\"2.5.0\"\n",
-        "!pip install deepspeed==\"0.13.1\""
+    "!pip install deepspeed==\"0.13.1\"!pip install mlflow==\"2.13.0\""
    ]
   },
   {
@@ -73,7 +73,6 @@
     "base_model: TinyLlama/TinyLlama-1.1B-intermediate-step-1431k-3T\n",
     "model_type: LlamaForCausalLM\n",
     "tokenizer_type: LlamaTokenizer\n",
-        "is_llama_derived_model: true\n",
     "\n",
     "load_in_8bit: false\n",
     "load_in_4bit: true\n",
@@ -89,8 +88,9 @@
     "adapter: qlora\n",
     "lora_model_dir:\n",
     "\n",
-        "sequence_len: 1096\n",
+    "sequence_len: 4096\n",
     "sample_packing: true\n",
+    "eval_sample_packing: false\n",
     "pad_to_sequence_len: true\n",
     "\n",
     "lora_r: 32\n",
@@ -106,20 +106,17 @@
     "wandb_name:\n",
     "wandb_log_model:\n",
     "\n",
-        "mlflow_experiment_name: colab-example\n",
-        "\n",
-        "gradient_accumulation_steps: 1\n",
-        "micro_batch_size: 1\n",
+    "gradient_accumulation_steps: 4\n",
+    "micro_batch_size: 2\n",
     "num_epochs: 4\n",
-        "max_steps: 20\n",
     "optimizer: paged_adamw_32bit\n",
     "lr_scheduler: cosine\n",
     "learning_rate: 0.0002\n",
     "\n",
     "train_on_inputs: false\n",
     "group_by_length: false\n",
-        "bf16: false\n",
-        "fp16: true\n",
+    "bf16: auto\n",
+    "fp16:\n",
     "tf32: false\n",
     "\n",
     "gradient_checkpointing: true\n",
@@ -128,11 +125,11 @@
     "local_rank:\n",
     "logging_steps: 1\n",
     "xformers_attention:\n",
-        "flash_attention: false\n",
+    "flash_attention: true\n",
     "\n",
     "warmup_steps: 10\n",
-        "evals_per_epoch:\n",
-        "saves_per_epoch:\n",
+    "evals_per_epoch: 4\n",
+    "saves_per_epoch: 1\n",
     "debug:\n",
     "deepspeed:\n",
     "weight_decay: 0.0\n",
@@ -204,13 +201,23 @@
    "provenance": []
   },
   "kernelspec": {
-      "display_name": "Python 3",
+   "display_name": "Python 3 (ipykernel)",
+   "language": "python",
    "name": "python3"
   },
   "language_info": {
-      "name": "python"
+   "codemirror_mode": {
+    "name": "ipython",
+    "version": 3
+   },
+   "file_extension": ".py",
+   "mimetype": "text/x-python",
+   "name": "python",
+   "nbconvert_exporter": "python",
+   "pygments_lexer": "ipython3",
+   "version": "3.12.1"
   }
  },
  "nbformat": 4,
-  "nbformat_minor": 0
+ "nbformat_minor": 4
 }