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 grouped_batch_sampler import GroupedBatchSampler, create_lengths_groups
  from lm_seqs_dataset import LmSeqsDataset
  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 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