Formatting code format for flake8.

dlrover/examples/torch_nanogpt_debug_job.yaml

@@ -22,7 +22,8 @@ spec:
  - -c
  - "dlrover-run --nnodes=$WORKER_NUM \
  --nproc_per_node=1 --max_restarts=1 \
- model_zoo/pytorch/nanoGPT/nanoGPT.py --data_dir '/data/nanogpt/'"
+ model_zoo/pytorch/nanoGPT/nanoGPT.py \
+ --data_dir '/data/nanogpt/'"
  resources:
  limits:
  cpu: "8"

dlrover/examples/torch_nanogpt_job.yaml

@@ -22,7 +22,8 @@ spec:
  - -c
  - "dlrover-run --nnodes=$WORKER_NUM \
  --nproc_per_node=1 --max_restarts=1 \
- model_zoo/pytorch/nanoGPT/nanoGPT.py --data_dir '/data/nanogpt/'"
+ model_zoo/pytorch/nanoGPT/nanoGPT.py \
+ --data_dir '/data/nanogpt/'"
  resources:
  limits:
  cpu: "8"

model_zoo/pytorch/nanoGPT/model.py

@@ -13,10 +13,9 @@
 
 """
 Full definition of a GPT Language Model, all of it in this single file.
-References:
-1) the official GPT-2 TensorFlow implementation released by OpenAI:
-https://github.com/openai/gpt-2/blob/master/src/model.py
-2) huggingface/transformers PyTorch implementation:
+References: 1) the official GPT-2 TensorFlow implementation released by
+OpenAI: https://github.com/openai/gpt-2/blob/master/src/model.py 2)
+huggingface/transformers PyTorch implementation:
 https://github.com/huggingface/transformers/blob/main/src/transformers/models/gpt2/modeling_gpt2.py
 """
 
@@ -30,7 +29,10 @@
 
 
 class LayerNorm(nn.Module):
- """LayerNorm but with an optional bias. PyTorch doesn't support simply bias=False"""
+ """
+ LayerNorm but with an optional bias. PyTorch doesn't support simply
+ bias=False
+ """
 
  def __init__(self, ndim, bias):
  super().__init__()
@@ -59,15 +61,18 @@ def __init__(self, config):
  self.n_head = config.n_head
  self.n_embd = config.n_embd
  self.dropout = config.dropout
- # flash attention make GPU go brrrrr but support is only in PyTorch >= 2.0
+ # flash attention make GPU go brrrrr but support is only in
+ # PyTorch >= 2.0
  self.flash = hasattr(
  torch.nn.functional, "scaled_dot_product_attention"
  )
  if not self.flash:
  print(
- "WARNING: using slow attention. Flash Attention requires PyTorch >= 2.0"
+ "WARNING: using slow attention. "
+ "Flash Attention requires PyTorch >= 2.0"
  )
- # causal mask to ensure that attention is only applied to the left in the input sequence
+ # causal mask to ensure that attention is only applied to
+ # the left in the input sequence
  self.register_buffer(
  "bias",
  torch.tril(
@@ -84,7 +89,8 @@ def forward(self, x):
  x.size()
  ) # batch size, sequence length, embedding dimensionality (n_embd)
 
- # calculate query, key, values for all heads in batch and move head forward to be the batch dim
+ # calculate query, key, values for all heads in batch and move
+ # head forward to be the batch dim
  q, k, v = self.c_attn(x).split(self.n_embd, dim=2)
  k = k.view(B, T, self.n_head, C // self.n_head).transpose(
  1, 2
@@ -96,7 +102,8 @@ def forward(self, x):
  1, 2
  ) # (B, nh, T, hs)
 
- # causal self-attention; Self-attend: (B, nh, T, hs) x (B, nh, hs, T) -> (B, nh, T, T)
+ # causal self-attention; Self-attend: (B, nh, T, hs) x (B, nh,
+ # hs, T) -> (B, nh, T, T)
  if self.flash:
  # efficient attention using Flash Attention CUDA kernels
  y = torch.nn.functional.scaled_dot_product_attention(
@@ -160,12 +167,16 @@ def forward(self, x):
 @dataclass
 class GPTConfig:
  block_size: int = 1024
- vocab_size: int = 50304 # GPT-2 vocab_size of 50257, padded up to nearest multiple of 64 for efficiency
+ # GPT-2 vocab_size of 50257, padded up to nearest multiple of 64 for
+ # efficiency
+ vocab_size: int = 50304
  n_layer: int = 12
  n_head: int = 12
  n_embd: int = 768
  dropout: float = 0.0
- bias: bool = True # True: bias in Linears and LayerNorms, like GPT-2. False: a bit better and faster
+ # True: bias in Linears and LayerNorms, like GPT-2. False: a bit
+ # better and faster
+ bias: bool = True
 
 
 class GPT(nn.Module):
@@ -187,17 +198,19 @@ def __init__(self, config):
  )
  )
  self.lm_head = nn.Linear(config.n_embd, config.vocab_size, bias=False)
- # with weight tying when using torch.compile() some warnings get generated:
- # "UserWarning: functional_call was passed multiple values for tied weights.
- # This behavior is deprecated and will be an error in future versions"
- # not 100% sure what this is, so far seems to be harmless. TODO investigate
+ # with weight tying when using torch.compile() some warnings get
+ # generated: "UserWarning: functional_call was passed multiple
+ # values for tied weights. This behavior is deprecated and will
+ # be an error in future versions" not 100% sure what this is, so
+ # far seems to be harmless. TODO investigate
  self.transformer.wte.weight = (
  self.lm_head.weight
  ) # https://paperswithcode.com/method/weight-tying
 
  # init all weights
  self.apply(self._init_weights)
- # apply special scaled init to the residual projections, per GPT-2 paper
+ # apply special scaled init to the residual projections, per
+ # GPT-2 paper
  for pn, p in self.named_parameters():
  if pn.endswith("c_proj.weight"):
  torch.nn.init.normal_(
@@ -209,10 +222,11 @@ def __init__(self, config):
 
  def get_num_params(self, non_embedding=True):
  """
- Return the number of parameters in the model.
- For non-embedding count (default), the position embeddings get subtracted.
- The token embeddings would too, except due to the parameter sharing these
- params are actually used as weights in the final layer, so we include them.
+ Return the number of parameters in the model. For non-embedding
+ count (default), the position embeddings get subtracted. The
+ token embeddings would too, except due to the parameter sharing
+ these params are actually used as weights in the final layer, so
+ we include them.
  """
  n_params = sum(p.numel() for p in self.parameters())
  if non_embedding:
@@ -230,33 +244,34 @@ def _init_weights(self, module):
  def forward(self, idx, targets=None):
  device = idx.device
  b, t = idx.size()
- assert (
- t <= self.config.block_size
- ), f"Cannot forward sequence of length {t}, block size is only {self.config.block_size}"
+ assert t <= self.config.block_size, (
+ f"Cannot forward sequence of length {t}, "
+ f"block size is only {self.config.block_size}"
+ )
  pos = torch.arange(0, t, dtype=torch.long, device=device) # shape (t)
 
- # forward the GPT model itself
- tok_emb = self.transformer.wte(
- idx
- ) # token embeddings of shape (b, t, n_embd)
- pos_emb = self.transformer.wpe(
- pos
- ) # position embeddings of shape (t, n_embd)
+ # forward the GPT model itself token embeddings of shape (b, t,
+ # n_embd)
+ tok_emb = self.transformer.wte(idx)
+ # position embeddings of shape (t, n_embd)
+ pos_emb = self.transformer.wpe(pos)
  x = self.transformer.drop(tok_emb + pos_emb)
  for block in self.transformer.h:
  x = block(x)
  x = self.transformer.ln_f(x)
 
  if targets is not None:
- # if we are given some desired targets also calculate the loss
+ # if we are given some desired targets also calculate the
+ # loss
  logits = self.lm_head(x)
  loss = F.cross_entropy(
  logits.view(-1, logits.size(-1)),
  targets.view(-1),
  ignore_index=-1,
  )
  else:
- # inference-time mini-optimization: only forward the lm_head on the very last position
+ # inference-time mini-optimization: only forward the lm_head
+ # on the very last position
  logits = self.lm_head(
  x[:, [-1], :]
  ) # note: using list [-1] to preserve the time dim
@@ -265,9 +280,10 @@ def forward(self, idx, targets=None):
  return logits, loss
 
  def crop_block_size(self, block_size):
- # model surgery to decrease the block size if necessary
- # e.g. we may load the GPT2 pretrained model checkpoint (block size 1024)
- # but want to use a smaller block size for some smaller, simpler model
+ # model surgery to decrease the block size if necessary e.g. we
+ # may load the GPT2 pretrained model checkpoint (block size
+ # 1024) but want to use a smaller block size for some smaller,
+ # simpler model
  assert block_size <= self.config.block_size
  self.config.block_size = block_size
  self.transformer.wpe.weight = nn.Parameter(
@@ -292,23 +308,18 @@ def from_pretrained(cls, model_type, override_args=None):
  # n_layer, n_head and n_embd are determined from model_type
  config_args = {
  "gpt2": dict(n_layer=12, n_head=12, n_embd=768), # 124M params
- "gpt2-medium": dict(
- n_layer=24, n_head=16, n_embd=1024
- ), # 350M params
- "gpt2-large": dict(
- n_layer=36, n_head=20, n_embd=1280
- ), # 774M params
- "gpt2-xl": dict(
- n_layer=48, n_head=25, n_embd=1600
- ), # 1558M params
+ # 350M params
+ "gpt2-medium": dict(n_layer=24, n_head=16, n_embd=1024),
+ # 774M params
+ "gpt2-large": dict(n_layer=36, n_head=20, n_embd=1280),
+ # 1558M params
+ "gpt2-xl": dict(n_layer=48, n_head=25, n_embd=1600),
  }[model_type]
  print("forcing vocab_size=50257, block_size=1024, bias=True")
- config_args[
- "vocab_size"
- ] = 50257 # always 50257 for GPT model checkpoints
- config_args[
- "block_size"
- ] = 1024 # always 1024 for GPT model checkpoints
+ # always 50257 for GPT model checkpoints
+ config_args["vocab_size"] = 50257
+ # always 1024 for GPT model checkpoints
+ config_args["block_size"] = 1024
  config_args["bias"] = True # always True for GPT model checkpoints
  # we can override the dropout rate, if desired
  if "dropout" in override_args:
@@ -327,7 +338,8 @@ def from_pretrained(cls, model_type, override_args=None):
  model_hf = GPT2LMHeadModel.from_pretrained(model_type)
  sd_hf = model_hf.state_dict()
 
- # copy while ensuring all of the parameters are aligned and match in names and shapes
+ # copy while ensuring all of the parameters are aligned and
+ # match in names and shapes
  sd_keys_hf = sd_hf.keys()
  sd_keys_hf = [
  k for k in sd_keys_hf if not k.endswith(".attn.masked_bias")
@@ -341,14 +353,16 @@ def from_pretrained(cls, model_type, override_args=None):
  "mlp.c_fc.weight",
  "mlp.c_proj.weight",
  ]
- # basically the openai checkpoints use a "Conv1D" module, but we only want to use a vanilla Linear
- # this means that we have to transpose these weights when we import them
+ # basically the openai checkpoints use a "Conv1D" module, but we
+ # only want to use a vanilla Linear this means that we have to
+ # transpose these weights when we import them
  assert len(sd_keys_hf) == len(
  sd_keys
  ), f"mismatched keys: {len(sd_keys_hf)} != {len(sd_keys)}"
  for k in sd_keys_hf:
  if any(k.endswith(w) for w in transposed):
- # special treatment for the Conv1D weights we need to transpose
+ # special treatment for the Conv1D weights we need to
+ # transpose
  assert sd_hf[k].shape[::-1] == sd[k].shape
  with torch.no_grad():
  sd[k].copy_(sd_hf[k].t())
@@ -367,8 +381,9 @@ def configure_optimizers(
  param_dict = {pn: p for pn, p in self.named_parameters()}
  # filter out those that do not require grad
  param_dict = {pn: p for pn, p in param_dict.items() if p.requires_grad}
- # create optim groups. Any parameters that is 2D will be weight decayed, otherwise no.
- # i.e. all weight tensors in matmuls + embeddings decay, all biases and layernorms don't.
+ # create optim groups. Any parameters that is 2D will be weight
+ # decayed, otherwise no. i.e. all weight tensors in matmuls +
+ # embeddings decay, all biases and layernorms don't.
  decay_params = [p for n, p in param_dict.items() if p.dim() >= 2]
  nodecay_params = [p for n, p in param_dict.items() if p.dim() < 2]
  optim_groups = [
@@ -378,12 +393,15 @@ def configure_optimizers(
  num_decay_params = sum(p.numel() for p in decay_params)
  num_nodecay_params = sum(p.numel() for p in nodecay_params)
  print(
- f"num decayed parameter tensors: {len(decay_params)}, with {num_decay_params:,} parameters"
+ f"num decayed parameter tensors: {len(decay_params)},"
+ f" with {num_decay_params:,} parameters"
  )
  print(
- f"num non-decayed parameter tensors: {len(nodecay_params)}, with {num_nodecay_params:,} parameters"
+ f"num non-decayed parameter tensors: {len(nodecay_params)},"
+ f" with {num_nodecay_params:,} parameters"
  )
- # Create AdamW optimizer and use the fused version if it is available
+ # Create AdamW optimizer and use the fused version if it is
+ # available
  fused_available = (
  "fused" in inspect.signature(torch.optim.AdamW).parameters
  )
@@ -397,9 +415,12 @@ def configure_optimizers(
  return optimizer
 
  def estimate_mfu(self, fwdbwd_per_iter, dt):
- """estimate model flops utilization (MFU) in units of A100 bfloat16 peak FLOPS"""
- # first estimate the number of flops we do per iteration.
- # see PaLM paper Appendix B as ref: https://arxiv.org/abs/2204.02311
+ """
+ estimate model flops utilization (MFU) in units of A100 bfloat16
+ peak FLOPS
+ """
+ # first estimate the number of flops we do per iteration. see
+ # PaLM paper Appendix B as ref: https://arxiv.org/abs/2204.02311
  N = self.get_num_params()
  cfg = self.config
  L, H, Q, T = (
@@ -411,7 +432,8 @@ def estimate_mfu(self, fwdbwd_per_iter, dt):
  flops_per_token = 6 * N + 12 * L * H * Q * T
  flops_per_fwdbwd = flops_per_token * T
  flops_per_iter = flops_per_fwdbwd * fwdbwd_per_iter
- # express our flops throughput as ratio of A100 bfloat16 peak flops
+ # express our flops throughput as ratio of A100 bfloat16 peak
+ # flops
  flops_achieved = flops_per_iter * (1.0 / dt) # per second
  flops_promised = 312e12 # A100 GPU bfloat16 peak flops is 312 TFLOPS
  mfu = flops_achieved / flops_promised
@@ -420,26 +442,32 @@ def estimate_mfu(self, fwdbwd_per_iter, dt):
  @torch.no_grad()
  def generate(self, idx, max_new_tokens, temperature=1.0, top_k=None):
  """
- Take a conditioning sequence of indices idx (LongTensor of shape (b,t)) and complete
- the sequence max_new_tokens times, feeding the predictions back into the model each time.
- Most likely you'll want to make sure to be in model.eval() mode of operation for this.
+ Take a conditioning sequence of indices idx (LongTensor of shape
+ (b,t)) and complete the sequence max_new_tokens times, feeding
+ the predictions back into the model each time. Most likely
+ you'll want to make sure to be in model.eval() mode of operation
+ for this.
  """
  for _ in range(max_new_tokens):
- # if the sequence context is growing too long we must crop it at block_size
+ # if the sequence context is growing too long we must crop
+ # it at block_size
  idx_cond = (
  idx
  if idx.size(1) <= self.config.block_size
  else idx[:, -self.config.block_size :]
  )
- # forward the model to get the logits for the index in the sequence
+ # forward the model to get the logits for the index in the
+ # sequence
  logits, _ = self(idx_cond)
- # pluck the logits at the final step and scale by desired temperature
+ # pluck the logits at the final step and scale by desired
+ # temperature
  logits = logits[:, -1, :] / temperature
  # optionally crop the logits to only the top k options
  if top_k is not None:
  v, _ = torch.topk(logits, min(top_k, logits.size(-1)))
  logits[logits < v[:, [-1]]] = -float("Inf")
- # apply softmax to convert logits to (normalized) probabilities
+ # apply softmax to convert logits to (normalized)
+ # probabilities
  probs = F.softmax(logits, dim=-1)
  # sample from the distribution
  idx_next = torch.multinomial(probs, num_samples=1)