### Start Training and Timing (Python)

Source: https://github.com/kellerjordan/modded-nanogpt/blob/master/records/092925_PolarExpress/df2b2147-08b3-48d4-8ac7-75d509768560.txt

Initializes the training data loader and prepares for the main training loop. It synchronizes the CUDA device and records the start time (`t0`) using `time.perf_counter` to measure the duration of the training process. This is the setup immediately preceding the actual training iterations.

```python
import time
import torch

# Assuming distributed_data_generator, args, grad_accum_steps are defined
# training_time_ms = 0
# start the clock
# torch.cuda.synchronize()
# t0 = time.perf_counter()
```

--------------------------------

### Training and Validation Loop Setup (Python)

Source: https://github.com/kellerjordan/modded-nanogpt/blob/master/records/082325_SparseAttnGate/4518e917-cec2-4c81-9c1a-53b0644c2326.txt

Sets up the data loader for the training set and prepares for the main training loop. It includes synchronization for accurate timing and records the start time.

```Python
train_loader = distributed_data_generator(args.train_files, args.train_seq_len, grad_accum_steps, align_to_bos=True)
training_time_ms = 0
# start the clock
torch.cuda.synchronize()
t0 = time.perf_counter()

```

--------------------------------

### Training and Validation Setup

Source: https://github.com/kellerjordan/modded-nanogpt/blob/master/records/091125_VectSigmoidBFloat16/0d451b7e-6500-41ae-ac9d-02352e611b88.txt

Sets up the data loader for training and starts a timer for measuring training performance. This section prepares the data pipeline and initiates the timing for the main training loop.

```python
train_loader = distributed_data_generator(
    args.train_files,
    args.train_batch_size,
    args.train_max_seq_len,
    grad_accum_steps=grad_accum_steps,
)
training_time_ms = 0
# start the clock
torch.cuda.synchronize()
t0 = time.perf_counter()
```

--------------------------------

### Start Training Loop Timer

Source: https://github.com/kellerjordan/modded-nanogpt/blob/master/records/120824_UNetValueEmbedsTweaks/23ed0b57-544b-4fc5-b780-062dcf3f6f2c.txt

This code prepares for the training loop by initializing a timer. It synchronizes the CUDA device and records the start time using `time.perf_counter` to measure the training duration.

```python
# Start training loop
training_time_ms = 0
# start the clock
torch.cuda.synchronize()
t0 = time.perf_counter()
```

--------------------------------

### Logging and Environment Setup

Source: https://github.com/kellerjordan/modded-nanogpt/blob/master/records/110924_Replicateleloykun/1621af10-aa0c-42af-bf54-8a773c63a2af.txt

Initializes logging by creating a unique run directory and log file. It logs the Python source code, hardware/software environment information (including nvidia-smi output), and starts a timer for training duration.

```python
import uuid
import os
import subprocess
import time

if master_process:
    run_id = str(uuid.uuid4())
    logdir = 'logs/%s/' % run_id
    os.makedirs(logdir, exist_ok=True)
    logfile = 'logs/%s.txt' % run_id
    with open(logfile, "w") as f:
        f.write('='*100 + '\n')
        f.write(code)
        f.write('='*100 + '\n')
        f.write(f"Running pytorch {torch.__version__} compiled for CUDA {torch.version.cuda}\n")
        result = subprocess.run(['nvidia-smi'], stdout=subprocess.PIPE, stderr=subprocess.PIPE, text=True)
        f.write(f'nvidia-smi:\n{result.stdout}\n')
        f.write('='*100 + '\n')

training_time_ms = 0
torch.cuda.synchronize()
t0 = time.time()
```

--------------------------------

### Training and Validation Loop Setup (Python)

Source: https://github.com/kellerjordan/modded-nanogpt/blob/master/records/082325_SparseAttnGate/020630eb-2191-4ba2-9ee4-4cdc94316943.txt

Sets up the data loader for training and initializes timers for performance measurement. It synchronizes CUDA devices and records the start time of the training process.

```Python
########################################
#        Training and validation       #
########################################

train_loader = distributed_data_generator(args.train_files, args.train_seq_len, grad_accum_steps, align_to_bos=True)
training_time_ms = 0
# start the clock
torch.cuda.synchronize()
t0 = time.perf_counter()
```

--------------------------------

### Training and Validation Setup with Performance Timing

Source: https://github.com/kellerjordan/modded-nanogpt/blob/master/records/082325_SparseAttnGate/ca042caf-b232-4a25-b28f-88e39a2009d3.txt

This section prepares the training data loader and starts the performance timer for the main training loop. It synchronizes CUDA events and records the start time. Requires `torch`, `time`, `distributed_data_generator`, and `args`.

```python
import torch
import time

########################################
#        Training and validation       #
########################################

train_loader = distributed_data_generator(args.train_files, args.train_seq_len, grad_accum_steps, align_to_bos=True)
training_time_ms = 0
# start the clock
torch.cuda.synchronize()
t0 = time.perf_counter()
```

--------------------------------

### Start Training and Timing (Python)

Source: https://github.com/kellerjordan/modded-nanogpt/blob/master/records/091525_AsyncDataLoadAttnFinalWindow/720dfef5-25c6-4e17-81b5-975ecfdd4d81.txt

Initializes the training data loader and starts the main training loop timer. It synchronizes CUDA operations before recording the start time to ensure accurate measurement of training duration. This is the entry point for the actual model training process.

```python
import torch
import time

# Assuming distributed_data_generator, args, grad_accum_steps are defined
# train_loader = distributed_data_generator(args.train_files, args.train_batch_size, args.train_max_seq_len, grad_accum_steps=grad_accum_steps)
training_time_ms = 0
# start the clock
torch.cuda.synchronize()
t0 = time.perf_counter()

# Placeholder for the actual start of training and timing
# This requires 'distributed_data_generator', 'args', and 'grad_accum_steps' to be defined.
```

--------------------------------

### Training and Validation Setup in Python

Source: https://github.com/kellerjordan/modded-nanogpt/blob/master/records/082325_SparseAttnGate/6df384bb-9c24-46b3-826b-f7c07168c27a.txt

Sets up the training data loader and initializes performance timing. It synchronizes CUDA events and records the start time (`t0`) before commencing the main training loop.

```python
train_loader = distributed_data_generator(args.train_files, args.train_seq_len, grad_accum_steps, align_to_bos=True)
training_time_ms = 0
torch.cuda.synchronize()
t0 = time.perf_counter()
```

--------------------------------

### Logging and Environment Information in Python

Source: https://github.com/kellerjordan/modded-nanogpt/blob/master/records/111024_ScaleShortcuts/d6b50d71-f419-4d26-bb39-a60d55ae7a04.txt

Initializes a logging directory and file for recording training progress and system information. It logs the Python script content, PyTorch version, CUDA version, and the output of `nvidia-smi` for environment details. It also starts a timer for training duration.

```python
import uuid
import os
import subprocess
import time

if master_process:
    run_id = str(uuid.uuid4())
    logdir = 'logs/%s/' % run_id
    os.makedirs(logdir, exist_ok=True)
    logfile = 'logs/%s.txt' % run_id
    with open(logfile, "w") as f:
        f.write('='*100 + '\n')
        # Assuming 'code' variable holds the script content
        # f.write(code)
        f.write('='*100 + '\n')
        f.write(f"Running pytorch {torch.version.__version__} compiled for CUDA {torch.version.cuda}\n")
        f.write('nvidia-smi:\n')
        result = subprocess.run(['nvidia-smi'], stdout=subprocess.PIPE, stderr=subprocess.PIPE, text=True)
        f.write(f'{result.stdout}\n')
        f.write('='*100 + '\n')

training_time_ms = 0
torch.cuda.synchronize()
t0 = time.time()
```

--------------------------------

### Initialize Training Timer (Python)

Source: https://github.com/kellerjordan/modded-nanogpt/blob/master/records/121024_MFUTweaks/0aa83756-53f0-4268-9721-db6d5985bc42.txt

Initializes variables for training time measurement, including a sliding window block counter and synchronization of the CUDA device. It then records the start time using `time.perf_counter()`. Dependencies include 'torch' and 'time'.

```Python
import torch
import time

# Initialize sliding window variables
sliding_window_num_blocks = torch.tensor(1, dtype=torch.int32, device="cuda")
sw_num_blocks_prev = 1

# Start training loop
training_time_ms = 0

# start the clock
torch.cuda.synchronize()
t0 = time.perf_counter()
```

--------------------------------

### Logging and Environment Information Setup

Source: https://github.com/kellerjordan/modded-nanogpt/blob/master/records/100924_SOAP/5bdc3988-496c-4232-b4ef-53764cb81c92.txt

Sets up logging for the training process, including creating a run directory and a log file. It logs the Python script itself, along with detailed environment information such as PyTorch version, CUDA version, and the output of `nvidia-smi`.

```python
import uuid
import os
import subprocess
import time

if master_process:
    run_id = str(uuid.uuid4())
    logdir = 'logs/%s/' % run_id
    os.makedirs(logdir, exist_ok=True)
    logfile = 'logs/%s.txt' % run_id
    # create the log file
    with open(logfile, "w") as f:
        # begin the log by printing this file (the Python code)
        f.write('='*100 + '\n')
        f.write(code) # Assuming 'code' variable holds the script content
        f.write('='*100 + '\n')
        # log information about the hardware/software environment this is running on
        # and print the full `nvidia-smi` to file
        f.write(f"Running pytorch {torch.version.__version__} compiled for CUDA {torch.version.cuda}\n")
        f.write('nvidia-smi:\n')
        result = subprocess.run(['nvidia-smi'], stdout=subprocess.PIPE, stderr=subprocess.PIPE, text=True)
        f.write(f'{result.stdout}\n')
        f.write('='*100 + '\n')

training_time_ms = 0
# start the clock
torch.cuda.synchronize()
t0 = time.time()
```

--------------------------------

### Logging Setup and Environment Information in Python

Source: https://github.com/kellerjordan/modded-nanogpt/blob/master/records/120424_ValueEmbed/1b4db08c-46d9-4e1b-a8e6-872a685061c3.txt

Initializes logging to a file, captures the current script's code, and logs hardware/software environment details including NVIDIA SMI output. This setup is crucial for reproducibility and debugging.

```python
logfile = None
if master_process:
    run_id = str(uuid.uuid4())
    logdir = 'logs/%s/' % run_id
    os.makedirs(logdir, exist_ok=True)
    logfile = 'logs/%s.txt' % run_id
    # create the log file
    with open(logfile, "w") as f:
        # begin the log by printing this file (the Python code)
        f.write(code)
        f.write('='*100 + '\n')
def print0(s, logonly=False):
    if master_process:
        with open(logfile, "a") as f:
            if not logonly:
                print(s)
            f.write(s+'\n')
# log information about the hardware/software environment this is running on
# and print the full `nvidia-smi` to file
print0(f"Running pytorch {torch.version.__version__} compiled for CUDA {torch.version.cuda}\nnvidia-smi:")
import subprocess
result = subprocess.run(['nvidia-smi'], stdout=subprocess.PIPE, stderr=subprocess.PIPE, text=True)
print0(f'{result.stdout}', logonly=True)
print0('='*100, logonly=True)
```

--------------------------------

### Logging Setup and Environment Information in Python

Source: https://github.com/kellerjordan/modded-nanogpt/blob/master/records/112424_WindowWarmup/4428858e-7cb8-4a25-a936-818d8f28de51.txt

Initializes logging by creating a log directory and file. It then captures and logs PyTorch and CUDA versions, along with the output of `nvidia-smi` for environment details. This setup ensures that training runs are reproducible and all relevant environment information is preserved.

```python
# begin logging
logfile = None
if master_process:
    run_id = str(uuid.uuid4())
    logdir = 'logs/%s/' % run_id
    os.makedirs(logdir, exist_ok=True)
    logfile = 'logs/%s.txt' % run_id
    # create the log file
    with open(logfile, "w") as f:
        # begin the log by printing this file (the Python code)
        f.write('='*100 + '\n')
        f.write(code)
        f.write('='*100 + '\n')
def print0(s, logonly=False):
    if master_process:
        with open(logfile, "a") as f:
            if not logonly:
                print(s)
            f.write(s+'\n')
# log information about the hardware/software environment this is running on
# and print the full `nvidia-smi` to file
print0(f"Running pytorch {torch.version.__version__} compiled for CUDA {torch.version.cuda}\nnvidia-smi:")
import subprocess
result = subprocess.run(['nvidia-smi'], stdout=subprocess.PIPE, stderr=subprocess.PIPE, text=True)
print0(f'{result.stdout}', logonly=True)
print0('='*100, logonly=True)
```

--------------------------------

### Starting Training Timer (Python)

Source: https://github.com/kellerjordan/modded-nanogpt/blob/master/records/071225_BosAlign/190aef37-dea2-40f7-8ece-1c5dcd1d4e89.txt

Initializes a new data loader for training and prepares to measure the training time. It synchronizes CUDA operations and records the start time using `time.perf_counter`.

```python
train_loader = distributed_data_generator(args.train_files, world_size * args.seq_len, rank, world_size, args.train_align_to_bos)
training_time_ms = 0
# start the clock
torch.cuda.synchronize()
t0 = time.perf_counter()
```

--------------------------------

### Start Training and Validation Loop (Python)

Source: https://github.com/kellerjordan/modded-nanogpt/blob/master/records/092725_BF16CE/b21c8cc7-c09c-401f-b654-23e947ad3e38.txt

Initializes the training data loader and starts the main training timer. It synchronizes CUDA operations and records the start time before entering the training loop. This sets up the environment for measuring training performance.

```Python
import torch
import time

# Assuming distributed_data_generator, args, grad_accum_steps, model, optimizers are defined

# Placeholder for distributed_data_generator
def distributed_data_generator(files, batch_size, max_seq_len, grad_accum_steps):
    print(f"Initializing data generator for files: {files}, batch_size: {batch_size}, max_seq_len: {max_seq_len}")
    # In a real scenario, this would yield batches of data
    def generator():
        for i in range(10): # Simulate a few batches
            yield (torch.randn(batch_size, max_seq_len), torch.randn(batch_size, max_seq_len), torch.randint(1, max_seq_len, (batch_size,)))
    return generator()

# Mock args and other variables
class MockArgs:
    train_batch_size = 4
    val_batch_size = 4
    num_iterations = 1000
    cooldown_frac = 0.1
    iteration_extension = 10
    ws_validate = 1024
    ws_schedule = [1024, 2048, 4096]
    train_files = ['file1.txt', 'file2.txt']
    train_max_seq_len = 1024

args = MockArgs()
grad_accum_steps = 1

########################################
#        Training and validation       #
########################################

train_loader = distributed_data_generator(args.train_files, args.train_batch_size, args.train_max_seq_len, grad_accum_steps=grad_accum_steps)
training_time_ms = 0

# start the clock
torch.cuda.synchronize()
t0 = time.perf_counter()
print("Training timer started.")

# The actual training loop would follow here
# Example of how the loop might start:
# for step in range(args.num_iterations):
#     inputs, targets, cum_seqlens = next(train_loader)
#     # ... training logic ...
#     if (step + 1) % 100 == 0: # Example validation frequency
#         # ... validation logic ...

```

--------------------------------

### Start Training and Timing in PyTorch

Source: https://github.com/kellerjordan/modded-nanogpt/blob/master/records/092925_PolarExpress/3bb6c2eb-1935-46d5-9f07-40b98223cfaa.txt

Sets up the training data loader and starts a timer using `time.perf_counter` after synchronizing CUDA. This marks the beginning of the actual training loop and measures the total training time.

```python
import time

# Assuming distributed_data_generator, args, grad_accum_steps are defined
# train_loader = distributed_data_generator(args.train_files, args.train_batch_size, args.train_max_seq_len,
#                                           grad_accum_steps=grad_accum_steps)
# training_time_ms = 0
# # start the clock
# torch.cuda.synchronize()
# t0 = time.perf_counter()
```

--------------------------------

### Python: Nanogpt Hyperparameters and DDP Setup

Source: https://github.com/kellerjordan/modded-nanogpt/blob/master/records/120824_UNetValueEmbedsTweaks/f2fda86f-4418-4ff8-a3a3-8112731cbcf9.txt

Defines hyperparameters for Nanogpt training and sets up the distributed data parallel (DDP) environment using PyTorch. This includes initializing the process group, determining device assignments, and setting up logging for the master process.

```python
import torch
import torch.distributed as dist
import os
import uuid
from dataclasses import dataclass

@dataclass
class Hyperparameters:
    # data hyperparams
    input_bin : str = 'data/fineweb10B/fineweb_train_*.bin' # input .bin to train on
    input_val_bin : str = 'data/fineweb10B/fineweb_val_*.bin' # input .bin to eval validation loss on
    # optimization hyperparams
    batch_size : int = 8 # batch size, in sequences, across all devices
    sequence_length : int = 64*1024 # sequence length, in tokens
    num_iterations : int = 1480 # number of iterations to run
    warmup_iters : int = 0
    cooldown_iters : int = 600 # number of iterations of linear warmup/cooldown for triangular or trapezoidal schedule
    weight_decay : float = 0
    # evaluation and logging hyperparams
    val_loss_every : int = 125 # every how many steps to evaluate val loss? 0 for only at the end
    val_tokens : int = 10485760 # how many tokens of validation data? it's important to keep this fixed for consistent comparisons
    save_every : int = 0 # every how many steps to save the checkpoint? 0 for only at the end
args = Hyperparameters()

# set up DDP (distributed data parallel). torchrun sets this env variable
assert torch.cuda.is_available()
dist.init_process_group(backend='nccl')
ddp_rank = int(os.environ['RANK'])
ddp_local_rank = int(os.environ['LOCAL_RANK'])
ddp_world_size = int(os.environ['WORLD_SIZE'])
device = f'cuda:{ddp_local_rank}'
torch.cuda.set_device(device)
print(f"using device: {device}")
master_process = (ddp_rank == 0) # this process will do logging, checkpointing etc.

# begin logging
logfile = None
if master_process:
    run_id = str(uuid.uuid4())
    logdir = 'logs/%s/' % run_id
    # os.makedirs(logdir, exist_ok=True)
    logfile = 'logs/%s.txt' % run_id
    # create the log file
    with open(logfile, "w") as f:
        # begin the log by printing this file (the Python code)
        f.write(code) # 'code' variable is not defined in this snippet, assuming it refers to the current script content
        f.write('='*100 + '\n')
def print0(s, logonly=False):
    if master_process:
        with open(logfile, "a") as f:
            if not logonly:
                print(s)
            f.write(s+'\n')
# log information about the hardware/software environment this is running on

```

--------------------------------

### Python Training Setup and Distributed Initialization

Source: https://github.com/kellerjordan/modded-nanogpt/blob/master/records/120824_UNetValueEmbedsTweaks/5fcbae90-4380-412e-8209-51fccc183040.txt

Demonstrates the setup for training, including defining hyperparameters using `dataclass`, initializing distributed process groups with PyTorch, and setting up the CUDA device based on environment variables.

```python
@dataclass
class Hyperparameters:
    # data hyperparams
    input_bin : str = 'data/fineweb10B/fineweb_train_*.bin' # input .bin to train on
    input_val_bin : str = 'data/fineweb10B/fineweb_val_*.bin' # input .bin to eval validation loss on
    # optimization hyperparams
    batch_size : int = 8 # batch size, in sequences, across all devices
    sequence_length : int = 64*1024 # sequence length, in tokens
    num_iterations : int = 1480 # number of iterations to run
    warmup_iters : int = 0
    cooldown_iters : int = 600 # number of iterations of linear warmup/cooldown for triangular or trapezoidal schedule
    weight_decay : float = 0
    # evaluation and logging hyperparams
    val_loss_every : int = 125 # every how many steps to evaluate val loss? 0 for only at the end
    val_tokens : int = 10485760 # how many tokens of validation data? it's important to keep this fixed for consistent comparisons
    save_every : int = 0 # every how many steps to save the checkpoint? 0 for only at the end
args = Hyperparameters()

# set up DDP (distributed data parallel). torchrun sets this env variable
assert torch.cuda.is_available()
dist.init_process_group(backend='nccl')
ddp_rank = int(os.environ['RANK'])
ddp_local_rank = int(os.environ['LOCAL_RANK'])
ddp_world_size = int(os.environ['WORLD_SIZE'])
device = f'cuda:{ddp_local_rank}'
torch.cuda.set_device(device)
print(f"using device: {device}")
master_process = (ddp_rank == 0) # this process will do logging, checkpointing etc.
```

--------------------------------

### Start Training Loop and Timing - Python

Source: https://github.com/kellerjordan/modded-nanogpt/blob/master/records/092725_BF16CE/5c44ff06-998a-4310-af6e-f0f5441452f4.txt

Initializes the training dataloader and starts the main training timer using `time.perf_counter`. It synchronizes CUDA events before starting the timer to ensure accurate measurement of the training duration. This marks the beginning of the model's training phase.

```python
import time
# Assuming distributed_data_generator, args, grad_accum_steps, optimizers, model are defined

train_loader = distributed_data_generator(args.train_files, args.train_batch_size, args.train_max_seq_len, grad_accum_steps=grad_accum_steps)
training_time_ms = 0
# start the clock
torch.cuda.synchronize()
t0 = time.perf_counter()
```

--------------------------------

### Distributed Setup and Logging - Python

Source: https://github.com/kellerjordan/modded-nanogpt/blob/master/records/121024_MFUTweaks/76ee9a40-0f61-42bb-849b-bf46c3a3e7c9.txt

Configures distributed training using PyTorch's DDP (Distributed Data Parallel) and sets up logging. It initializes the process group, determines the device, and establishes a master process for logging and checkpointing. A helper function `print0` is provided for synchronized printing and logging.

```python
@dataclass
class Hyperparameters:
    # data hyperparams
    input_bin : str = 'data/fineweb10B/fineweb_train_*.bin' # input .bin to train on
    input_val_bin : str = 'data/fineweb10B/fineweb_val_*.bin' # input .bin to eval validation loss on
    # optimization hyperparams
    batch_size : int = 8 # batch size, in sequences, across all devices
    sequence_length : int = 64*1024 # sequence length, in tokens
    num_iterations : int = 1480 # number of iterations to run
    warmup_iters : int = 0
    cooldown_iters : int = 600 # number of iterations of linear warmup/cooldown for triangular or trapezoidal schedule
    weight_decay : float = 0
    # evaluation and logging hyperparams
    val_loss_every : int = 125 # every how many steps to evaluate val loss? 0 for only at the end
    val_tokens : int = 10485760 # how many tokens of validation data? it's important to keep this fixed for consistent comparisons
    save_every : int = 0 # every how many steps to save the checkpoint? 0 for only at the end
args = Hyperparameters()

# set up DDP (distributed data parallel). torchrun sets this env variable
ddp_rank = int(os.environ['RANK'])
ddp_local_rank = int(os.environ['LOCAL_RANK'])
ddp_world_size = int(os.environ['WORLD_SIZE'])
assert torch.cuda.is_available()
device = torch.device(f"cuda:{ddp_local_rank}")
torch.cuda.set_device(device)
print(f"using device: {device}")
dist.init_process_group(backend='nccl', device_id=device)
dist.barrier()
master_process = (ddp_rank == 0) # this process will do logging, checkpointing etc.

# begin logging
logfile = None
if master_process:
    run_id = uuid.uuid4()
    logdir = Path("logs") / f"{run_id}"
    logdir.mkdir(exist_ok=True)
    logfile = Path("logs") / f"{run_id}.txt"
    print(logfile.stem)
    # create the log file
    with logfile.open("w") as f:
        # begin the log by printing this file (the Python code)
        print(code, file=f)
        print("=" * 100, file=f)
def print0(s, logonly=False):
    if master_process:
        with logfile.open("a") as f:
            if not logonly:
                print(s)
            print(s, file=f)
# log information about the hardware/software environment this is running on

```

--------------------------------

### Start Training and Validation Phase - Python

Source: https://github.com/kellerjordan/modded-nanogpt/blob/master/records/092125_DropAttn/278c1540-3b42-4ab0-94ed-273830bdfa11.txt

Prepares the training data loader and initializes timers for the training and validation phase. It synchronizes CUDA events and records the start time to measure the overall duration of the training process.

```Python
import time
import torch

# Assuming distributed_data_generator and args are defined
# Dummy classes and functions for demonstration purposes
def distributed_data_generator(files, batch_size, max_seq_len, grad_accum_steps):
    def generator():
        for _ in range(10):
            yield torch.randint(0, 50257, (10, max_seq_len)), torch.randint(0, 50257, (10, max_seq_len)), torch.arange(11)
    return generator()

class Args:
    def __init__(self):
        self.train_batch_size = 4
        self.val_batch_size = 4
        self.num_iterations = 10000
        self.cooldown_frac = 0.1
        self.iteration_extension = 0
        self.ws_validate = 1
        self.ws_validate_final_layer = 1
        self.ws_schedule = [1, 2, 4, 8, 16]
        self.train_files = []
        self.val_files = []
        self.train_max_seq_len = 1024
args = Args()

grad_accum_steps = 1

train_loader = distributed_data_generator(args.train_files, args.train_batch_size, args.train_max_seq_len, grad_accum_steps=grad_accum_steps)
training_time_ms = 0

# start the clock
torch.cuda.synchronize()
t0 = time.perf_counter()
```

--------------------------------

### Start Training Timer (Python)

Source: https://github.com/kellerjordan/modded-nanogpt/blob/master/records/092725_BF16CE/11dd9171-d060-4279-a6e5-5ba91fb7758e.txt

Initializes the training data loader and synchronizes CUDA devices before starting a performance timer. This setup is preparatory for the main training loop, ensuring that the timer capture only the actual training execution time, excluding setup and warm-up phases.

```python
import time
import torch

# Assuming distributed_data_generator, args, and grad_accum_steps are defined
train_loader = distributed_data_generator(args.train_files, args.train_batch_size, args.train_max_seq_len, grad_accum_steps=grad_accum_steps)
training_time_ms = 0
# start the clock
torch.cuda.synchronize()
t0 = time.perf_counter()
```

--------------------------------

### Setup Distributed Training Environment with PyTorch

Source: https://github.com/kellerjordan/modded-nanogpt/blob/master/records/121024_MFUTweaks/0aa83756-53f0-4268-9721-db6d5985bc42.txt

Initializes the distributed training environment using PyTorch's distributed backend (NCCL). It sets up CUDA device, initializes process groups, and identifies the master process for logging and checkpointing. This is a standard setup for multi-GPU and multi-node training.

```python
import torch
import os
import uuid
import distutils.util
from pathlib import Path
from dataclasses import dataclass
import torch.distributed as dist

@dataclass
class Hyperparameters:
    # data hyperparams
    input_bin : str = 'data/fineweb10B/fineweb_train_*.bin' # input .bin to train on
    input_val_bin : str = 'data/fineweb10B/fineweb_val_*.bin' # input .bin to eval validation loss on
    # optimization hyperparams
    batch_size : int = 8 # batch size, in sequences, across all devices
    sequence_length : int = 64*1024 # sequence length, in tokens
    num_iterations : int = 1480 # number of iterations to run
    warmup_iters : int = 0
    cooldown_iters : int = 600 # number of iterations of linear warmup/cooldown for triangular or trapezoidal schedule
    weight_decay : float = 0
    # evaluation and logging hyperparams
    val_loss_every : int = 125 # every how many steps to evaluate val loss? 0 for only at the end
    val_tokens : int = 10485760 # how many tokens of validation data? it's important to keep this fixed for consistent comparisons
    save_every : int = 0 # every how many steps to save the checkpoint? 0 for only at the end
args = Hyperparameters()

# set up DDP (distributed data parallel). torchrun sets this env variable
ddp_rank = int(os.environ['RANK'])
ddp_local_rank = int(os.environ['LOCAL_RANK'])
ddp_world_size = int(os.environ['WORLD_SIZE'])
assert torch.cuda.is_available()
device = torch.device(f"cuda:{ddp_local_rank}")
torch.cuda.set_device(device)
print(f"using device: {device}")
dist.init_process_group(backend='nccl', device_id=device)
dist.barrier()
master_process = (ddp_rank == 0) # this process will do logging, checkpointing etc.

# begin logging
logfile = None
if master_process:
    run_id = uuid.uuid4()
    logdir = Path("logs") / f"{run_id}"
    logdir.mkdir(exist_ok=True)
    logfile = Path("logs") / f"{run_id}.txt"
    print(logfile.stem)
    # create the log file
    with logfile.open("w") as f:
        # begin the log by printing this file (the Python code)
        # NOTE: 'code' variable is assumed to be defined elsewhere, likely containing the script's source
        # print(code, file=f) 
        print("=" * 100, file=f)
def print0(s, logonly=False):
    if master_process:
        with logfile.open("a") as f:
            if not logonly:
                print(s)
            print(s, file=f)
# log information about the hardware/software environment this is running on
```

--------------------------------

### Python Logging Setup and Patching

Source: https://github.com/kellerjordan/modded-nanogpt/blob/master/records/041625_GPT2Medium_Record7/223_3310d0b1-b24d-48ee-899f-d5c2a254a195.txt

Sets up file-based logging for distributed training. It defines a `print0` function to write messages to a log file and optionally to the console, and patches `torch._inductor.codecache.trace_structured` and `torch._inductor.graph.trace_structured` to log specific inductor events.

```python
if master_process:
    run_id_full = f"{run_id:03d}_{uuid.uuid4()}"
    os.makedirs("logs", exist_ok=True)
    logfile = f"logs/{run_id_full}.txt"
    print(logfile)
def print0(s, console=False):
    if master_process:
        with open(logfile, "a") as f:
            if console:
                print(s)
            print(s, file=f)
from torch._logging._internal import trace_structured # noqa: E402
import torch._inductor.codecache # noqa: E402
import torch._inductor.graph # noqa: E402
def _patched_trace_structured(name, metadata_fn, **kwargs):
    if name == "inductor_output_code":
        print0(f"inductor_output_code: {metadata_fn().get("filename", "Unknown")}")
    trace_structured(name, metadata_fn, **kwargs)
torch._inductor.codecache.trace_structured = _patched_trace_structured
torch._inductor.graph.trace_structured = _patched_trace_structured

# begin by printing this file (the Python code)
print0(code)
print0("="*100)
# log information about the hardware/software environment this is running on
print0(f"Running Python {sys.version}")
print0(f"Running PyTorch {torch.version.__version__} compiled for CUDA {torch.version.cuda}")
def nvidia_smi():
    import subprocess  # avoid top level import
    return subprocess.run(["nvidia-smi"], stdout=subprocess.PIPE, stderr=subprocess.PIPE, text=True).stdout
print0(nvidia_smi())
print0("="*100)
```

--------------------------------

### Main Function Entry Point for Training Script

Source: https://github.com/kellerjordan/modded-nanogpt/blob/master/records/082325_SparseAttnGate/020630eb-2191-4ba2-9ee4-4cdc94316943.txt

Sets up the data path for training and validation files by reading the DATA_PATH environment variable or defaulting to the current directory. This prepares the arguments for subsequent data loading and model initialization. Dependencies include `os`.

```python
# int main


data_path = os.environ.get("DATA_PATH", ".")
args.train_files = os.path.join(data_path, args.train_files)
args.val_files = os.path.join(data_path, args.val_files)
```

--------------------------------

### Start Training and Validation Clock

Source: https://github.com/kellerjordan/modded-nanogpt/blob/master/records/092125_DropAttn/79383017-eb05-4857-842a-7866b00571b4.txt

Prepares the training data loader and initializes a timer before starting the main training loop. It synchronizes CUDA operations and records the start time using `time.perf_counter()`. This setup is essential for accurately measuring the duration of the training and validation phases.

```python
import torch
import time
from nanogpt.distributed import distributed_data_generator

# Assuming args, grad_accum_steps, model, optimizers are defined
########################################
#        Training and validation       #
########################################

train_loader = distributed_data_generator(args.train_files, args.train_batch_size, args.train_max_seq_len, grad_accum_steps=grad_accum_steps)
training_time_ms = 0
# start the clock
torch.cuda.synchronize()
t0 = time.perf_counter()
```

--------------------------------

### Distributed Training Initialization and Logging Setup (Python)

Source: https://github.com/kellerjordan/modded-nanogpt/blob/master/records/121724_SparsifyEmbeds/165384d5-5a41-4fba-9b79-8ecc372616ea.txt

This section sets up distributed data parallel training and logging. It initializes the process group, sets the CUDA device, determines the master process, and configures a file logger to capture output, including the source code itself.

```python
import torch
import os
import dist
import uuid
from pathlib import Path
from dataclasses import dataclass

@dataclass
class Hyperparameters:
    # data hyperparams
    input_bin : str = 'data/fineweb10B/fineweb_train_*.bin' # input .bin to train on
    input_val_bin : str = 'data/fineweb10B/fineweb_val_*.bin' # input .bin to eval validation loss on
    # optimization hyperparams
    batch_size : int = 8 # batch size, in sequences, across all devices
    sequence_length : int = 64*1024 # sequence length, in tokens
    num_iterations : int = 1490 # number of iterations to run
    warmup_iters : int = 0
    cooldown_iters : int = 600 # number of iterations of linear warmup/cooldown for triangular or trapezoidal schedule
    weight_decay : float = 0
    # evaluation and logging hyperparams
    val_loss_every : int = 125 # every how many steps to evaluate val loss? 0 for only at the end
    val_tokens : int = 10485760 # how many tokens of validation data? it's important to keep this fixed for consistent comparisons
    save_every : int = 0 # every how many steps to save the checkpoint? 0 for only at the end
args = Hyperparameters()

# set up DDP (distributed data parallel). torchrun sets this env variable
rank = int(os.environ['RANK'])
local_rank = int(os.environ['LOCAL_RANK'])
world_size = int(os.environ['WORLD_SIZE'])
assert torch.cuda.is_available()
device = torch.device(f"cuda:{local_rank}")
torch.cuda.set_device(device)
print(f"using device: {device}")
dist.init_process_group(backend='nccl', device_id=device)
dist.barrier()
master_process = (rank == 0) # this process will do logging, checkpointing etc.

# begin logging
logfile = None
if master_process:
    run_id = uuid.uuid4()
    logdir = Path("logs") / f"{run_id}"
    # logdir.mkdir(parents=True, exist_ok=True)
    logfile = Path("logs") / f"{run_id}.txt"
    print(logfile.stem)
    # create the log file
    with logfile.open("w") as f:
        # begin the log by printing this file (the Python code)
        # Assume 'code' variable holds the content of this script
        # print(code, file=f)
        print("=" * 100, file=f)
def print0(s, logonly=False):
    if master_process:
        with logfile.open("a") as f:
            # if not logonly:
            #     print(s)
            print(s, file=f)
# log information about the hardware/software environment this is running on

```

--------------------------------

### Start Training Timer (Python)

Source: https://github.com/kellerjordan/modded-nanogpt/blob/master/records/092925_PolarExpress/f62629d6-8c01-4154-96d1-85945920514a.txt

Initializes the data loader for training and starts a timer using `time.perf_counter()` after synchronizing CUDA events. This is the beginning of the main training loop. Requires `time.perf_counter` and `torch.cuda.synchronize`.

```Python
import time
import torch

# Assuming distributed_data_generator, args, grad_accum_steps are defined
# train_loader = distributed_data_generator(args.train_files, args.train_batch_size, args.train_max_seq_len,
#                                           grad_accum_steps=grad_accum_steps)
# training_time_ms = 0
# start the clock
# torch.cuda.synchronize()
# t0 = time.perf_counter()

```

--------------------------------

### Main Execution Setup and Argument Parsing (Python)

Source: https://github.com/kellerjordan/modded-nanogpt/blob/master/records/082325_SparseAttnGate/c6be54c1-12d0-45a3-83cb-41cad0868d15.txt

Configures the data path for training and validation files based on environment variables and command-line arguments. It sets up the `args` object with the appropriate file paths, preparing the environment for the main training loop.

```Python
# int main


data_path = os.environ.get("DATA_PATH", ".")
args.train_files = os.path.join(data_path, args.train_files)
args.val_files = os.path.join(data_path, args.val_files)

```

--------------------------------

### Shell: Installing PyTorch Nightly and FA3 Wheel

Source: https://github.com/kellerjordan/modded-nanogpt/blob/master/records/090325_FA3/README.md

This shell script demonstrates how to install a specific nightly build of PyTorch (version 2.9.0.dev20250718) compatible with CUDA 12.6, followed by the installation of a prebuilt Flash Attention v3 wheel. This is recommended for exact reproduction of the project's setup.

```shell
pip install --pre "torch==2.9.0.dev20250718+cu126" --index-url https://download.pytorch.org/whl/nightly/cu126

pip install /path/to/flash_attn_3-3.0.0b1-cp39-abi3-linux_x86_64.whl
```

--------------------------------

### Start Training and Validation Clock (Python)

Source: https://github.com/kellerjordan/modded-nanogpt/blob/master/records/091825_Smear/692ba682-7466-4b59-a31a-7e0adcb55b4b.txt

Initializes the training data loader and synchronizes the CUDA device to accurately measure the start time of the training process using `time.perf_counter`. This is the setup phase before the main training loop begins.

```python
import time

# Assuming distributed_data_generator, args, grad_accum_steps, and optimizers are defined

# Mock distributed_data_generator for demonstration
def distributed_data_generator(files, batch_size, max_seq_len, grad_accum_steps):
    while True:
        yield torch.randint(0, 1000, (batch_size, max_seq_len)), torch.randint(0, 1000, (batch_size, max_seq_len)), torch.randint(0, max_seq_len, (batch_size,))

train_loader = distributed_data_generator(args.train_files, args.train_batch_size, args.train_max_seq_len, grad_accum_steps=grad_accum_steps)
training_time_ms = 0

# start the clock
torch.cuda.synchronize()
t0 = time.perf_counter()
```

--------------------------------

### Distributed Training Setup and Hyperparameters

Source: https://github.com/kellerjordan/modded-nanogpt/blob/master/records/050925_SkipMLPBlocks/1858912a-2697-4461-9edb-e5ee4246ee3d.txt

Initializes the distributed training environment using PyTorch DDP, sets up hyperparameters using a dataclass, and configures environment variables for data paths and distributed ranks. It also handles logging setup and environment information printing.

```python
import os
import sys
import uuid

import torch
import torch.distributed as dist
import torch.nn as nn
import triton

from dataclasses import dataclass
from typing import Optional, Tuple

# Assuming GPT, DistAdam, Muon, and code are defined elsewhere or imported
# from .model import GPT
# from .optim import DistAdam, Muon

# Dummy definitions for demonstration purposes
class GPT(nn.Module):
    def __init__(self, vocab_size, num_layers, num_heads, model_dim, max_seq_len):
        super().__init__()
        # Simplified model structure for example
        self.embed = nn.Embedding(vocab_size, model_dim)
        self.blocks = nn.ModuleList([nn.Module() for _ in range(num_layers)]) # Placeholder
        self.lm_head = nn.Linear(model_dim, vocab_size)

    def modules(self):
        # Mock modules method
        return [self.embed, self.lm_head]

    def parameters(self):
        # Mock parameters method
        return self.lm_head.parameters()

    def named_parameters(self):
        # Mock named_parameters method
        return [('lm_head.weight', self.lm_head.weight)]

    def bfloat16(self):
        pass # Mock method

class DistAdam(torch.optim.Optimizer):
    def __init__(self, params, lr, betas, eps, weight_decay):
        super().__init__(params, {'lr': lr, 'betas': betas, 'eps': eps, 'weight_decay': weight_decay})

class Muon(torch.optim.Optimizer):
    def __init__(self, params, lr, momentum, weight_decay):
        super().__init__(params, {'lr': lr, 'momentum': momentum, 'weight_decay': weight_decay})

code = "# This is a placeholder for the Python code being logged."

@dataclass
class Hyperparameters:
    # data
    train_files: str = "data/fineweb10B/fineweb_train_*.bin"
    val_files: str = "data/fineweb10B/fineweb_val_*.bin"
    val_tokens: int = 10485760
    train_batch_size: int = 2048 * 24 * 8
    train_max_seq_len: int = 128 * 16
    val_batch_size: int = 4 * 64 * 1024 * 8
    # optimization
    num_iterations: int = 1705
    cooldown_frac: int = 0.45
    # evaluation and logging
    run_id: str = str(uuid.uuid4())
    val_loss_every: int = 125
    save_checkpoint: bool = False
    # attention masking
    block_size: int = 128
    ws_schedule: tuple = (3, 7, 11)

args = Hyperparameters()

data_path = os.environ.get("DATA_PATH", ".")
args.train_files = os.path.join(data_path, args.train_files)
args.val_files = os.path.join(data_path, args.val_files)

# torchrun sets these env variables
rank = int(os.environ["RANK"])
world_size = int(os.environ["WORLD_SIZE"])
assert 8 % world_size == 0, "world_size must be a divisor of 8"
grad_accum_steps = 8 // world_size
assert torch.cuda.is_available()
device = torch.device("cuda", int(os.environ["LOCAL_RANK"]))
torch.cuda.set_device(device)
dist.init_process_group(backend="nccl", device_id=device)
dist.barrier()
master_process = (rank == 0) # this process will do logging, checkpointing etc.

# begin logging
logfile = None
if master_process:
    run_id = args.run_id
    os.makedirs("logs", exist_ok=True)
    logfile = f"logs/{run_id}.txt"
    print(logfile)
def print0(s, console=False):
    if master_process:
        with open(logfile, "a") as f:
            if console:
                print(s)
            print(s, file=f)

# begin by printing this file (the Python code)
print0(code)
print0("="*100)
# log information about the hardware/software environment this is running on
print0(f"Running Python {sys.version}")
print0(f"Running PyTorch {torch.version.__version__} compiled for CUDA {torch.version.cuda}")
print0(f"Running Triton version {triton.__version__}")

def nvidia_smi():
    import subprocess  # avoid top level import
    return subprocess.run(["nvidia-smi"], stdout=subprocess.PIPE, stderr=subprocess.PIPE, text=True).stdout
print0(nvidia_smi())
print0("="*100)

```

--------------------------------

### Start Training Timer in Python

Source: https://github.com/kellerjordan/modded-nanogpt/blob/master/records/071825_TritonMuon/record.txt

Initializes the training data loader and prepares for the main training loop by synchronizing CUDA and recording the start time. This setup is crucial for accurate performance measurement of the training process. It uses PyTorch's CUDA synchronization and Python's `time.perf_counter`.

```python
train_loader = distributed_data_generator(args.train_files, args.train_seq_len, grad_accum_steps, align_to_bos=True)
training_time_ms = 0
# start the clock
torch.cuda.synchronize()
t0 = time.perf_counter()
```

--------------------------------

### Distributed Training Setup and Initialization (Python)

Source: https://github.com/kellerjordan/modded-nanogpt/blob/master/records/082325_SparseAttnGate/020630eb-2191-4ba2-9ee4-4cdc94316943.txt

Initializes distributed training environment using torchrun, sets up device and process group, and identifies the master process for logging and checkpointing. It also configures logging to a file and prints environment details.

```Python
rank = int(os.environ["RANK"])
world_size = int(os.environ["WORLD_SIZE"])
assert 8 % world_size == 0, "world_size must be a divisor of 8"
grad_accum_steps = 8 // world_size
assert torch.cuda.is_available()
device = torch.device("cuda", int(os.environ["LOCAL_RANK"]))
torch.cuda.set_device(device)
dist.init_process_group(backend="nccl", device_id=device)
dist.barrier()
master_process = (rank == 0) # this process will do logging, checkpointing etc.

# begin logging
logfile = None
if master_process:
    run_id = args.run_id
    os.makedirs("logs", exist_ok=True)
    logfile = f"logs/{run_id}.txt"
    print(logfile)
def print0(s, console=False):
    if master_process:
        with open(logfile, "a") as f:
            if console:
                print(s)
            print(s, file=f)

# begin by printing this file (the Python code)
print0(code)
print0("="*100)
# log information about the hardware/software environment this is running on
print0(f"Running Python {sys.version}")
print0(f"Running PyTorch {torch.version.__version__} compiled for CUDA {torch.version.cuda}")
print0(f"Running Triton version {triton.__version__}")
def nvidia_smi():
    import subprocess  # avoid top level import
    return subprocess.run(["nvidia-smi"], stdout=subprocess.PIPE, stderr=subprocess.PIPE, text=True).stdout
print0(nvidia_smi())
print0("="*100)
```

--------------------------------

### Training Loop Initialization (Python)

Source: https://github.com/kellerjordan/modded-nanogpt/blob/master/records/082325_SparseAttnGate/c6be54c1-12d0-45a3-83cb-41cad0868d15.txt

Initializes the training data loader and starts a timer for measuring training duration. This is the setup for the main training and validation loop.

```python
# Assuming 'args', 'GPT', 'DistAdam', 'Muon', 'distributed_data_generator', 'next_multiple_of_n' are defined elsewhere

########################################
#        Training and validation       #
########################################

train_loader = distributed_data_generator(args.train_files, args.train_seq_len, grad_accum_steps, align_to_bos=True)
training_time_ms = 0
# start the clock
torch.cuda.synchronize()
t0 = time.perf_counter()

```