This file contains bidirectional Unicode text that may be interpreted or compiled differently than what appears below. To review, open the file in an editor that reveals hidden Unicode characters.
Learn more about bidirectional Unicode characters
# 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, |
This file contains bidirectional Unicode text that may be interpreted or compiled differently than what appears below. To review, open the file in an editor that reveals hidden Unicode characters.
Learn more about bidirectional Unicode characters
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) |
This file contains bidirectional Unicode text that may be interpreted or compiled differently than what appears below. To review, open the file in an editor that reveals hidden Unicode characters.
Learn more about bidirectional Unicode characters
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 |