Update notebook for auto flash attention detection

finetune.ipynb

@@ -254,6 +254,9 @@
     "# This cannot be greater than model's max sequence length\n",
     "max_sequence_length = launch_parameters.max_length\n",
     "\n",
+    "# If to drop sequences that are longer than max_sequence_length\n",
+    "drop_long_sequences = False\n",
+    "\n",
     "# Batch size per GPU. \n",
     "# Increasing this will increase GPU memory requirement and training time\n",
     "micro_batch_size = launch_parameters.batch_size\n",
@@ -353,7 +356,8 @@
     "import torch\n",
     "\n",
     "# Mixed Precision Training. We automatically select the precision based on GPU capability\n",
-    "mixed_precision = \"bf16\" if torch.cuda.is_bf16_supported() else \"fp16\"\n",
+    "is_ampere_or_newer = torch.cuda.get_device_capability(device=0) >= (8, 0)\n",
+    "mixed_precision = \"bf16\" if is_ampere_or_newer and torch.cuda.is_bf16_supported() else \"fp16\"\n",
     "\n",
     "COMMAND = f\"\"\"\n",
     "accelerate launch \\\n",
@@ -362,7 +366,6 @@
     "train.py \\\n",
     "config-base.yaml \\\n",
     "--deepspeed ./deepspeed_configs/3_ds_z2_config.json \\\n",
-    "--flash_attention True \\\n",
     "--gradient_checkpointing unsloth \\\n",
     "--base_model {model_id} \\\n",
     "--output_dir {output_dir} \\\n",
@@ -372,6 +375,7 @@
     "--val_set_size {eval_size} \\\n",
     "--max_steps {max_steps} \\\n",
     "--sequence_len {max_sequence_length} \\\n",
+    "--drop_long_sequences {drop_long_sequences} \\\n",
     "--train_on_inputs False \\\n",
     "--sample_packing {sample_packing} \\\n",
     "--pad_to_sequence_len True \\\n",

notebook-requirements.txt

@@ -1,3 +1,3 @@
 -r requirements.txt
-jupyter-server-proxy==4.1.1
+jupyter-server-proxy==4.1.2
 jupyter_app_launcher @ git+https://github.com/truefoundry/jupyter_app_launcher@9a959b894542995fc763ed07324bbd274e96610d