Plotting loss curves.

train_simple.py

@@ -144,19 +144,19 @@ def shorten_value(value) -> str:
 def main(
  batch_size: int = 32,
  max_ctx: int = 1024,
- ds_name: str = "sciq",
+ ds_name: str = "glue_cola",
  loss: str = "xent",
  n_docs: int = 20000,
  n_test_docs: int = 10000,
  model_size: str = "gpt2",
- lr: Optional[float] = None,
+ lr: Optional[float] = 5e-05,
  optim: Optional[str] = None,
  epochs: int = 2,
  force_retrain: bool = False,
  seed: int = 0,
  minibatch_size_per_device: Optional[float] = None,
  train_with_dropout: bool = False,
- results_folder: str = "/tmp/results",
+ results_folder: str = "./results",
  linear_probe: bool = False,
  lr_schedule: str = "cosine_anneal",
  # Note: you can pass either weak_model_size or weak_labels_path. If you pass
@@ -297,6 +297,8 @@ def main(
  eval_every=eval_every,
  )
 
+ torch.cuda.empty_cache()
+
  if weak_ds is not None:
  weak_ds.save_to_disk(save_path + "/" + "weak_labels")
 

train_weak_to_strong.py

@@ -117,7 +117,7 @@
 def main(
  batch_size: int = 32,
  max_ctx: int = 1024,
- ds_name: str = "sciq",
+ ds_name: str = "paws_labeled_final",
  transfer_loss: Union[str, Sequence[str]] = "xent,logconf",
  n_docs: int = 10000,
  n_test_docs: int = 200,

weak_to_strong/datasets.py

@@ -161,6 +161,26 @@ def format_boolq(ex, rng):
  ),
 )
 
+def format_paws(ex, rng):
+ txt = f"Sentence 1: {ex['sentence1']} Sentence 2: {ex['sentence2']}"
+ hard_label = int(ex['label'])
+ return dict(txt=txt, hard_label=hard_label)
+
+register_dataset(
+ "paws_labeled_final", # Unique name for the dataset registration.
+ DatasetConfig(
+ loader=hf_loader("paws", "labeled_final", split_names=dict(test="validation")), 
+ formatter=format_paws
+ ),
+)
+
+def format_glue_cola(ex, rng):
+ return dict(txt=ex['sentence'], hard_label=ex['label'])
+
+register_dataset(
+ "glue_cola", 
+ DatasetConfig(loader=hf_loader("glue", "cola"), formatter=format_glue_cola),
+)
 
 VALID_DATASETS: list[str] = list(_REGISTRY.keys())
 

weak_to_strong/train.py

@@ -10,11 +10,13 @@
 import torch
 import torch_optimizer as toptim
 from transformers.modeling_utils import load_sharded_checkpoint
+from torch.utils.data import DataLoader
+from safetensors.torch import load_model
 
 import weak_to_strong.logger as logger
 from weak_to_strong.common import clear_mem
 from weak_to_strong.eval import eval_model_acc
-from weak_to_strong.loss import xent_loss
+from weak_to_strong.loss import xent_loss, logconf_loss_fn
 from weak_to_strong.model import TransformerWithHead
 
 
@@ -28,18 +30,30 @@ class ModelConfig:
  model_parallel: bool = False
  default_optimizer: str = "adam"
 
+def pad_collate(batch):
+ """
+ Custom collate function to pad sequences to the same length within a batch.
+ """
+ input_ids = [torch.tensor(item["input_ids"]) for item in batch]
+ soft_labels = [torch.tensor(item["soft_label"]) for item in batch]
+
+ padded_input_ids = torch.nn.utils.rnn.pad_sequence(input_ids, batch_first=True)
+ padded_soft_labels = torch.nn.utils.rnn.pad_sequence(soft_labels, batch_first=True)
+
+ return {"input_ids": padded_input_ids, "soft_label": padded_soft_labels}
 
 def train_model(
  model: torch.nn.Module,
  ds: datasets.Dataset,
  batch_size: int,
- lr: float = 1e-5,
+ lr: float = 1e-05,
  loss_fn: Callable = xent_loss,
  log_every: int = 10,
  eval_every: int = 100,
  eval_batch_size: int = 256,
  minibatch_size: int = 8,
  eval_ds: Optional[datasets.Dataset] = None,
+ test_ds: Optional[datasets.Dataset] = None,
  gradient_checkpointing: bool = False,
  train_with_dropout: bool = False,
  epochs: int = 1,
@@ -48,16 +62,14 @@ def train_model(
 ):
  print("LR", lr, "batch_size", batch_size, "minibatch_size", minibatch_size)
  assert batch_size % minibatch_size == 0, "batch size must be divisible by minibatch size"
- # we purposefully turn off dropout, for determinism
- # this seems to help for 1 epoch finetuning anyways
+
  if train_with_dropout:
  model.train()
  else:
  model.eval()
+
  if gradient_checkpointing:
- (
- model if hasattr(model, "gradient_checkpointing_enable") else model.module
- ).gradient_checkpointing_enable()
+ model.gradient_checkpointing_enable() if hasattr(model, "gradient_checkpointing_enable") else model.module.gradient_checkpointing_enable()
 
  nsteps = len(ds) * epochs // batch_size
 
@@ -73,55 +85,62 @@ def lr_schedule_fn(step):
  optimizer = toptim.Adafactor(model.parameters(), lr=lr)
  else:
  assert False, f"invalid optimizer {optimizer_name}, must be adam or adafactor"
+
  if lr_schedule == "cosine_anneal":
  lr_scheduler = torch.optim.lr_scheduler.CosineAnnealingLR(optimizer, nsteps)
  else:
  lr_scheduler = torch.optim.lr_scheduler.LambdaLR(optimizer, lr_schedule_fn)
+
  step = 0
  it = itertools.chain.from_iterable(itertools.repeat(ds, epochs))
  losses = []
  accuracies = []
  eval_acc_dict = {}
 
- # If the model is wrapped by DataParallel, it doesn't have a device. In this case,
- # we use GPU 0 as the output device. This sadly means that this device will store
- # a bit more data than other ones, but hopefully should not be too big of a deal.
- io_device = model.device if hasattr(model, "device") else 0
+ io_device = torch.device("cuda" if torch.cuda.is_available() else "cpu")
+ model = model.to(io_device)
+
+ def calculate_loss(model, dataset, loss_fn, batch_size, step_frac=0):
+ if dataset is None or len(dataset) == 0:
+ return float('nan')
+
+ model.eval()
+ dataloader = DataLoader(dataset, batch_size=batch_size, collate_fn=pad_collate)
+ total_loss = 0
+ count = 0
+
+ with torch.no_grad():
+ for batch in dataloader:
+ input_ids = batch["input_ids"].to(io_device)
+ labels = batch["soft_label"].to(io_device)
+ outputs = model(input_ids)
+ loss = loss_fn(outputs, labels, step_frac)
+ total_loss += loss.item()
+ count += 1
+ return total_loss / count if count > 0 else float('nan')
+
+ # Initialize val_loss and test_loss to None
+ val_loss, test_loss = None, None
 
  while step < nsteps:
  loss_tot = 0
- if eval_every and (step + 1) % eval_every == 0:
- eval_results = eval_model_acc(model, eval_ds, eval_batch_size)
- if gradient_checkpointing:
- (
- model if hasattr(model, "gradient_checkpointing_enable") else model.module
- ).gradient_checkpointing_enable()
- if train_with_dropout:
- model.train()
- eval_accs = np.mean([r["acc"] for r in eval_results])
- eval_acc_dict[step] = eval_accs
- logger.logkv("eval_accuracy", eval_accs)
  all_logits = []
  all_labels = []
  for i in range(batch_size // minibatch_size):
  try:
  mbatch = [next(it) for _ in range(minibatch_size)]
  except StopIteration:
  break
- input_ids = (
- torch.nn.utils.rnn.pad_sequence([torch.tensor(ex["input_ids"]) for ex in mbatch])
- .transpose(
- 0,
- 1,
- )
- .to(io_device)
- )
+ input_ids = torch.nn.utils.rnn.pad_sequence(
+ [torch.tensor(ex["input_ids"]) for ex in mbatch], batch_first=True
+ ).to(io_device)
  labels = torch.tensor([ex["soft_label"] for ex in mbatch]).to(io_device)
 
  logits = model(input_ids)
 
  all_logits.extend(logits.to(io_device))
  all_labels.extend(labels)
+
  all_logits = torch.stack(all_logits)
  all_labels = torch.stack(all_labels)
  loss = loss_fn(all_logits, all_labels, step_frac=step / nsteps)
@@ -130,9 +149,7 @@ def lr_schedule_fn(step):
  losses.append(loss_tot)
  accuracies.append(
  torch.mean(
- (torch.argmax(all_logits, dim=1) == torch.argmax(all_labels, dim=1)).to(
- torch.float32
- )
+ (torch.argmax(all_logits, dim=1) == torch.argmax(all_labels, dim=1)).to(torch.float32)
  ).item()
  )
  logger.logkvs(
@@ -147,22 +164,38 @@ def lr_schedule_fn(step):
  optimizer.step()
  optimizer.zero_grad()
  lr_scheduler.step()
+
  if log_every and step % log_every == 0:
+ # Calculate and print validation and test losses
+ if eval_ds is not None:
+ val_loss = calculate_loss(model, eval_ds, loss_fn, eval_batch_size, step_frac=step / nsteps)
+ if test_ds is not None:
+ test_loss = calculate_loss(model, test_ds, loss_fn, eval_batch_size, step_frac=step / nsteps)
+
  print(
- f"Step: {step}/{nsteps} Recent losses: {np.mean(losses)} {np.mean(accuracies)} {len(losses)}"
+ f"Step: {step}/{nsteps} Recent training losses: {np.mean(losses)} {np.mean(accuracies)} {len(losses)}"
  )
+ if val_loss is not None:
+ print(f"Step: {step}/{nsteps} Recent validation losses: {val_loss}")
+ if test_loss is not None:
+ print(f"Step: {step}/{nsteps} Recent test losses: {test_loss}")
+
  losses = []
  accuracies = []
+
  step += 1
  logger.dumpkvs()
+ torch.cuda.empty_cache()
+
  final_eval_results = None
- if eval_every:
+ if eval_ds is not None:
  print("Final evaluation:")
  final_eval_results = eval_model_acc(model, eval_ds, eval_batch_size)
  logger.logkv("eval_accuracy", np.mean([r["acc"] for r in final_eval_results]))
  logger.dumpkvs()
+
  return final_eval_results
-
+ 
 
 def train_and_save_model(
  model_config: ModelConfig,
@@ -195,22 +228,25 @@ def train_and_save_model(
  custom_kwargs = model_config.custom_kwargs or {}
 
  def maybe_load_model(model):
- if os.path.exists(os.path.join(save_path, "results.pkl")) and not force_retrain:
+ print("Save path: {}".format(save_path))
+ if os.path.exists(os.path.join(save_path, "results.txt")) and not force_retrain:
  print("loading from", save_path)
- checkpoint_path = os.path.join(save_path, "pytorch_model.bin")
- if not os.path.exists(checkpoint_path):
- # Assume this means we have a sharded checkpoint, and load it appropriately
- load_sharded_checkpoint(model, checkpoint_path)
- else:
- state_dict = torch.load(os.path.join(save_path, "pytorch_model.bin"))
- state_dict = {
- k.replace("transformer.module", "transformer"): v
- for (k, v) in state_dict.items()
- }
- custom_kwargs["state_dict"] = state_dict
- return True
+ checkpoint_path = os.path.join(save_path, "model.safetensors")
+ try:
+ if not os.path.exists(checkpoint_path):
+ print("using load_sharded_checkpoint")
+ load_sharded_checkpoint(model, checkpoint_path)
+ else:
+ print("using load_model")
+ state_dict = torch.load(checkpoint_path)
+ model.load_state_dict(state_dict)
+ model.to("cuda") # Ensure the model is on the correct device
+ return True
+ except Exception as e:
+ print(f"error loading model: {e}")
+ return False
  return False
-
+ 
  already_trained = False
  # Load the model
  if model_config.model_parallel:
@@ -222,6 +258,7 @@ def maybe_load_model(model):
  linear_probe=linear_probe,
  **custom_kwargs,
  )
+ model.to("cuda")
  already_trained = maybe_load_model(model)
  # slight misnomer, more like minibatch_size_per_dp_replica
  minibatch_size = minibatch_size_per_device
@@ -256,6 +293,7 @@ def maybe_load_model(model):
  lr=lr,
  epochs=epochs,
  eval_ds=test_ds,
+ test_ds=inference_ds,
  gradient_checkpointing=gradient_checkpointing,
  loss_fn=loss_fn,
  eval_batch_size=eval_batch_size,
@@ -269,7 +307,7 @@ def maybe_load_model(model):
  if save_path:
  # Note: If the model is wrapped by DataParallel, we need to unwrap it before saving
  (model if hasattr(model, "save_pretrained") else model.module).save_pretrained(
- save_path
+ save_path, safe_serialization=False
  )
  print("saved", save_path)