TEMPLAR Submission

Submission Not Found

This submission could not be found.

Evaluation Results

Median MFU:

42.15%

#	MFU	TPS	Wall Time	Success	Date
1	42.31%	8,920	4.59s	Yes	Jan 7, 10:32
2	41.98%	8,845	4.63s	Yes	Jan 7, 10:33
3	42.15%	8,880	4.61s	Yes	Jan 7, 10:34

Submission Code

Python

Hash: a1b2c3d4e5f67890...

"""
Optimized training code for Templar Tournament.
"""

from collections.abc import Iterator
from dataclasses import dataclass

import torch
import torch.nn.functional as F


@dataclass
class InnerStepsResult:
    """Required return type from inner_steps function."""

    final_logits: torch.Tensor
    total_tokens: int
    final_loss: float


def inner_steps(
    model: torch.nn.Module,
    data_iterator: Iterator[torch.Tensor],
    optimizer: torch.optim.Optimizer,
    num_steps: int,
    device: torch.device,
) -> InnerStepsResult:
    """Run training for num_steps and return results."""
    torch.backends.cudnn.deterministic = True
    torch.backends.cudnn.benchmark = False

    total_tokens = 0
    final_logits = None
    final_loss = 0.0

    for step in range(num_steps):
        batch = next(data_iterator)
        batch = batch.to(device, dtype=torch.long)

        input_ids = batch[:, :-1]
        labels = batch[:, 1:]

        with torch.autocast(device_type=device.type, dtype=torch.bfloat16):
            outputs = model(input_ids)
            logits = outputs.logits if hasattr(outputs, 'logits') else outputs

            loss = F.cross_entropy(
                logits.reshape(-1, logits.size(-1)),
                labels.reshape(-1),
                ignore_index=-100,
            )

        loss.backward()
        optimizer.step()
        optimizer.zero_grad()

        total_tokens += batch.numel()
        final_logits = logits.detach().float()
        final_loss = loss.item()

    if torch.cuda.is_available():
        torch.cuda.synchronize()

    return InnerStepsResult(
        final_logits=final_logits,
        total_tokens=total_tokens,
        final_loss=final_loss,
    )