File size: 42,425 Bytes
60616b8
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
1
2
3
4
5
6
7
8
9
10
11
12
13
14
15
16
17
18
19
20
21
22
23
24
25
26
27
28
29
30
31
32
33
34
35
36
37
38
39
40
41
42
43
44
45
46
47
48
49
50
51
52
53
54
55
56
57
58
59
60
61
62
63
64
65
66
67
68
69
70
71
72
73
74
75
76
77
78
79
80
81
82
83
84
85
86
87
88
89
90
91
92
93
94
95
96
97
98
99
100
101
102
103
104
105
106
107
108
109
110
111
112
113
114
115
116
117
118
119
120
121
122
123
124
125
126
127
128
129
130
131
132
133
134
135
136
137
138
139
140
141
142
143
144
145
146
147
148
149
150
151
152
153
154
155
156
157
158
159
160
161
162
163
164
165
166
167
168
169
170
171
172
173
174
175
176
177
178
179
180
181
182
183
184
185
186
187
188
189
190
191
192
193
194
195
196
197
198
199
200
201
202
203
204
205
206
207
208
209
210
211
212
213
214
215
216
217
218
219
220
221
222
223
224
225
226
227
228
229
230
231
232
233
234
235
236
237
238
239
240
241
242
243
244
245
246
247
248
249
250
251
252
253
254
255
256
257
258
259
260
261
262
263
264
265
266
267
268
269
270
271
272
273
274
275
276
277
278
279
280
281
282
283
284
285
286
287
288
289
290
291
292
293
294
295
296
297
298
299
300
301
302
303
304
305
306
307
308
309
310
311
312
313
314
315
316
317
318
319
320
321
322
323
324
325
326
327
328
329
330
331
332
333
334
335
336
337
338
339
340
341
342
343
344
345
346
347
348
349
350
351
352
353
354
355
356
357
358
359
360
361
362
363
364
365
366
367
368
369
370
371
372
373
374
375
376
377
378
379
380
381
382
383
384
385
386
387
388
389
390
391
392
393
394
395
396
397
398
399
400
401
402
403
404
405
406
407
408
409
410
411
412
413
414
415
416
417
418
419
420
421
422
423
424
425
426
427
428
429
430
431
432
433
434
435
436
437
438
439
440
441
442
443
444
445
446
447
448
449
450
451
452
453
454
455
456
457
458
459
460
461
462
463
464
465
466
467
468
469
470
471
472
473
474
475
476
477
478
479
480
481
482
483
484
485
486
487
488
489
490
491
492
493
494
495
496
497
498
499
500
501
502
503
504
505
506
507
508
509
510
511
512
513
514
515
516
517
518
519
520
521
522
523
524
525
526
527
528
529
530
531
532
533
534
535
536
537
538
539
540
541
542
543
544
545
546
547
548
549
550
551
552
553
554
555
556
557
558
559
560
561
562
563
564
565
566
567
568
569
570
571
572
573
574
575
576
577
578
579
580
581
582
583
584
585
586
587
588
589
590
591
592
593
594
595
596
597
598
599
600
601
602
603
604
605
606
607
608
609
610
611
612
613
614
615
616
617
618
619
620
621
622
623
624
625
626
627
628
629
630
631
632
633
634
635
636
637
638
639
640
641
642
643
644
645
646
647
648
649
650
651
652
653
654
655
656
657
658
659
660
661
662
{
 "cells": [
  {
   "cell_type": "code",
   "execution_count": 3,
   "metadata": {},
   "outputs": [
    {
     "data": {
      "text/plain": [
       "True"
      ]
     },
     "execution_count": 3,
     "metadata": {},
     "output_type": "execute_result"
    }
   ],
   "source": [
    "import torch\n",
    "torch.cuda.is_available()"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 4,
   "metadata": {},
   "outputs": [],
   "source": [
    "import glob\n",
    "import math\n",
    "import sys\n",
    "import time\n",
    "from pathlib import Path\n",
    "from typing import Optional, Tuple, Union\n",
    "\n",
    "import lightning as L\n",
    "import torch\n",
    "from lightning.fabric.loggers import CSVLogger\n",
    "from lightning.fabric.strategies import FSDPStrategy\n",
    "from torch.utils.data import DataLoader\n",
    "\n",
    "# # support running without installing as a package\n",
    "# wd = Path(__file__).parent.parent.resolve()\n",
    "# sys.path.append(str(wd))\n",
    "\n",
    "from tsai_gpt.model import GPT, Block, Config\n",
    "from tsai_gpt.packed_dataset import CombinedDataset, PackedDataset\n",
    "from tsai_gpt.speed_monitor import SpeedMonitorBase, estimate_flops, measure_flops\n",
    "from tsai_gpt.speed_monitor import SpeedMonitorFabric as SpeedMonitor\n",
    "from tsai_gpt.utils import chunked_cross_entropy, get_default_supported_precision, num_parameters, load_checkpoint"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 5,
   "metadata": {},
   "outputs": [],
   "source": [
    "model_name = \"pythia-160m\"\n",
    "name = \"redpajama\"\n",
    "out_dir = Path(\"out\") / name\n",
    "save_interval = 1000\n",
    "eval_interval = 1000\n",
    "eval_iters = 100\n",
    "log_interval = 100"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 6,
   "metadata": {},
   "outputs": [],
   "source": [
    "# Hyperparameters\n",
    "learning_rate = 6e-3\n",
    "batch_size = 32\n",
    "micro_batch_size = 8\n",
    "gradient_accumulation_steps = batch_size // micro_batch_size\n",
    "assert gradient_accumulation_steps > 0\n",
    "#max_iters = 600000  # num_epochs * (epoch_size // micro_batch_size) // devices\n",
    "max_iters = 15000\n",
    "weight_decay = 1e-1\n",
    "beta1 = 0.9\n",
    "beta2 = 0.95\n",
    "grad_clip = 1.0\n",
    "decay_lr = True\n",
    "warmup_iters = 2000\n",
    "lr_decay_iters = max_iters\n",
    "min_lr = 6e-6"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 7,
   "metadata": {},
   "outputs": [],
   "source": [
    "# Data proportions from https://arxiv.org/pdf/2302.13971.pdf Table 1\n",
    "data_config = [\n",
    "    (\"arxiv\", 2.5),\n",
    "    (\"book\", 4.5),\n",
    "    (\"c4\", 15.0),\n",
    "    (\"cc\", 67.0),\n",
    "    (\"github\", 4.5),\n",
    "    (\"stackexchange\", 2.0),\n",
    "    (\"wikipedia\", 4.5),\n",
    "]"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 8,
   "metadata": {},
   "outputs": [],
   "source": [
    "hparams = {k: v for k, v in locals().items() if isinstance(v, (int, float, str)) and not k.startswith(\"_\")}\n",
    "logger = CSVLogger(\"out\", name, flush_logs_every_n_steps=log_interval)\n",
    "\n",
    "\n",
    "def setup(\n",
    "    devices: int = 4,\n",
    "    train_data_dir: Path = Path(\"data/redpajama_sample\"),\n",
    "    val_data_dir: Optional[Path] = None,\n",
    "    precision: Optional[str] = None,\n",
    "    resume: Union[bool, Path] = False,\n",
    ") -> None:\n",
    "    precision = precision or get_default_supported_precision(training=True)\n",
    "\n",
    "    if devices > 1:\n",
    "        strategy = FSDPStrategy(\n",
    "            auto_wrap_policy={Block},\n",
    "            activation_checkpointing_policy={Block},\n",
    "            state_dict_type=\"full\",\n",
    "            limit_all_gathers=True,\n",
    "            cpu_offload=False,\n",
    "        )\n",
    "    else:\n",
    "        strategy = \"auto\"\n",
    "\n",
    "    fabric = L.Fabric(devices=devices, strategy=strategy, precision=precision, loggers=logger)\n",
    "    fabric.print(hparams)\n",
    "    fabric.launch(main, train_data_dir, val_data_dir, resume)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 8,
   "metadata": {},
   "outputs": [],
   "source": [
    "model_copy = None"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 9,
   "metadata": {},
   "outputs": [],
   "source": [
    "def main(fabric: L.Fabric, train_data_dir: Path, val_data_dir: Path, resume: Union[bool, Path]) -> None:\n",
    "    global model_copy\n",
    "    speed_monitor = SpeedMonitor(fabric, window_size=50, time_unit=\"seconds\")\n",
    "\n",
    "    if fabric.global_rank == 0:\n",
    "        out_dir.mkdir(parents=True, exist_ok=True)\n",
    "\n",
    "    config = Config.from_name(model_name)\n",
    "\n",
    "    train_dataloader, val_dataloader = create_dataloaders(\n",
    "        batch_size=micro_batch_size,\n",
    "        block_size=config.block_size,\n",
    "        fabric=fabric,\n",
    "        train_data_dir=train_data_dir,\n",
    "        val_data_dir=val_data_dir,\n",
    "        seed=(1337 + fabric.global_rank),\n",
    "    )\n",
    "    if val_dataloader is None:\n",
    "        train_dataloader = fabric.setup_dataloaders(train_dataloader)\n",
    "    else:\n",
    "        train_dataloader, val_dataloader = fabric.setup_dataloaders(train_dataloader, val_dataloader)\n",
    "\n",
    "    fabric.seed_everything(1337)  # same seed for every process to init model (FSDP)\n",
    "\n",
    "    fabric.print(f\"Loading model with {config.__dict__}\")\n",
    "    t0 = time.perf_counter()\n",
    "    import torch\n",
    "    import torch.nn as nn\n",
    "    def _init_weights(module: nn.Module) -> None:\n",
    "            \"\"\"Meant to be used with `gpt.apply(gpt._init_weights)`.\"\"\"\n",
    "            if isinstance(module, nn.Linear):\n",
    "                torch.nn.init.normal_(module.weight, mean=0.0, std=0.02)\n",
    "                if module.bias is not None:\n",
    "                    torch.nn.init.zeros_(module.bias)\n",
    "            elif isinstance(module, nn.Embedding):\n",
    "                torch.nn.init.normal_(module.weight, mean=0.0, std=0.02)\n",
    "            \n",
    "    with fabric.init_module(empty_init=True):\n",
    "        model = GPT(config)\n",
    "        model.apply(_init_weights)\n",
    "    model.apply(_init_weights)\n",
    "\n",
    "    \n",
    "    # checkpoint_path = Path(\"out/redpajama/iter-000999-ckpt.pth\")\n",
    "\n",
    "    # load_checkpoint(fabric, model, checkpoint_path)\n",
    "        \n",
    "    # print(model.transformer.h[0].mlp.fc.weight)\n",
    "\n",
    "    fabric.print(f\"Time to instantiate model: {time.perf_counter() - t0:.02f} seconds.\")\n",
    "    fabric.print(f\"Total parameters {num_parameters(model):,}\")\n",
    "\n",
    "    model = fabric.setup(model)\n",
    "    optimizer = torch.optim.AdamW(\n",
    "        model.parameters(), lr=learning_rate, weight_decay=weight_decay, betas=(beta1, beta2), foreach=False\n",
    "    )\n",
    "\n",
    "    # model_copy = model\n",
    "\n",
    "    optimizer = fabric.setup_optimizers(optimizer)\n",
    "\n",
    "    state = {\"model\": model, \"optimizer\": optimizer, \"hparams\": hparams, \"iter_num\": 0, \"step_count\": 0}\n",
    "\n",
    "    if resume is True:\n",
    "        resume = max(out_dir.glob(\"*.pth\"), key=lambda p: int(p.name.split(\"-\")[1]))\n",
    "    if resume:\n",
    "        fabric.print(f\"Resuming training from {resume}\")\n",
    "        fabric.load(resume, state)\n",
    "\n",
    "    train_time = time.perf_counter()\n",
    "    train(fabric, state, train_dataloader, val_dataloader, speed_monitor)\n",
    "    fabric.print(f\"Training time: {(time.perf_counter()-train_time):.2f}s\")\n",
    "    if fabric.device.type == \"cuda\":\n",
    "        fabric.print(f\"Memory used: {torch.cuda.max_memory_allocated() / 1e9:.02f} GB\")\n",
    "\n"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 10,
   "metadata": {},
   "outputs": [],
   "source": [
    "def train(\n",
    "    fabric: L.Fabric,\n",
    "    state: dict,\n",
    "    train_dataloader: DataLoader,\n",
    "    val_dataloader: DataLoader,\n",
    "    speed_monitor: SpeedMonitorBase,\n",
    ") -> None:\n",
    "    model = state[\"model\"]\n",
    "    optimizer = state[\"optimizer\"]\n",
    "\n",
    "    if val_dataloader is not None:\n",
    "        validate(fabric, model, val_dataloader)  # sanity check\n",
    "\n",
    "    with torch.device(\"meta\"):\n",
    "        meta_model = GPT(model.config)\n",
    "        # \"estimated\" is not as precise as \"measured\". Estimated is optimistic but widely used in the wild.\n",
    "        # When comparing MFU or FLOP numbers with other projects that use estimated FLOPs,\n",
    "        # consider passing `SpeedMonitor(flops_per_batch=estimated_flops)` instead\n",
    "        estimated_flops = estimate_flops(meta_model) * micro_batch_size\n",
    "        fabric.print(f\"Estimated TFLOPs: {estimated_flops * fabric.world_size / 1e12:.2f}\")\n",
    "        x = torch.randint(0, 1, (micro_batch_size, model.max_seq_length))\n",
    "        measured_flops = measure_flops(meta_model, x)\n",
    "        fabric.print(f\"Measured TFLOPs: {measured_flops * fabric.world_size / 1e12:.2f}\")\n",
    "        del meta_model, x\n",
    "\n",
    "    total_lengths = 0\n",
    "    total_t0 = time.perf_counter()\n",
    "\n",
    "    for state[\"iter_num\"], train_data in enumerate(train_dataloader, state[\"iter_num\"]):\n",
    "        if state[\"iter_num\"] >= max_iters:\n",
    "            checkpoint_path = out_dir / f\"iter-{state['iter_num']:06d}-ckpt.pth\"\n",
    "            fabric.print(f\"Saving checkpoint to {str(checkpoint_path)!r}\")\n",
    "            fabric.save(checkpoint_path, state)\n",
    "            break\n",
    "\n",
    "        # determine and set the learning rate for this iteration\n",
    "        lr = get_lr(state[\"iter_num\"]) if decay_lr else learning_rate\n",
    "        for param_group in optimizer.param_groups:\n",
    "            param_group[\"lr\"] = lr\n",
    "\n",
    "        iter_t0 = time.perf_counter()\n",
    "\n",
    "        input_ids = train_data[:, 0 : model.max_seq_length].contiguous()\n",
    "        targets = train_data[:, 1 : model.max_seq_length + 1].contiguous()\n",
    "\n",
    "        is_accumulating = (state[\"iter_num\"] + 1) % gradient_accumulation_steps != 0\n",
    "        with fabric.no_backward_sync(model, enabled=is_accumulating):\n",
    "            logits = model(input_ids)\n",
    "            loss = chunked_cross_entropy(logits, targets, chunk_size=0)\n",
    "            fabric.backward(loss / gradient_accumulation_steps)\n",
    "        \n",
    "        # return \n",
    "\n",
    "        if not is_accumulating:\n",
    "            fabric.clip_gradients(model, optimizer, max_norm=grad_clip)\n",
    "            optimizer.step()\n",
    "            optimizer.zero_grad()\n",
    "            state[\"step_count\"] += 1\n",
    "\n",
    "        t1 = time.perf_counter()\n",
    "        total_lengths += input_ids.size(1)\n",
    "        speed_monitor.on_train_batch_end(\n",
    "            (state[\"iter_num\"] + 1) * micro_batch_size,\n",
    "            t1 - total_t0,\n",
    "            # this assumes that device FLOPs are the same and that all devices have the same batch size\n",
    "            fabric.world_size,\n",
    "            flops_per_batch=measured_flops,\n",
    "            lengths=total_lengths,\n",
    "        )\n",
    "        if state[\"iter_num\"] % log_interval == 0:\n",
    "            fabric.print(\n",
    "                f\"iter {state['iter_num']} step {state['step_count']}: loss {loss.item():.4f}, LR: {lr:.6f}, iter time:\"\n",
    "                f\" {(t1 - iter_t0) * 1000:.2f}ms{' (optimizer.step)' if not is_accumulating else ''}\"\n",
    "            )\n",
    "\n",
    "        if val_dataloader is not None and not is_accumulating and state[\"step_count\"] % eval_interval == 0:\n",
    "            t0 = time.perf_counter()\n",
    "            val_loss = validate(fabric, model, val_dataloader)\n",
    "            t1 = time.perf_counter() - t0\n",
    "            speed_monitor.eval_end(t1)\n",
    "            fabric.print(f\"step {state['iter_num']}: val loss {val_loss.item():.4f}, val time: {t1 * 1000:.2f}ms\")\n",
    "            fabric.barrier()\n",
    "        if not is_accumulating and state[\"step_count\"] % save_interval == 0:\n",
    "            checkpoint_path = out_dir / f\"iter-{state['iter_num']:06d}-ckpt.pth\"\n",
    "            fabric.print(f\"Saving checkpoint to {str(checkpoint_path)!r}\")\n",
    "            fabric.save(checkpoint_path, state)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 11,
   "metadata": {},
   "outputs": [],
   "source": [
    "@torch.inference_mode()\n",
    "def validate(fabric: L.Fabric, model: torch.nn.Module, val_dataloader: DataLoader) -> torch.Tensor:\n",
    "    fabric.print(\"Validating ...\")\n",
    "    model.eval()\n",
    "\n",
    "    losses = torch.zeros(eval_iters, device=fabric.device)\n",
    "    for k, val_data in enumerate(val_dataloader):\n",
    "        input_ids = val_data[:, 0 : model.max_seq_length].contiguous()\n",
    "        targets = val_data[:, 1 : model.max_seq_length + 1].contiguous()\n",
    "        logits = model(input_ids)\n",
    "        losses[k] = chunked_cross_entropy(logits, targets, chunk_size=0)\n",
    "    out = losses.mean()\n",
    "\n",
    "    model.train()\n",
    "    return out"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 12,
   "metadata": {},
   "outputs": [],
   "source": [
    "def create_dataloader(\n",
    "    batch_size: int, block_size: int, data_dir: Path, fabric: L.Fabric, shuffle: bool = True, seed: int = 12345\n",
    ") -> DataLoader:\n",
    "    datasets = []\n",
    "    for prefix, _ in data_config:\n",
    "        filenames = glob.glob(str(data_dir / f\"{prefix}*\"))\n",
    "        dataset = PackedDataset(\n",
    "            filenames,\n",
    "            n_chunks=4,\n",
    "            block_size=block_size,\n",
    "            shuffle=shuffle,\n",
    "            seed=seed,\n",
    "            num_processes=fabric.world_size,\n",
    "            process_rank=fabric.global_rank,\n",
    "        )\n",
    "        datasets.append(dataset)\n",
    "\n",
    "    if not datasets:\n",
    "        raise RuntimeError(\n",
    "            f\"No data found at {data_dir}. Make sure you ran prepare_redpajama.py to create the dataset.\"\n",
    "        )\n",
    "\n",
    "    weights = [weight for _, weight in data_config]\n",
    "    sum_weights = sum(weights)\n",
    "    weights = [el / sum_weights for el in weights]\n",
    "\n",
    "    combined_dataset = CombinedDataset(datasets=datasets, seed=seed, weights=weights)\n",
    "\n",
    "    return DataLoader(combined_dataset, batch_size=batch_size, shuffle=False, pin_memory=True)\n"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 13,
   "metadata": {},
   "outputs": [],
   "source": [
    "def create_dataloaders(\n",
    "    batch_size: int,\n",
    "    block_size: int,\n",
    "    fabric: L.Fabric,\n",
    "    train_data_dir: Path = Path(\"data/redpajama_sample\"),\n",
    "    val_data_dir: Optional[Path] = None,\n",
    "    seed: int = 12345,\n",
    ") -> Tuple[DataLoader, DataLoader]:\n",
    "    # Increase by one because we need the next word as well\n",
    "    effective_block_size = block_size + 1\n",
    "    train_dataloader = create_dataloader(\n",
    "        batch_size=batch_size,\n",
    "        block_size=effective_block_size,\n",
    "        fabric=fabric,\n",
    "        data_dir=train_data_dir,\n",
    "        shuffle=True,\n",
    "        seed=seed,\n",
    "    )\n",
    "    val_dataloader = (\n",
    "        create_dataloader(\n",
    "            batch_size=batch_size,\n",
    "            block_size=effective_block_size,\n",
    "            fabric=fabric,\n",
    "            data_dir=val_data_dir,\n",
    "            shuffle=False,\n",
    "            seed=seed,\n",
    "        )\n",
    "        if val_data_dir\n",
    "        else None\n",
    "    )\n",
    "    return train_dataloader, val_dataloader"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 14,
   "metadata": {},
   "outputs": [],
   "source": [
    "def get_lr(it: int) -> float:\n",
    "    # 1) linear warmup for warmup_iters steps\n",
    "    if it < warmup_iters:\n",
    "        return learning_rate * it / warmup_iters\n",
    "    # 2) if it > lr_decay_iters, return min learning rate\n",
    "    if it > lr_decay_iters:\n",
    "        return min_lr\n",
    "    # 3) in between, use cosine decay down to min learning rate\n",
    "    decay_ratio = (it - warmup_iters) / (lr_decay_iters - warmup_iters)\n",
    "    assert 0 <= decay_ratio <= 1\n",
    "    coeff = 0.5 * (1.0 + math.cos(math.pi * decay_ratio))  # coeff ranges 0..1\n",
    "    return min_lr + coeff * (learning_rate - min_lr)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 16,
   "metadata": {},
   "outputs": [
    {
     "name": "stderr",
     "output_type": "stream",
     "text": [
      "Using bfloat16 Automatic Mixed Precision (AMP)\n",
      "Seed set to 1337\n"
     ]
    },
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "{'model_name': 'pythia-160m', 'name': 'redpajama', 'save_interval': 1000, 'eval_interval': 1000, 'eval_iters': 100, 'log_interval': 100, 'learning_rate': 0.006, 'batch_size': 32, 'micro_batch_size': 8, 'gradient_accumulation_steps': 4, 'max_iters': 15000, 'weight_decay': 0.1, 'beta1': 0.9, 'beta2': 0.95, 'grad_clip': 1.0, 'decay_lr': True, 'warmup_iters': 2000, 'lr_decay_iters': 15000, 'min_lr': 6e-06}\n",
      "Loading model with {'name': 'pythia-160m', 'hf_config': {'org': 'EleutherAI', 'name': 'pythia-160m-deduped'}, 'block_size': 2048, 'vocab_size': 50254, 'padding_multiple': 128, 'padded_vocab_size': 50304, 'n_layer': 12, 'n_head': 12, 'n_embd': 768, 'rotary_percentage': 0.25, 'parallel_residual': True, 'bias': True, 'lm_head_bias': False, 'n_query_groups': 12, 'shared_attention_norm': False, '_norm_class': 'LayerNorm', 'norm_eps': 1e-05, '_mlp_class': 'GptNeoxMLP', 'gelu_approximate': 'none', 'intermediate_size': 3072, 'rope_condense_ratio': 1, 'rope_base': 10000, 'head_size': 64, 'rope_n_elem': 16}\n",
      "Time to instantiate model: 0.03 seconds.\n",
      "Total parameters 162,322,944\n",
      "Estimated TFLOPs: 22.14\n",
      "Measured TFLOPs: 15.86\n",
      "iter 0 step 0: loss 10.9790, LR: 0.000000, iter time: 993.87ms\n",
      "iter 100 step 25: loss 7.3972, LR: 0.000300, iter time: 56.71ms\n",
      "iter 200 step 50: loss 5.9952, LR: 0.000600, iter time: 50.39ms\n",
      "iter 300 step 75: loss 5.8594, LR: 0.000900, iter time: 50.64ms\n",
      "iter 400 step 100: loss 6.0047, LR: 0.001200, iter time: 56.22ms\n",
      "iter 500 step 125: loss 5.9611, LR: 0.001500, iter time: 55.67ms\n",
      "iter 600 step 150: loss 5.7425, LR: 0.001800, iter time: 56.04ms\n",
      "iter 700 step 175: loss 5.6345, LR: 0.002100, iter time: 56.77ms\n",
      "iter 800 step 200: loss 5.4736, LR: 0.002400, iter time: 56.59ms\n",
      "iter 900 step 225: loss 5.3942, LR: 0.002700, iter time: 55.94ms\n",
      "iter 1000 step 250: loss 5.3758, LR: 0.003000, iter time: 55.97ms\n",
      "iter 1100 step 275: loss 5.4347, LR: 0.003300, iter time: 50.85ms\n",
      "iter 1200 step 300: loss 5.6140, LR: 0.003600, iter time: 55.73ms\n",
      "iter 1300 step 325: loss 5.1840, LR: 0.003900, iter time: 50.78ms\n",
      "iter 1400 step 350: loss 5.5744, LR: 0.004200, iter time: 55.55ms\n",
      "iter 1500 step 375: loss 4.8744, LR: 0.004500, iter time: 56.58ms\n",
      "iter 1600 step 400: loss 5.2784, LR: 0.004800, iter time: 56.40ms\n",
      "iter 1700 step 425: loss 4.6915, LR: 0.005100, iter time: 56.13ms\n",
      "iter 1800 step 450: loss 5.1188, LR: 0.005400, iter time: 51.62ms\n",
      "iter 1900 step 475: loss 5.1236, LR: 0.005700, iter time: 54.36ms\n",
      "iter 2000 step 500: loss 4.8566, LR: 0.006000, iter time: 56.13ms\n",
      "iter 2100 step 525: loss 4.7733, LR: 0.005999, iter time: 56.35ms\n",
      "iter 2200 step 550: loss 4.8778, LR: 0.005997, iter time: 50.88ms\n",
      "iter 2300 step 575: loss 4.4962, LR: 0.005992, iter time: 56.58ms\n",
      "iter 2400 step 600: loss 4.7976, LR: 0.005986, iter time: 55.87ms\n",
      "iter 2500 step 625: loss 4.3805, LR: 0.005978, iter time: 57.95ms\n",
      "iter 2600 step 650: loss 4.5287, LR: 0.005969, iter time: 59.25ms\n",
      "iter 2700 step 675: loss 4.3517, LR: 0.005957, iter time: 52.44ms\n",
      "iter 2800 step 700: loss 4.6585, LR: 0.005944, iter time: 51.33ms\n",
      "iter 2900 step 725: loss 4.4335, LR: 0.005929, iter time: 58.50ms\n",
      "iter 3000 step 750: loss 4.4874, LR: 0.005913, iter time: 57.96ms\n",
      "iter 3100 step 775: loss 4.3373, LR: 0.005895, iter time: 57.91ms\n",
      "iter 3200 step 800: loss 4.1922, LR: 0.005875, iter time: 56.81ms\n",
      "iter 3300 step 825: loss 4.5304, LR: 0.005853, iter time: 57.85ms\n",
      "iter 3400 step 850: loss 4.1766, LR: 0.005830, iter time: 58.52ms\n",
      "iter 3500 step 875: loss 4.2740, LR: 0.005805, iter time: 57.50ms\n",
      "iter 3600 step 900: loss 4.2820, LR: 0.005779, iter time: 58.88ms\n",
      "iter 3700 step 925: loss 4.2292, LR: 0.005751, iter time: 57.29ms\n",
      "iter 3800 step 950: loss 4.4272, LR: 0.005721, iter time: 51.97ms\n",
      "iter 3900 step 975: loss 4.2242, LR: 0.005690, iter time: 56.50ms\n",
      "Saving checkpoint to 'out/redpajama/iter-003999-ckpt.pth'\n",
      "iter 4000 step 1000: loss 4.2313, LR: 0.005657, iter time: 52.22ms\n",
      "iter 4100 step 1025: loss 3.7856, LR: 0.005622, iter time: 58.16ms\n",
      "iter 4200 step 1050: loss 3.6729, LR: 0.005586, iter time: 57.64ms\n",
      "iter 4300 step 1075: loss 3.3938, LR: 0.005549, iter time: 57.79ms\n",
      "iter 4400 step 1100: loss 4.0303, LR: 0.005510, iter time: 59.47ms\n",
      "iter 4500 step 1125: loss 4.0657, LR: 0.005469, iter time: 57.41ms\n",
      "iter 4600 step 1150: loss 4.1320, LR: 0.005428, iter time: 52.20ms\n",
      "iter 4700 step 1175: loss 3.7960, LR: 0.005384, iter time: 58.86ms\n",
      "iter 4800 step 1200: loss 4.2554, LR: 0.005340, iter time: 58.40ms\n",
      "iter 4900 step 1225: loss 3.8772, LR: 0.005294, iter time: 56.95ms\n",
      "iter 5000 step 1250: loss 3.6817, LR: 0.005246, iter time: 57.29ms\n",
      "iter 5100 step 1275: loss 3.9250, LR: 0.005198, iter time: 57.37ms\n",
      "iter 5200 step 1300: loss 4.3746, LR: 0.005148, iter time: 57.65ms\n",
      "iter 5300 step 1325: loss 4.0608, LR: 0.005096, iter time: 56.79ms\n",
      "iter 5400 step 1350: loss 4.0989, LR: 0.005044, iter time: 56.71ms\n",
      "iter 5500 step 1375: loss 3.7373, LR: 0.004990, iter time: 56.18ms\n",
      "iter 5600 step 1400: loss 3.7889, LR: 0.004936, iter time: 49.99ms\n",
      "iter 5700 step 1425: loss 4.0925, LR: 0.004880, iter time: 54.77ms\n",
      "iter 5800 step 1450: loss 3.8588, LR: 0.004823, iter time: 55.27ms\n",
      "iter 5900 step 1475: loss 4.1029, LR: 0.004765, iter time: 50.15ms\n",
      "iter 6000 step 1500: loss 3.7252, LR: 0.004705, iter time: 49.90ms\n",
      "iter 6100 step 1525: loss 3.4831, LR: 0.004645, iter time: 55.53ms\n",
      "iter 6200 step 1550: loss 3.9866, LR: 0.004584, iter time: 54.90ms\n",
      "iter 6300 step 1575: loss 3.7148, LR: 0.004522, iter time: 55.46ms\n",
      "iter 6400 step 1600: loss 3.5724, LR: 0.004459, iter time: 50.00ms\n",
      "iter 6500 step 1625: loss 3.7567, LR: 0.004396, iter time: 55.14ms\n",
      "iter 6600 step 1650: loss 3.6850, LR: 0.004331, iter time: 54.25ms\n",
      "iter 6700 step 1675: loss 3.8028, LR: 0.004266, iter time: 54.40ms\n",
      "iter 6800 step 1700: loss 3.8552, LR: 0.004200, iter time: 54.60ms\n",
      "iter 6900 step 1725: loss 3.8593, LR: 0.004133, iter time: 54.42ms\n",
      "iter 7000 step 1750: loss 3.9515, LR: 0.004066, iter time: 54.44ms\n",
      "iter 7100 step 1775: loss 3.6975, LR: 0.003998, iter time: 49.42ms\n",
      "iter 7200 step 1800: loss 3.8195, LR: 0.003929, iter time: 54.56ms\n",
      "iter 7300 step 1825: loss 3.7486, LR: 0.003860, iter time: 49.36ms\n",
      "iter 7400 step 1850: loss 3.8587, LR: 0.003790, iter time: 54.44ms\n",
      "iter 7500 step 1875: loss 3.6892, LR: 0.003720, iter time: 51.14ms\n",
      "iter 7600 step 1900: loss 3.1359, LR: 0.003650, iter time: 57.86ms\n",
      "iter 7700 step 1925: loss 3.4292, LR: 0.003579, iter time: 56.81ms\n",
      "iter 7800 step 1950: loss 3.4693, LR: 0.003508, iter time: 50.99ms\n",
      "iter 7900 step 1975: loss 3.2231, LR: 0.003436, iter time: 56.80ms\n",
      "Saving checkpoint to 'out/redpajama/iter-007999-ckpt.pth'\n",
      "iter 8000 step 2000: loss 3.8785, LR: 0.003364, iter time: 50.98ms\n",
      "iter 8100 step 2025: loss 3.3586, LR: 0.003292, iter time: 56.01ms\n",
      "iter 8200 step 2050: loss 3.5540, LR: 0.003220, iter time: 69.23ms\n",
      "iter 8300 step 2075: loss 3.5014, LR: 0.003148, iter time: 64.78ms\n",
      "iter 8400 step 2100: loss 4.1081, LR: 0.003075, iter time: 64.00ms\n",
      "iter 8500 step 2125: loss 3.4069, LR: 0.003003, iter time: 70.03ms\n",
      "iter 8600 step 2150: loss 3.4040, LR: 0.002931, iter time: 50.92ms\n",
      "iter 8700 step 2175: loss 3.6502, LR: 0.002858, iter time: 56.68ms\n",
      "iter 8800 step 2200: loss 3.7889, LR: 0.002786, iter time: 51.22ms\n",
      "iter 8900 step 2225: loss 3.5215, LR: 0.002714, iter time: 56.65ms\n",
      "iter 9000 step 2250: loss 3.5122, LR: 0.002642, iter time: 56.33ms\n",
      "iter 9100 step 2275: loss 3.2583, LR: 0.002570, iter time: 55.80ms\n",
      "iter 9200 step 2300: loss 3.6411, LR: 0.002498, iter time: 50.97ms\n",
      "iter 9300 step 2325: loss 3.3789, LR: 0.002427, iter time: 56.73ms\n",
      "iter 9400 step 2350: loss 3.4600, LR: 0.002356, iter time: 56.51ms\n",
      "iter 9500 step 2375: loss 3.5573, LR: 0.002286, iter time: 56.18ms\n",
      "iter 9600 step 2400: loss 3.8386, LR: 0.002216, iter time: 56.41ms\n",
      "iter 9700 step 2425: loss 3.6447, LR: 0.002146, iter time: 50.92ms\n",
      "iter 9800 step 2450: loss 3.7056, LR: 0.002077, iter time: 56.08ms\n",
      "iter 9900 step 2475: loss 3.6347, LR: 0.002008, iter time: 56.52ms\n",
      "iter 10000 step 2500: loss 3.6540, LR: 0.001940, iter time: 56.33ms\n",
      "iter 10100 step 2525: loss 2.9236, LR: 0.001873, iter time: 56.44ms\n",
      "iter 10200 step 2550: loss 3.3438, LR: 0.001806, iter time: 56.55ms\n",
      "iter 10300 step 2575: loss 2.8825, LR: 0.001740, iter time: 56.82ms\n",
      "iter 10400 step 2600: loss 3.2412, LR: 0.001675, iter time: 56.60ms\n",
      "iter 10500 step 2625: loss 3.7394, LR: 0.001610, iter time: 51.23ms\n",
      "iter 10600 step 2650: loss 3.0055, LR: 0.001547, iter time: 56.51ms\n",
      "iter 10700 step 2675: loss 3.0301, LR: 0.001484, iter time: 55.63ms\n",
      "iter 10800 step 2700: loss 3.7498, LR: 0.001422, iter time: 56.42ms\n",
      "iter 10900 step 2725: loss 3.3228, LR: 0.001361, iter time: 56.71ms\n",
      "iter 11000 step 2750: loss 3.7291, LR: 0.001301, iter time: 56.78ms\n",
      "iter 11100 step 2775: loss 3.2531, LR: 0.001241, iter time: 50.94ms\n",
      "iter 11200 step 2800: loss 3.5005, LR: 0.001183, iter time: 51.03ms\n",
      "iter 11300 step 2825: loss 3.6694, LR: 0.001126, iter time: 55.72ms\n",
      "iter 11400 step 2850: loss 3.6801, LR: 0.001070, iter time: 56.28ms\n",
      "iter 11500 step 2875: loss 3.5435, LR: 0.001016, iter time: 56.29ms\n",
      "iter 11600 step 2900: loss 3.1989, LR: 0.000962, iter time: 56.46ms\n",
      "iter 11700 step 2925: loss 3.7118, LR: 0.000910, iter time: 56.78ms\n",
      "iter 11800 step 2950: loss 3.4840, LR: 0.000858, iter time: 56.68ms\n",
      "iter 11900 step 2975: loss 3.3171, LR: 0.000808, iter time: 58.20ms\n",
      "Saving checkpoint to 'out/redpajama/iter-011999-ckpt.pth'\n",
      "iter 12000 step 3000: loss 3.4315, LR: 0.000760, iter time: 133.55ms\n",
      "iter 12100 step 3025: loss 3.2361, LR: 0.000712, iter time: 50.96ms\n",
      "iter 12200 step 3050: loss 3.2333, LR: 0.000666, iter time: 56.34ms\n"
     ]
    },
    {
     "ename": "KeyboardInterrupt",
     "evalue": "",
     "output_type": "error",
     "traceback": [
      "\u001b[0;31m---------------------------------------------------------------------------\u001b[0m",
      "\u001b[0;31mKeyboardInterrupt\u001b[0m                         Traceback (most recent call last)",
      "\u001b[1;32m/home/raghu/work/ERA-V1-assignments/assignment-22/main.ipynb Cell 14\u001b[0m line \u001b[0;36m2\n\u001b[1;32m      <a href='vscode-notebook-cell://wsl%2Bubuntu/home/raghu/work/ERA-V1-assignments/assignment-22/main.ipynb#X16sdnNjb2RlLXJlbW90ZQ%3D%3D?line=0'>1</a>\u001b[0m torch\u001b[39m.\u001b[39mset_float32_matmul_precision(\u001b[39m\"\u001b[39m\u001b[39mmedium\u001b[39m\u001b[39m\"\u001b[39m)\n\u001b[0;32m----> <a href='vscode-notebook-cell://wsl%2Bubuntu/home/raghu/work/ERA-V1-assignments/assignment-22/main.ipynb#X16sdnNjb2RlLXJlbW90ZQ%3D%3D?line=1'>2</a>\u001b[0m setup(\n\u001b[1;32m      <a href='vscode-notebook-cell://wsl%2Bubuntu/home/raghu/work/ERA-V1-assignments/assignment-22/main.ipynb#X16sdnNjb2RlLXJlbW90ZQ%3D%3D?line=2'>3</a>\u001b[0m     devices\u001b[39m=\u001b[39;49m\u001b[39m1\u001b[39;49m,\n\u001b[1;32m      <a href='vscode-notebook-cell://wsl%2Bubuntu/home/raghu/work/ERA-V1-assignments/assignment-22/main.ipynb#X16sdnNjb2RlLXJlbW90ZQ%3D%3D?line=3'>4</a>\u001b[0m     \u001b[39m#train_data_dir=Path(\"data/lit-redpajama-sample\")\u001b[39;49;00m\n\u001b[1;32m      <a href='vscode-notebook-cell://wsl%2Bubuntu/home/raghu/work/ERA-V1-assignments/assignment-22/main.ipynb#X16sdnNjb2RlLXJlbW90ZQ%3D%3D?line=4'>5</a>\u001b[0m     train_data_dir\u001b[39m=\u001b[39;49mPath(\u001b[39m\"\u001b[39;49m\u001b[39m/home/raghu/work/data/redpajama/data/lit-redpajama-sample\u001b[39;49m\u001b[39m\"\u001b[39;49m)\n\u001b[1;32m      <a href='vscode-notebook-cell://wsl%2Bubuntu/home/raghu/work/ERA-V1-assignments/assignment-22/main.ipynb#X16sdnNjb2RlLXJlbW90ZQ%3D%3D?line=5'>6</a>\u001b[0m )\n",
      "\u001b[1;32m/home/raghu/work/ERA-V1-assignments/assignment-22/main.ipynb Cell 14\u001b[0m line \u001b[0;36m2\n\u001b[1;32m     <a href='vscode-notebook-cell://wsl%2Bubuntu/home/raghu/work/ERA-V1-assignments/assignment-22/main.ipynb#X16sdnNjb2RlLXJlbW90ZQ%3D%3D?line=24'>25</a>\u001b[0m fabric \u001b[39m=\u001b[39m L\u001b[39m.\u001b[39mFabric(devices\u001b[39m=\u001b[39mdevices, strategy\u001b[39m=\u001b[39mstrategy, precision\u001b[39m=\u001b[39mprecision, loggers\u001b[39m=\u001b[39mlogger)\n\u001b[1;32m     <a href='vscode-notebook-cell://wsl%2Bubuntu/home/raghu/work/ERA-V1-assignments/assignment-22/main.ipynb#X16sdnNjb2RlLXJlbW90ZQ%3D%3D?line=25'>26</a>\u001b[0m fabric\u001b[39m.\u001b[39mprint(hparams)\n\u001b[0;32m---> <a href='vscode-notebook-cell://wsl%2Bubuntu/home/raghu/work/ERA-V1-assignments/assignment-22/main.ipynb#X16sdnNjb2RlLXJlbW90ZQ%3D%3D?line=26'>27</a>\u001b[0m fabric\u001b[39m.\u001b[39;49mlaunch(main, train_data_dir, val_data_dir, resume)\n",
      "File \u001b[0;32m~/.local/lib/python3.10/site-packages/lightning/fabric/fabric.py:834\u001b[0m, in \u001b[0;36mFabric.launch\u001b[0;34m(self, function, *args, **kwargs)\u001b[0m\n\u001b[1;32m    829\u001b[0m \u001b[39melif\u001b[39;00m \u001b[39misinstance\u001b[39m(\u001b[39mself\u001b[39m\u001b[39m.\u001b[39mstrategy\u001b[39m.\u001b[39mlauncher, (_MultiProcessingLauncher, _XLALauncher)):\n\u001b[1;32m    830\u001b[0m     \u001b[39mraise\u001b[39;00m \u001b[39mTypeError\u001b[39;00m(\n\u001b[1;32m    831\u001b[0m         \u001b[39mf\u001b[39m\u001b[39m\"\u001b[39m\u001b[39mTo use the `\u001b[39m\u001b[39m{\u001b[39;00m\u001b[39mtype\u001b[39m(\u001b[39mself\u001b[39m\u001b[39m.\u001b[39mstrategy)\u001b[39m.\u001b[39m\u001b[39m__name__\u001b[39m\u001b[39m}\u001b[39;00m\u001b[39m` strategy, `.launch()` needs to be called with a function\u001b[39m\u001b[39m\"\u001b[39m\n\u001b[1;32m    832\u001b[0m         \u001b[39m\"\u001b[39m\u001b[39m that contains the code to launch in processes.\u001b[39m\u001b[39m\"\u001b[39m\n\u001b[1;32m    833\u001b[0m     )\n\u001b[0;32m--> 834\u001b[0m \u001b[39mreturn\u001b[39;00m \u001b[39mself\u001b[39;49m\u001b[39m.\u001b[39;49m_wrap_and_launch(function, \u001b[39mself\u001b[39;49m, \u001b[39m*\u001b[39;49margs, \u001b[39m*\u001b[39;49m\u001b[39m*\u001b[39;49mkwargs)\n",
      "File \u001b[0;32m~/.local/lib/python3.10/site-packages/lightning/fabric/fabric.py:920\u001b[0m, in \u001b[0;36mFabric._wrap_and_launch\u001b[0;34m(self, to_run, *args, **kwargs)\u001b[0m\n\u001b[1;32m    918\u001b[0m \u001b[39mif\u001b[39;00m (launcher \u001b[39m:=\u001b[39m \u001b[39mself\u001b[39m\u001b[39m.\u001b[39m_strategy\u001b[39m.\u001b[39mlauncher) \u001b[39mis\u001b[39;00m \u001b[39mnot\u001b[39;00m \u001b[39mNone\u001b[39;00m:\n\u001b[1;32m    919\u001b[0m     \u001b[39mreturn\u001b[39;00m launcher\u001b[39m.\u001b[39mlaunch(to_run, \u001b[39m*\u001b[39margs, \u001b[39m*\u001b[39m\u001b[39m*\u001b[39mkwargs)\n\u001b[0;32m--> 920\u001b[0m \u001b[39mreturn\u001b[39;00m to_run(\u001b[39m*\u001b[39;49margs, \u001b[39m*\u001b[39;49m\u001b[39m*\u001b[39;49mkwargs)\n",
      "File \u001b[0;32m~/.local/lib/python3.10/site-packages/lightning/fabric/fabric.py:925\u001b[0m, in \u001b[0;36mFabric._wrap_with_setup\u001b[0;34m(self, to_run, *args, **kwargs)\u001b[0m\n\u001b[1;32m    923\u001b[0m \u001b[39mself\u001b[39m\u001b[39m.\u001b[39m_strategy\u001b[39m.\u001b[39msetup_environment()\n\u001b[1;32m    924\u001b[0m \u001b[39mwith\u001b[39;00m _replace_dunder_methods(DataLoader, \u001b[39m\"\u001b[39m\u001b[39mdataset\u001b[39m\u001b[39m\"\u001b[39m), _replace_dunder_methods(BatchSampler):\n\u001b[0;32m--> 925\u001b[0m     \u001b[39mreturn\u001b[39;00m to_run(\u001b[39m*\u001b[39;49margs, \u001b[39m*\u001b[39;49m\u001b[39m*\u001b[39;49mkwargs)\n",
      "\u001b[1;32m/home/raghu/work/ERA-V1-assignments/assignment-22/main.ipynb Cell 14\u001b[0m line \u001b[0;36m7\n\u001b[1;32m     <a href='vscode-notebook-cell://wsl%2Bubuntu/home/raghu/work/ERA-V1-assignments/assignment-22/main.ipynb#X16sdnNjb2RlLXJlbW90ZQ%3D%3D?line=67'>68</a>\u001b[0m     fabric\u001b[39m.\u001b[39mload(resume, state)\n\u001b[1;32m     <a href='vscode-notebook-cell://wsl%2Bubuntu/home/raghu/work/ERA-V1-assignments/assignment-22/main.ipynb#X16sdnNjb2RlLXJlbW90ZQ%3D%3D?line=69'>70</a>\u001b[0m train_time \u001b[39m=\u001b[39m time\u001b[39m.\u001b[39mperf_counter()\n\u001b[0;32m---> <a href='vscode-notebook-cell://wsl%2Bubuntu/home/raghu/work/ERA-V1-assignments/assignment-22/main.ipynb#X16sdnNjb2RlLXJlbW90ZQ%3D%3D?line=70'>71</a>\u001b[0m train(fabric, state, train_dataloader, val_dataloader, speed_monitor)\n\u001b[1;32m     <a href='vscode-notebook-cell://wsl%2Bubuntu/home/raghu/work/ERA-V1-assignments/assignment-22/main.ipynb#X16sdnNjb2RlLXJlbW90ZQ%3D%3D?line=71'>72</a>\u001b[0m fabric\u001b[39m.\u001b[39mprint(\u001b[39mf\u001b[39m\u001b[39m\"\u001b[39m\u001b[39mTraining time: \u001b[39m\u001b[39m{\u001b[39;00m(time\u001b[39m.\u001b[39mperf_counter()\u001b[39m-\u001b[39mtrain_time)\u001b[39m:\u001b[39;00m\u001b[39m.2f\u001b[39m\u001b[39m}\u001b[39;00m\u001b[39ms\u001b[39m\u001b[39m\"\u001b[39m)\n\u001b[1;32m     <a href='vscode-notebook-cell://wsl%2Bubuntu/home/raghu/work/ERA-V1-assignments/assignment-22/main.ipynb#X16sdnNjb2RlLXJlbW90ZQ%3D%3D?line=72'>73</a>\u001b[0m \u001b[39mif\u001b[39;00m fabric\u001b[39m.\u001b[39mdevice\u001b[39m.\u001b[39mtype \u001b[39m==\u001b[39m \u001b[39m\"\u001b[39m\u001b[39mcuda\u001b[39m\u001b[39m\"\u001b[39m:\n",
      "\u001b[1;32m/home/raghu/work/ERA-V1-assignments/assignment-22/main.ipynb Cell 14\u001b[0m line \u001b[0;36m5\n\u001b[1;32m     <a href='vscode-notebook-cell://wsl%2Bubuntu/home/raghu/work/ERA-V1-assignments/assignment-22/main.ipynb#X16sdnNjb2RlLXJlbW90ZQ%3D%3D?line=51'>52</a>\u001b[0m \u001b[39m# return \u001b[39;00m\n\u001b[1;32m     <a href='vscode-notebook-cell://wsl%2Bubuntu/home/raghu/work/ERA-V1-assignments/assignment-22/main.ipynb#X16sdnNjb2RlLXJlbW90ZQ%3D%3D?line=53'>54</a>\u001b[0m \u001b[39mif\u001b[39;00m \u001b[39mnot\u001b[39;00m is_accumulating:\n\u001b[0;32m---> <a href='vscode-notebook-cell://wsl%2Bubuntu/home/raghu/work/ERA-V1-assignments/assignment-22/main.ipynb#X16sdnNjb2RlLXJlbW90ZQ%3D%3D?line=54'>55</a>\u001b[0m     fabric\u001b[39m.\u001b[39;49mclip_gradients(model, optimizer, max_norm\u001b[39m=\u001b[39;49mgrad_clip)\n\u001b[1;32m     <a href='vscode-notebook-cell://wsl%2Bubuntu/home/raghu/work/ERA-V1-assignments/assignment-22/main.ipynb#X16sdnNjb2RlLXJlbW90ZQ%3D%3D?line=55'>56</a>\u001b[0m     optimizer\u001b[39m.\u001b[39mstep()\n\u001b[1;32m     <a href='vscode-notebook-cell://wsl%2Bubuntu/home/raghu/work/ERA-V1-assignments/assignment-22/main.ipynb#X16sdnNjb2RlLXJlbW90ZQ%3D%3D?line=56'>57</a>\u001b[0m     optimizer\u001b[39m.\u001b[39mzero_grad()\n",
      "File \u001b[0;32m~/.local/lib/python3.10/site-packages/lightning/fabric/fabric.py:455\u001b[0m, in \u001b[0;36mFabric.clip_gradients\u001b[0;34m(self, module, optimizer, clip_val, max_norm, norm_type, error_if_nonfinite)\u001b[0m\n\u001b[1;32m    453\u001b[0m     \u001b[39mreturn\u001b[39;00m \u001b[39mNone\u001b[39;00m\n\u001b[1;32m    454\u001b[0m \u001b[39mif\u001b[39;00m max_norm \u001b[39mis\u001b[39;00m \u001b[39mnot\u001b[39;00m \u001b[39mNone\u001b[39;00m:\n\u001b[0;32m--> 455\u001b[0m     \u001b[39mreturn\u001b[39;00m \u001b[39mself\u001b[39;49m\u001b[39m.\u001b[39;49mstrategy\u001b[39m.\u001b[39;49mclip_gradients_norm(\n\u001b[1;32m    456\u001b[0m         _unwrap_objects(module),\n\u001b[1;32m    457\u001b[0m         _unwrap_objects(optimizer),\n\u001b[1;32m    458\u001b[0m         max_norm\u001b[39m=\u001b[39;49mmax_norm,\n\u001b[1;32m    459\u001b[0m         norm_type\u001b[39m=\u001b[39;49mnorm_type,\n\u001b[1;32m    460\u001b[0m         error_if_nonfinite\u001b[39m=\u001b[39;49merror_if_nonfinite,\n\u001b[1;32m    461\u001b[0m     )\n\u001b[1;32m    462\u001b[0m \u001b[39mraise\u001b[39;00m \u001b[39mValueError\u001b[39;00m(\u001b[39m\"\u001b[39m\u001b[39mYou have to specify either `clip_val` or `max_norm` to do gradient clipping!\u001b[39m\u001b[39m\"\u001b[39m)\n",
      "File \u001b[0;32m~/.local/lib/python3.10/site-packages/lightning/fabric/strategies/strategy.py:380\u001b[0m, in \u001b[0;36mStrategy.clip_gradients_norm\u001b[0;34m(self, module, optimizer, max_norm, norm_type, error_if_nonfinite)\u001b[0m\n\u001b[1;32m    378\u001b[0m \u001b[39mself\u001b[39m\u001b[39m.\u001b[39mprecision\u001b[39m.\u001b[39munscale_gradients(optimizer)\n\u001b[1;32m    379\u001b[0m parameters \u001b[39m=\u001b[39m \u001b[39mself\u001b[39m\u001b[39m.\u001b[39mprecision\u001b[39m.\u001b[39mmain_params(optimizer)\n\u001b[0;32m--> 380\u001b[0m \u001b[39mreturn\u001b[39;00m torch\u001b[39m.\u001b[39;49mnn\u001b[39m.\u001b[39;49mutils\u001b[39m.\u001b[39;49mclip_grad_norm_(\n\u001b[1;32m    381\u001b[0m     parameters, max_norm\u001b[39m=\u001b[39;49mmax_norm, norm_type\u001b[39m=\u001b[39;49mnorm_type, error_if_nonfinite\u001b[39m=\u001b[39;49merror_if_nonfinite\n\u001b[1;32m    382\u001b[0m )\n",
      "File \u001b[0;32m~/.local/lib/python3.10/site-packages/torch/nn/utils/clip_grad.py:63\u001b[0m, in \u001b[0;36mclip_grad_norm_\u001b[0;34m(parameters, max_norm, norm_type, error_if_nonfinite, foreach)\u001b[0m\n\u001b[1;32m     59\u001b[0m             norms\u001b[39m.\u001b[39mextend([torch\u001b[39m.\u001b[39mlinalg\u001b[39m.\u001b[39mvector_norm(g, norm_type) \u001b[39mfor\u001b[39;00m g \u001b[39min\u001b[39;00m grads])\n\u001b[1;32m     61\u001b[0m     total_norm \u001b[39m=\u001b[39m torch\u001b[39m.\u001b[39mlinalg\u001b[39m.\u001b[39mvector_norm(torch\u001b[39m.\u001b[39mstack([norm\u001b[39m.\u001b[39mto(first_device) \u001b[39mfor\u001b[39;00m norm \u001b[39min\u001b[39;00m norms]), norm_type)\n\u001b[0;32m---> 63\u001b[0m \u001b[39mif\u001b[39;00m error_if_nonfinite \u001b[39mand\u001b[39;00m torch\u001b[39m.\u001b[39mlogical_or(total_norm\u001b[39m.\u001b[39misnan(), total_norm\u001b[39m.\u001b[39misinf()):\n\u001b[1;32m     64\u001b[0m     \u001b[39mraise\u001b[39;00m \u001b[39mRuntimeError\u001b[39;00m(\n\u001b[1;32m     65\u001b[0m         \u001b[39mf\u001b[39m\u001b[39m'\u001b[39m\u001b[39mThe total norm of order \u001b[39m\u001b[39m{\u001b[39;00mnorm_type\u001b[39m}\u001b[39;00m\u001b[39m for gradients from \u001b[39m\u001b[39m'\u001b[39m\n\u001b[1;32m     66\u001b[0m         \u001b[39m'\u001b[39m\u001b[39m`parameters` is non-finite, so it cannot be clipped. To disable \u001b[39m\u001b[39m'\u001b[39m\n\u001b[1;32m     67\u001b[0m         \u001b[39m'\u001b[39m\u001b[39mthis error and scale the gradients by the non-finite norm anyway, \u001b[39m\u001b[39m'\u001b[39m\n\u001b[1;32m     68\u001b[0m         \u001b[39m'\u001b[39m\u001b[39mset `error_if_nonfinite=False`\u001b[39m\u001b[39m'\u001b[39m)\n\u001b[1;32m     69\u001b[0m clip_coef \u001b[39m=\u001b[39m max_norm \u001b[39m/\u001b[39m (total_norm \u001b[39m+\u001b[39m \u001b[39m1e-6\u001b[39m)\n",
      "\u001b[0;31mKeyboardInterrupt\u001b[0m: "
     ]
    }
   ],
   "source": [
    "torch.set_float32_matmul_precision(\"medium\")\n",
    "setup(\n",
    "    devices=1,\n",
    "    #train_data_dir=Path(\"data/lit-redpajama-sample\")\n",
    "    train_data_dir=Path(\"/home/raghu/work/data/redpajama/data/lit-redpajama-sample\")\n",
    ")"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": []
  }
 ],
 "metadata": {
  "kernelspec": {
   "display_name": "base",
   "language": "python",
   "name": "python3"
  },
  "language_info": {
   "codemirror_mode": {
    "name": "ipython",
    "version": 3
   },
   "file_extension": ".py",
   "mimetype": "text/x-python",
   "name": "python",
   "nbconvert_exporter": "python",
   "pygments_lexer": "ipython3",
   "version": "3.10.12"
  }
 },
 "nbformat": 4,
 "nbformat_minor": 2
}