antferdom

# see https://github.com/pytorch/torchsnapshot/blob/main/benchmarks/fsdp/main.py
import torch
from torch import distributed as dist, nn

def create_model() -> nn.Module:
    # 7.8GB model, 1.9B parameters
    model = nn.Transformer(
        d_model=864,
        num_encoder_layers=1,
        num_decoder_layers=20,

antferdom

import torch
    
model_size = sum(
    p.numel() * p.element_size() for p in model.parameters() if p.requires_grad
)
model_params = sum(p.numel() for p in model.parameters() if p.requires_grad)

antferdom

from typing import List, Dict, Tuple, Union, Optional

def largest_prime_factor(n: int):
    """Return the largest prime factor of n. Assume n > 1 and is not a prime.
    Q: What/How is a prime factor of n?
    A: A prime factor of n is a prime number that divides n.
    Q: What/how a number is not prime?
    A: A number is not prime if it has more than two factors.
    Examples:
    Q: 15