docstr: One optimization step: forward of student AND teacher, backward on the loss (for gradient accumulation),
  and possibly a parameter update (depending on the gradient accumulation).
  
  Input:
  ------
  input_ids: `torch.tensor(bs, seq_length)` - The token ids.
  attention_mask: `torch.tensor(bs, seq_length)` - The attention mask for self attention.
  lm_labels: `torch.tensor(bs, seq_length)` - The language modeling labels (mlm labels for MLM and clm labels for CLM).
docstr_tok: ["One","optimization","step",":","forward","of","student","AND","teacher",",","backward","on","the","loss","(","for","gradient","accumulation",")",",","and","possibly","a","parameter","update","(","depending","on","the","gradient","accumulation",")",".","Input",":","--","--","--","input_ids",":","`","torch.tensor","(","bs",",","seq_length",")","`","-","The","token","ids",".","attention_mask",":","`","torch.tensor","(","bs",",","seq_length",")","`","-","The","attention","mask","for","self","attention",".","lm_labels",":","`","torch.tensor","(","bs",",","seq_length",")","`","-","The","language","modeling","labels","(","mlm","labels","for","MLM","and","clm","labels","for","CLM",")","."]
code: def step(
      self,
      input_ids: torch.tensor,
      attention_mask: torch.tensor,
      lm_labels: torch.tensor,
  ):
      """
      One optimization step: forward of student AND teacher, backward on the loss (for gradient accumulation),
      and possibly a parameter update (depending on the gradient accumulation).
  
      Input:
      ------
      input_ids: `torch.tensor(bs, seq_length)` - The token ids.
      attention_mask: `torch.tensor(bs, seq_length)` - The attention mask for self attention.
      lm_labels: `torch.tensor(bs, seq_length)` - The language modeling labels (mlm labels for MLM and clm labels for CLM).
      """
      if self.mlm:
          student_outputs = self.student(
              input_ids=input_ids, attention_mask=attention_mask
          )  # (bs, seq_length, voc_size)
          with torch.no_grad():
              teacher_outputs = self.teacher(
                  input_ids=input_ids, attention_mask=attention_mask
              )  # (bs, seq_length, voc_size)
      else:
          student_outputs = self.student(
              input_ids=input_ids, attention_mask=None
          )  # (bs, seq_length, voc_size)
          with torch.no_grad():
              teacher_outputs = self.teacher(
                  input_ids=input_ids, attention_mask=None
              )  # (bs, seq_length, voc_size)
      s_logits, s_hidden_states = (
          student_outputs["logits"],
          student_outputs["hidden_states"],
      )
      t_logits, t_hidden_states = (
          teacher_outputs["logits"],
          teacher_outputs["hidden_states"],
      )
      assert s_logits.size() == t_logits.size()
  
      # https://github.com/peterliht/knowledge-distillation-pytorch/blob/master/model/net.py#L100
      # https://github.com/peterliht/knowledge-distillation-pytorch/issues/2
      if self.params.restrict_ce_to_mask:
          mask = (
              (lm_labels > -1).unsqueeze(-1).expand_as(s_logits)
          )  # (bs, seq_length, voc_size)
      else:
          mask = attention_mask.unsqueeze(-1).expand_as(
              s_logits
          )  # (bs, seq_length, voc_size)
      s_logits_slct = torch.masked_select(
          s_logits, mask
      )  # (bs * seq_length * voc_size) modulo the 1s in mask
      s_logits_slct = s_logits_slct.view(
          -1, s_logits.size(-1)
      )  # (bs * seq_length, voc_size) modulo the 1s in mask
      t_logits_slct = torch.masked_select(
          t_logits, mask
      )  # (bs * seq_length * voc_size) modulo the 1s in mask
      t_logits_slct = t_logits_slct.view(
          -1, s_logits.size(-1)
      )  # (bs * seq_length, voc_size) modulo the 1s in mask
      assert t_logits_slct.size() == s_logits_slct.size()
  
      loss_ce = (
          self.ce_loss_fct(
              nn.functional.log_softmax(
                  s_logits_slct / self.temperature, dim=-1
              ),
              nn.functional.softmax(
                  t_logits_slct / self.temperature, dim=-1
              ),
          )
          * (self.temperature) ** 2
      )
      loss = self.alpha_ce * loss_ce
  
      if self.alpha_mlm > 0.0:
          loss_mlm = self.lm_loss_fct(
              s_logits.view(-1, s_logits.size(-1)), lm_labels.view(-1)
          )
          loss += self.alpha_mlm * loss_mlm
      if self.alpha_clm > 0.0:
          shift_logits = s_logits[..., :-1, :].contiguous()
          shift_labels = lm_labels[..., 1:].contiguous()
          loss_clm = self.lm_loss_fct(
              shift_logits.view(-1, shift_logits.size(-1)),
              shift_labels.view(-1),
          )
          loss += self.alpha_clm * loss_clm
  
      if self.alpha_mse > 0.0:
          loss_mse = self.mse_loss_fct(
              s_logits_slct, t_logits_slct
          ) / s_logits_slct.size(
              0
          )  # Reproducing batchmean reduction
          loss += self.alpha_mse * loss_mse
      if self.alpha_cos > 0.0:
          s_hidden_states = s_hidden_states[-1]  # (bs, seq_length, dim)
          t_hidden_states = t_hidden_states[-1]  # (bs, seq_length, dim)
          mask = attention_mask.unsqueeze(-1).expand_as(
              s_hidden_states
          )  # (bs, seq_length, dim)
          assert s_hidden_states.size() == t_hidden_states.size()
          dim = s_hidden_states.size(-1)
  
          s_hidden_states_slct = torch.masked_select(
              s_hidden_states, mask
          )  # (bs * seq_length * dim)
          s_hidden_states_slct = s_hidden_states_slct.view(
              -1, dim
          )  # (bs * seq_length, dim)
          t_hidden_states_slct = torch.masked_select(
              t_hidden_states, mask
          )  # (bs * seq_length * dim)
          t_hidden_states_slct = t_hidden_states_slct.view(
              -1, dim
          )  # (bs * seq_length, dim)
  
          target = s_hidden_states_slct.new(
              s_hidden_states_slct.size(0)
          ).fill_(
              1
          )  # (bs * seq_length,)
          loss_cos = self.cosine_loss_fct(
              s_hidden_states_slct, t_hidden_states_slct, target
          )
          loss += self.alpha_cos * loss_cos
  
      self.total_loss_epoch += loss.item()
      self.last_loss = loss.item()
      self.last_loss_ce = loss_ce.item()
      if self.alpha_mlm > 0.0:
          self.last_loss_mlm = loss_mlm.item()
      if self.alpha_clm > 0.0:
          self.last_loss_clm = loss_clm.item()
      if self.alpha_mse > 0.0:
          self.last_loss_mse = loss_mse.item()
      if self.alpha_cos > 0.0:
          self.last_loss_cos = loss_cos.item()
  
      self.optimize(loss)
  
      self.n_sequences_epoch += input_ids.size(0)
  
imports: import math
  import os
  import time
  import psutil
  import torch
  from torch import nn
  from torch.optim import AdamW
  from torch.utils.data import BatchSampler, DataLoader, RandomSampler
  from torch.utils.data.distributed import DistributedSampler
  from tqdm import tqdm
  from grouped_batch_sampler import GroupedBatchSampler, create_lengths_groups
  from lm_seqs_dataset import LmSeqsDataset
  from transformers import get_linear_schedule_with_warmup
  from utils import logger
label_desc: Use image node_id 6 for calling the Distiller obj's underlying member method code with example usage: obj.step(input_ids, attention_mask, lm_labels) without return types