![\"Open](\"https://colab.research.google.com/assets/colab-badge.svg\"){kind=icon}
"
]
},
{
"cell_type": "markdown",
"source": [
"# Fine-tune Llama 2 in Google Colab\n",
"> 🗣️ Large Language Model Course\n",
"\n",
"❤️ Created by [@maximelabonne](), based on Younes Belkada's [GitHub Gist](https://gist.github.com/younesbelkada/9f7f75c94bdc1981c8ca5cc937d4a4da).\n",
"\n",
"This notebook runs on a T4 GPU with high RAM. (Last update: 30 Jul 2023)\n"
],
"metadata": {
"id": "OSHlAbqzDFDq"
}
},
{
"cell_type": "code",
"execution_count": null,
"metadata": {
"id": "GLXwJqbjtPho"
},
"outputs": [],
"source": [
"!pip install -q accelerate==0.21.0 peft==0.4.0 bitsandbytes==0.40.2 transformers==4.31.0 trl==0.4.7"
]
},
{
"cell_type": "code",
"source": [
"import os\n",
"import torch\n",
"from datasets import load_dataset\n",
"from transformers import (\n",
" AutoModelForCausalLM,\n",
" AutoTokenizer,\n",
" BitsAndBytesConfig,\n",
" HfArgumentParser,\n",
" TrainingArguments,\n",
" pipeline,\n",
" logging,\n",
")\n",
"from peft import LoraConfig, PeftModel\n",
"from trl import SFTTrainer"
],
"metadata": {
"id": "nAMzy_0FtaUZ"
},
"execution_count": null,
"outputs": []
},
{
"cell_type": "code",
"source": [
"# The model that you want to train from the Hugging Face hub\n",
"model_name = \"NousHermes/llama-2-7b-chat-hf\"\n",
"\n",
"# The instruction dataset to use\n",
"dataset_name = \"mlabonne/guanaco-llama2-1k\"\n",
"\n",
"# Fine-tuned model name\n",
"new_model = \"llama-2-7b-miniguanaco\"\n",
"\n",
"################################################################################\n",
"# QLoRA parameters\n",
"################################################################################\n",
"\n",
"# LoRA attention dimension\n",
"lora_r = 64\n",
"\n",
"# Alpha parameter for LoRA scaling\n",
"lora_alpha = 16\n",
"\n",
"# Dropout probability for LoRA layers\n",
"lora_dropout = 0.1\n",
"\n",
"################################################################################\n",
"# bitsandbytes parameters\n",
"################################################################################\n",
"\n",
"# Activate 4-bit precision base model loading\n",
"use_4bit = True\n",
"\n",
"# Compute dtype for 4-bit base models\n",
"bnb_4bit_compute_dtype = \"float16\"\n",
"\n",
"# Quantization type (fp4 or nf4)\n",
"bnb_4bit_quant_type = \"nf4\"\n",
"\n",
"# Activate nested quantization for 4-bit base models (double quantization)\n",
"use_nested_quant = False\n",
"\n",
"################################################################################\n",
"# TrainingArguments parameters\n",
"################################################################################\n",
"\n",
"# Output directory where the model predictions and checkpoints will be stored\n",
"output_dir = \"./results\"\n",
"\n",
"# Number of training epochs\n",
"num_train_epochs = 1\n",
"\n",
"# Enable fp16/bf16 training (set bf16 to True with an A100)\n",
"fp16 = False\n",
"bf16 = False\n",
"\n",
"# Batch size per GPU for training\n",
"per_device_train_batch_size = 4\n",
"\n",
"# Batch size per GPU for evaluation\n",
"per_device_eval_batch_size = 4\n",
"\n",
"# Number of update steps to accumulate the gradients for\n",
"gradient_accumulation_steps = 1\n",
"\n",
"# Enable gradient checkpointing\n",
"gradient_checkpointing = True\n",
"\n",
"# Maximum gradient normal (gradient clipping)\n",
"max_grad_norm = 0.3\n",
"\n",
"# Initial learning rate (AdamW optimizer)\n",
"learning_rate = 1e-5\n",
"\n",
"# Weight decay to apply to all layers except bias/LayerNorm weights\n",
"weight_decay = 0.001\n",
"\n",
"# Optimizer to use\n",
"optim = \"paged_adamw_32bit\"\n",
"\n",
"# Learning rate schedule\n",
"lr_scheduler_type = \"cosine\"\n",
"\n",
"# Number of training steps (overrides num_train_epochs)\n",
"max_steps = -1\n",
"\n",
"# Ratio of steps for a linear warmup (from 0 to learning rate)\n",
"warmup_ratio = 0.03\n",
"\n",
"# Group sequences into batches with same length\n",
"# Saves memory and speeds up training considerably\n",
"group_by_length = True\n",
"\n",
"# Save checkpoint every X updates steps\n",
"save_steps = 0\n",
"\n",
"# Log every X updates steps\n",
"logging_steps = 25\n",
"\n",
"################################################################################\n",
"# SFT parameters\n",
"################################################################################\n",
"\n",
"# Maximum sequence length to use\n",
"max_seq_length = None\n",
"\n",
"# Pack multiple short examples in the same input sequence to increase efficiency\n",
"packing = False\n",
"\n",
"# Load the entire model on the GPU 0\n",
"device_map = {\"\": 0}"
],
"metadata": {
"id": "ib_We3NLtj2E"
},
"execution_count": null,
"outputs": []
},
{
"cell_type": "code",
"source": [
"# Load dataset (you can process it here)\n",
"dataset = load_dataset(dataset_name, split=\"train\")\n",
"\n",
"# Load tokenizer and model with QLoRA configuration\n",
"compute_dtype = getattr(torch, bnb_4bit_compute_dtype)\n",
"\n",
"bnb_config = BitsAndBytesConfig(\n",
" load_in_4bit=use_4bit,\n",
" bnb_4bit_quant_type=bnb_4bit_quant_type,\n",
" bnb_4bit_compute_dtype=compute_dtype,\n",
" bnb_4bit_use_double_quant=use_nested_quant,\n",
")\n",
"\n",
"# Check GPU compatibility with bfloat16\n",
"if compute_dtype == torch.float16 and use_4bit:\n",
" major, _ = torch.cuda.get_device_capability()\n",
" if major >= 8:\n",
" print(\"=\" * 80)\n",
" print(\"Your GPU supports bfloat16: accelerate training with bf16=True\")\n",
" print(\"=\" * 80)\n",
"\n",
"# Load base model\n",
"model = AutoModelForCausalLM.from_pretrained(\n",
" model_name,\n",
" quantization_config=bnb_config,\n",
" device_map=device_map\n",
")\n",
"model.config.use_cache = False\n",
"model.config.pretraining_tp = 1\n",
"\n",
"# Load LLaMA tokenizer\n",
"tokenizer = AutoTokenizer.from_pretrained(model_name, trust_remote_code=True)\n",
"tokenizer.pad_token = tokenizer.eos_token\n",
"tokenizer.padding_side = \"right\" # Fix weird overflow issue with fp16 training\n",
"\n",
"# Load LoRA configuration\n",
"peft_config = LoraConfig(\n",
" lora_alpha=lora_alpha,\n",
" lora_dropout=lora_dropout,\n",
" r=lora_r,\n",
" bias=\"none\",\n",
" task_type=\"CAUSAL_LM\",\n",
")\n",
"\n",
"# Set training parameters\n",
"training_arguments = TrainingArguments(\n",
" output_dir=output_dir,\n",
" num_train_epochs=num_train_epochs,\n",
" per_device_train_batch_size=per_device_train_batch_size,\n",
" gradient_accumulation_steps=gradient_accumulation_steps,\n",
" optim=optim,\n",
" save_steps=save_steps,\n",
" logging_steps=logging_steps,\n",
" learning_rate=learning_rate,\n",
" weight_decay=weight_decay,\n",
" fp16=fp16,\n",
" bf16=bf16,\n",
" max_grad_norm=max_grad_norm,\n",
" max_steps=max_steps,\n",
" warmup_ratio=warmup_ratio,\n",
" group_by_length=group_by_length,\n",
" lr_scheduler_type=lr_scheduler_type,\n",
" report_to=\"tensorboard\"\n",
")\n",
"\n",
"# Set supervised fine-tuning parameters\n",
"trainer = SFTTrainer(\n",
" model=model,\n",
" train_dataset=dataset,\n",
" peft_config=peft_config,\n",
" dataset_text_field=\"text\",\n",
" max_seq_length=max_seq_length,\n",
" tokenizer=tokenizer,\n",
" args=training_arguments,\n",
" packing=packing,\n",
")\n",
"\n",
"# Train model\n",
"trainer.train()\n",
"\n",
"# Save trained model\n",
"trainer.model.save_pretrained(new_model)"
],
"metadata": {
"id": "OJXpOgBFuSrc",
"colab": {
"base_uri": "https://localhost:8080/",
"height": 1000,
"referenced_widgets": [
"721399b15bb54fbf96b30cdfbd745d67",
"51b9aeef1c484d5e8157098f12e2b80e",
"df84af09ed494acfbe0c35a4b09be9ea",
"e5eabfa5f9ec4bb78a768cc0023fa7d6",
"ddeb064c99eb4c98ba13ca9d5614eb38",
"fd4af61367b34412be06b99ff73664e5",
"b70661cf67d54816a921985efd0296f3",
"376c32de2a6a408eb7d1f04a20d93e31",
"356dcfe1bd55462bbcd5bdfea8f88ded",
"75602449f24a4fe08798529993cd6dd9",
"1920aa1224764e13b6850d47e5011405",
"75565360e4bf4e3dafa2d38eb093f7c5",
"fbc22ad1b318493b86351c7669c67364",
"a43cfcd604d04478ac54fffbb0272cda",
"3e26777495204c0794bd15e67f068468",
"98f13286addf4cd7b28345c29d17f2a0",
"954cc945033d4824beefc340f9b1a9b0",
"f989a371eb0f481fb325c41e501e8cb3",
"c0b13bc5448f43e98497b3255b874946",
"5f61953aa4ac4a439d3b66989e077be5",
"81a79dbb2ff54a05a00172087cd42498",
"842208929edb467a85083ec6e23cbb81",
"48c160ccc64b4a3b8a2cba89f1ebf240",
"7e711c9917134054a233b38f696dbe49",
"0d0fcdad38294753badbe11cffd0e151",
"7e74ce97b71140dbbd85cf00dc7492aa",
"5d107a70134a4f3c84cf4b73e16d5fe2",
"a86fcef000ba46ddb7f2b442500d062d",
"d25cc83962074b038b35c5de9f73dca5",
"2cd8a6810af849159b68a81e4a9f1725",
"9d297f07606a4fd88a96f36583ff80ce",
"4b7429d78b534a17961b00d91c42904a",
"ef33893275a84e1c9817be5b2add635d",
"e230554ebab348f0981b5baf0363a903",
"135203f78ff24d05a8db58cf718faab6",
"8f245e60bb1d47ed9d52df310609f836",
"7854bb526fbc443da830613751b6314d",
"73a132907275438f93039c9e561c89a6",
"dcb43471c4a747d186a173cf24babc0e",
"3573701cab144e2c94cc3898522ed3ba",
"e078fc5b071048c88c601012346b8a6e",
"66c20be0eb9e405bb2209f45d09bdb3c",
"b4f135c27df2414d80bd4918ebfffe57",
"1c15a2e4d4834e40964d828c66f3e4b3",
"ab6b34e2db854b18bb7780989815b0ab",
"ebabcad7ea0e4b69a3c7d0ada4b26bae",
"d533fcf22a8f4a1bb140aee6e0899ba5",
"43c68091402e4f97814a96a9c8032383",
"3105439ed9c64cf191b43d1f8ae255af",
"95e8aeb7d70b441f90b3cf5e302b5f61",
"12a688de04a9431694d2617ded1b9223",
"639868f764894a5d9be140354868251b",
"ade266560ff74d0a9e40a99d06953053",
"977db907e3ec492b88972372a33694f4",
"2abfd6fa9dc04e3388f6fff6d80b41d7",
"2b039f7b746a45369e5d851614175bbc",
"3cd6bb7087bc47c98796639d4cff7bb3",
"07238381ecd1479da7e963e8fb8a91ec",
"f0543285053e451e95f7778803a84b70",
"b24cfa8bc11548c8aa6f6add241c5eb7",
"3ca72bcb3a284a6c9f3602b2b0b6388f",
"7faf0bc6810445b581c9287d25b08e47",
"a4df012ff43d481d896872f632f1b5b2",
"6998be1c589640f7b284e2151edf8d4e",
"ef855e09789440fcb2d9e2fbe9874836",
"dd0bc49e48e841bd9facff437d87ca8f",
"f5cf90ecc0ef45f7a9136b52f6acc785",
"e1b2181c04fa4343af127756baee6757",
"9dc472303b7f486e8d403405a60f57f5",
"a9ac60f5010947b0b7143ae6a0e6e3b8",
"7cdcdf80508841f0ac20b66e63f64423",
"e570736c249e499e96f697a9b300efe1",
"65c221c962664344b8ec274c7289fada",
"e0379e76d58e4bcf92a56be5731bf307",
"9526d2a549404ea7a7c6df3a3327683e",
"efdfd35123ad4833b2f59759bf05fe06",
"d0e87897a938489caef358ac0a53d101",
"43fd6a31e1d54018a389e2a4287e270f",
"cd52aee35064492d99d6dc5bc1328aab",
"e3c18d78f48a4edabfc08033abac504a",
"a59aab32b6c14610b396a02da210c3a6",
"c1d05eb2a2b34dbdbe91753f65d55cf0",
"3a3670b48c3346aab4b3345db62a6ba0",
"ce7ab3917f1f4879b87a204bb9e53eaf",
"301940e95a2c487fbe30d36534291554",
"bf69301dc5f3456195fdb864ea5bd6ba",
"aed04ad4eb064357855cc6ea1b9e3786",
"73bcdeae0b3c4780838e69773680a7ca",
"31411a98a2bd4a5896c451046005f307",
"e2c08fc80c01430b881a9c01a86dbe70",
"566c1be3fcc24524bcdf3527c39b38b7",
"f2794ccb8f6e4e5681d188789887167a",
"58effb6659c24cd681b6a81223708ec4",
"4f17dd57c8e34d3f9be864d66ea7060c",
"2ec2cb89b86b479a8eb2fb5dee0d5751",
"5d4f102ca6c44192bb702426cae8d8db",
"e596fdad6c714b2596fc9c376a7b9f8b",
"bfcbe40a637447d388f0b1ab6461794c",
"fe42120088f946b992eaca504b73d9a1",
"c0dacdb08b2348c1970187ca2c0f9bd5",
"b225ca1fd1d94f5698116993800725a0",
"b2e5e869506f436983741349d8a43d6a",
"196295d82e96496f9b4a5def1b5fe166",
"10b8da377bce4b79a3bb6c71edfb65e9",
"9224829e54e04f7d98fe1a1a5c18e1b3",
"b48295cc6af94b27b5c0c5a69909b86f",
"64525e54b15944ff9143c6a358fe61b3",
"31bbff2160b1449a86493311bbf278fe",
"695f46a469fb401ea9b8904d32cd2072",
"6c269b9b047b4cad84e213a51caf8fec",
"0d8b9d12a18d4b67bc10934a6a8c7765",
"85c9184551ba4ec8ade6655b089ae9d5",
"4c777d36137d4687a4511241096d3828",
"2c2cdd0b468a4ee2a327f2a23219da90",
"83076a1d1f4d495e9d41f77d2f02e120",
"9ba00f5f2a9e4ed5976d28096bdb48f7",
"874a08b9948a4b0898f4a564cf188097",
"579dc3fbdbbb4cee9e75300304c22b46",
"cb234f57b6a6445cbaa3d9aed02b7f10",
"4bdd3198a0c3429ba17f474e477f63e1",
"a357f57266f549e2959c123f0583c9f6",
"867bdade5af74b99b8b97ef5da9b4ae6",
"ac0327a4bcac4c81b8ec9ade2eb640be",
"b5e16baac1224bddb2aaec9378366da9",
"9bd34ecdf0634a89b2eafc6beb1c917f",
"e4a6189749234b3e91ca46faef60646e",
"228fa4bcaba84a1ea899b4825ce5d942",
"64ac789db1f04dbe891d274779c8c190",
"88434bfdfc974e988d05b1e1e97ede32",
"10ff1fb0713946ec86727d073602bbcc",
"12018acbdaec430fa80d25abb8bbce85",
"c07fb58c80e04b1c959954425a4398c3",
"6d5a3219be21454eb68c93556355250c",
"6bc5ec464c164f36891cf768dc8f41c0",
"e28c9aa01646434a969be4cd82ed2507",
"9088bcc3c8174f8b841df3adea7c6015",
"68efa6c26a05437aae10709bed1e74f7",
"13e6248a25be4c2aa8b4da3ac18f3b6f",
"76601f0b77f04db181c02f7f0325f0fc",
"7250f29cf9c547cdb7ff9a301d6d9df6",
"c47e7303700c4e448d45ddd85965bf4c",
"51ad2d170407408fb0b7d9244045ec5e",
"70ffe5adb54f48398b5a0feac3875118",
"f591ade1aa5a4dcf8f672721bfdcd60f",
"24859e5319304f82a45a6950e485932e",
"97af72dff95846f4940d8430a84d5c75",
"3a5ff03511124ccfbd69838ea58f8712",
"c460031571464d96a10404c723988cb5",
"e21cb6fd80f749d5a289e59cf4fdff07",
"7037f03347c546e7b18a54409e19685f",
"f9ef7139260d45df9540b5a720507b60",
"67d8ea99e1834096a480bb8fef83298e",
"bc617d2b349544228f1a838d7dc2b641",
"f478281283fd45d89234d2a031a6bbff",
"e3a08b4c37c64aedaf3326caae0a731f",
"fb72964872354303a30087965d781f7a",
"42e97ad6a01841d8906af61aea1d29d1",
"35eab2a8f3ad4393be74df9f1754ef85",
"d00c372e2b8343deb61729b5f3a8dfb5",
"05248306adb24e94bf290327c7baf35a",
"c8d7eb2318bc4744b03e34a891fba887",
"70bbd89b244644b4a9ec2ce2f579f8e1",
"9f424d2d6180423eb8a5792d8fe00ef0",
"94995f77750e4be6800569c9af23740f",
"5a8c65c76a0140dfb9c54d197e3090de",
"b9330b4a429440d1be69aea1b5d19077",
"6fcc368bdf234cb4bb4db998f793e902",
"43bd0adf678448c8ab6ce498b3a39d38",
"81a6ebccfaa247d79746dd86100ee72f",
"00e8c2da722b4505b5a33c97fe277973",
"7d1bcf2847dd40d084a9612cc703d39f",
"19f7fafff5284847b167039aab8eeaf2",
"b34a8f400bd74069a49e170bb97c76f8",
"64aed29227b04702b80b21ce90d9aa64",
"fa7a87e7467641f2b2d7e71f5183d715",
"c0caea34cfd24568840cfb6da9cf0145",
"9fee3103bb5f4722828b1254f8642162",
"661fada8f57e449abb65a525f9085f27",
"b31161d089ce4f8794ba72c6580f0a50",
"e9f1d5714057409c86bd3695c2232a29",
"c9dc39195bf04c7cb6755b4195bcba3a",
"e00fba1fd87e4bcf9cd2e49a7133a7bd",
"91d7de910b67440387660223c5f222ea",
"66f809dd6a2645e6b5b839b6a36ab22d",
"0d2f087bbe444457a91b6f71ad589912",
"0899df69633c4c1b9aea3d89be739bc2",
"fa7e2dbc7cf34a1f998ee0a9289d8ed3"
]
},
"outputId": "a29901b3-5257-45a7-983c-f3f5f4eb7b94"
},
"execution_count": null,
"outputs": [
{
"output_type": "display_data",
"data": {
"text/plain": [
"Downloading readme: 0%| | 0.00/782 [00:00"
],
"text/html": [
"\n",
" | Step | \n", "Training Loss | \n", "
|---|---|
| 25 | \n", "1.353500 | \n", "
| 50 | \n", "1.625300 | \n", "
| 75 | \n", "1.208700 | \n", "
| 100 | \n", "1.438300 | \n", "
| 125 | \n", "1.176000 | \n", "
| 150 | \n", "1.358800 | \n", "
| 175 | \n", "1.171700 | \n", "
| 200 | \n", "1.455300 | \n", "
| 225 | \n", "1.153900 | \n", "
| 250 | \n", "1.528100 | \n", "
"
]
},
"metadata": {}
}
]
},
{
"cell_type": "code",
"source": [
"# %load_ext tensorboard\n",
"# %tensorboard --logdir results/runs"
],
"metadata": {
"id": "crj9svNe4hU5"
},
"execution_count": null,
"outputs": []
},
{
"cell_type": "code",
"source": [
"# Ignore warnings\n",
"logging.set_verbosity(logging.CRITICAL)\n",
"\n",
"# Run text generation pipeline with our next model\n",
"prompt = \"What is a large language model?\"\n",
"pipe = pipeline(task=\"text-generation\", model=model, tokenizer=tokenizer, max_length=200)\n",
"result = pipe(f\"[INST] {prompt} [/INST]\")\n",
"print(result[0]['generated_text'])"
],
"metadata": {
"id": "frlSLPin4IJ4",
"colab": {
"base_uri": "https://localhost:8080/"
},
"outputId": "757676f4-1d4c-4777-e25b-3c92d37c51d8"
},
"execution_count": null,
"outputs": [
{
"output_type": "stream",
"name": "stderr",
"text": [
"/usr/local/lib/python3.10/dist-packages/transformers/generation/utils.py:1270: UserWarning: You have modified the pretrained model configuration to control generation. This is a deprecated strategy to control generation and will be removed soon, in a future version. Please use a generation configuration file (see https://huggingface.co/docs/transformers/main_classes/text_generation )\n",
" warnings.warn(\n",
"/usr/local/lib/python3.10/dist-packages/torch/utils/checkpoint.py:31: UserWarning: None of the inputs have requires_grad=True. Gradients will be None\n",
" warnings.warn(\"None of the inputs have requires_grad=True. Gradients will be None\")\n"
]
},
{
"output_type": "stream",
"name": "stdout",
"text": [
"[INST] What is a large language model? [/INST] A large language model is a type of artificial intelligence that is trained on a large dataset of text to generate human-like language. It is typically trained on a dataset of text that is much larger than the dataset used for smaller language models. The large dataset allows the model to learn more complex patterns in language, which can result in more accurate and natural-sounding language generation.\n",
"\n",
"Large language models are often used for tasks such as text generation, language translation, and question answering. They are also used in applications such as chatbots, virtual assistants, and language translation software.\n",
"\n",
"Some examples of large language models include:\n",
"\n",
"* BERT (Bidirectional Encoder Representations from Transformers): A large language model developed by Google that is trained on a large dataset of text and is used for a wide range of natural language processing tasks.\n",
"* RoBERTa (\n"
]
}
]
},
{
"cell_type": "markdown",
"source": [
"There is a problem with the VRAM here despite `del model` and emptying the VRAM. You probably need to restart the runtime (Runtime > Restart runtime), re-execute the three first cells, and then execute this one. Please contact me if you have a fix!"
],
"metadata": {
"id": "6WjzALHtSfdb"
}
},
{
"cell_type": "code",
"source": [
"# Reload model in FP16 and merge it with LoRA weights\n",
"base_model = AutoModelForCausalLM.from_pretrained(\n",
" model_name,\n",
" low_cpu_mem_usage=True,\n",
" return_dict=True,\n",
" torch_dtype=torch.float16,\n",
" device_map=device_map,\n",
")\n",
"model = PeftModel.from_pretrained(base_model, new_model)\n",
"model = model.merge_and_unload()\n",
"\n",
"# Reload tokenizer to save it\n",
"tokenizer = AutoTokenizer.from_pretrained(model_name, trust_remote_code=True)\n",
"tokenizer.pad_token = tokenizer.eos_token\n",
"tokenizer.padding_side = \"right\""
],
"metadata": {
"id": "QQn30cRtAZ-P",
"colab": {
"base_uri": "https://localhost:8080/",
"height": 49,
"referenced_widgets": [
"5251de6649944ecd910c588ca62ba875",
"987f74c748744a2ba606d1b9e1ae3ae2",
"f5624f09a7524de38a021464516b745d",
"5875d29480704bf298cf8563a1687bae",
"47df5b8d0212476d87d86fd49d9a791d",
"8185344c705949aba78c604e38ad98da",
"c455d30e801a41818b1d05d322d5bd08",
"1297c38157f24ac5819c95b1b88dc616",
"3094677e36f441cd8944cddbce5be7ae",
"c290a1d5effb41df9156d6324f88fde5",
"31702ac6175649c592f23f1443ca986c"
]
},
"outputId": "463dab09-f8e9-4d99-ac87-19df4a51dbb5"
},
"execution_count": null,
"outputs": [
{
"output_type": "display_data",
"data": {
"text/plain": [
"Loading checkpoint shards: 0%| | 0/2 [00:00