Training improvements #17

evanhanders · 2024-08-22T22:31:31Z

A bunch of small changes to make training smoother & a bit more robust; most importantly:

use_single_loss is now true and default optimizer arguments for Adam use beta = (0.9, 0.9), as opposed to (0.9, 0.999). I find that this makes training with a single loss step more robust, but please check this for your cases!
progress bars are now reused during training and a log is output below them instead of them being recreated every epoch.
removes redundancy between training_args definitions
puts SIIT loss calculation into its own function like IIT loss.
Changes how nodes are sampled for SIIT and makes this a training argument; now there's an option so that SIIT can randomly sample from all of the nodes that aren't in the circuit and ablate those, rather than just ablating one.

mypy type checking and pytest tests/ passes, but it's possible some downstream stuff broke? Unclear.

… type-check.

…cls to training_args

cybershiptrooper · 2024-08-23T05:38:42Z

iit/model_pairs/base_model_pair.py

@@ -18,6 +20,24 @@
 from iit.utils.index import Ix, TorchIndex
 from iit.utils.metric import MetricStoreCollection, MetricType

+def in_notebook() -> bool:


I think we can do this by just importing tqdm: much cleaner that way. (at least according to this)

I tried using just tqdm, but it definitely didn't work in notebook mode. I think I hunted down all of the print statements, too.

Cleaner compromise than what's here now: moved this block to utils/tqdm.py, and added from iit.utils.tqdm import tqdm.

cybershiptrooper · 2024-08-23T05:38:46Z

iit/model_pairs/base_model_pair.py

@@ -177,7 +197,7 @@ def get_IIT_loss_over_batch(
        hl_output, ll_output = self.do_intervention(base_input, ablation_input, hl_node)
        label_idx = self.get_label_idxs()
        # IIT loss is only computed on the tokens we care about
-        loss = loss_fn(ll_output[label_idx.as_index], hl_output[label_idx.as_index])
+        loss = loss_fn(ll_output[label_idx.as_index].to(hl_output.device), hl_output[label_idx.as_index])


We should probably just raise if dataset, hl_model and ll_model aren't on the same device during init/starting training. This usually just hides the main problem and makes it harder to find bugs.

Makes sense, I'll add an assert to the beginning of train() and remove all of these.

cybershiptrooper · 2024-08-23T05:39:16Z

iit/model_pairs/base_model_pair.py


-            if early_stop and self._check_early_stop_condition(test_metrics):
-                break
+                    epoch_pbar.update(1)


Would be nicer if we can move the entire logic to _print_and_log_metrics. current_epoch_log can remain there. And logging it to wandb might be useful as well!

Moved this logic to _print_and_log_metrics. I think everything that makes up the string is already being logged to wandb.

cybershiptrooper · 2024-08-23T05:41:11Z

iit/model_pairs/base_model_pair.py

        for metric in metrics:
-            print(metric, end=", ")
+            if metric.type == MetricType.ACCURACY:
+                current_epoch_log += f"{metric.get_name()}: {metric.get_value():.2f}, "


str(metric) does this automatically

changed to current_epoch_log += str(metric) + ", "

iit/model_pairs/iit_behavior_model_pair.py

cybershiptrooper · 2024-08-23T05:41:38Z

iit/model_pairs/strict_iit_model_pair.py

@@ -21,14 +21,9 @@ def __init__(
            training_args: dict = {}
            ):
        default_training_args = {
-            "batch_size": 256,


Would be really helpful if we could maintain the default args as they were before. Or at least store the default hyperparams we used before in some config for reproducibility.

I think all the defaults are preserved (they were just set in multiple places) EXCEPT I did change use_single_loss and optimizer_kwargs. I'll change those back to the defaults from before.

cybershiptrooper · 2024-08-23T05:41:43Z

iit/model_pairs/strict_iit_model_pair.py

            "strict_weight": 1.0,
-            "clip_grad_norm": 1.0,
+            "strict_weight_schedule" : lambda s, i: s,


This is a cool idea!

Maybe it is better to implement it as

@property def strict_weight_at_epoch(self): return self.training_args.strict_weight_schedule(<args_from_self>)

Instead of changing the strict weight variable after each epoch? (or a method like strict_weight_for_epoch = self.get_scheduled_strict_weight() and then calculate the loss).

This lambda is also throwing me off a bit, maybe renaming the args will make it clearer...

It also seems like this is achievable by using different optimisers/lrs for each loss (and not using single loss)? No idea which one's better though...

Hm, unclear to me which is the right way to go right now. I think it's best to remove it (right now it's not doing anything) and if you find that this is a useful idea down the road you can add it how you see fit?

cybershiptrooper · 2024-08-23T05:41:54Z

iit/model_pairs/strict_iit_model_pair.py

-        iit_loss = 0
-        ll_loss = 0
-        behavior_loss = 0
+        iit_loss = t.zeros(1)


I'm not completely sure why this is needed- You can usually add floats and tensors without messing up the grads, right?

This is for mypy type-checking. step_on_loss expects a Tensor instead of a float and the .item() call at the end is a type error if this isn't a Tensor.

I think the right way to resolve this is to remove the if isintance(Tensor) logic at the end of the function, since it's now always a tensor. I'll do that.

cybershiptrooper · 2024-08-23T05:42:51Z

iit/model_pairs/base_model_pair.py

@@ -8,6 +9,7 @@
 from torch.utils.data import DataLoader
 from tqdm import tqdm # type: ignore
 from transformer_lens.hook_points import HookedRootModule, HookPoint # type: ignore
+from IPython.display import clear_output


Why is this needed here? Don't think it is being used...

It's not! Good catch, that was leftover from getting tqdm stuff working.

cybershiptrooper · 2024-08-23T05:48:41Z

iit/model_pairs/base_model_pair.py

-            )
+
+        # Set seed before iterating on loaders for reproduceablility.
+        t.manual_seed(training_args["seed"])


Is it possible to use a generator for loaders like we do for numpy? I think I used to set this once globally in the training script before- my bad. :(

I'm not totally sure? I got this solution here. It seems like the random operation is set when the dataloader is turned into an iterable, and someone could use torch functions between initializing and training the model pair, which could hinder reproducibility without putting something here.

cybershiptrooper · 2024-08-23T05:53:19Z

mypy type checking and pytest tests/ passes, but it's possible some downstream stuff broke? Unclear.

It doesn't look particularly problematic. I'll have a more careful look in a while.
Might be worth checking if circuits benchmark tests pass after updating its poetry version...

Thanks for adding these!

Will definitely check my cases though. This seems important in general- somehow I can't reproduce the 4 new trained cases after pulling the newer PRs. Maybe these changes help. :")

…nch to break

evanhanders · 2024-08-23T16:27:34Z

OK! I think I responded to all of your changes and pushed updates. Also found a problem that was causing circuits-bench tests to fail and fixed that so they all are passing on my end.

I'll be offline for the next three weeks starting in a few hours, so if there are other problems / stylistic things, please feel free to edit the branch of my repo / this PR to get those fixed!

evanhanders · 2024-08-23T16:43:58Z

Also I just added back in one .to(device) in the eval() step. It's really helpful for me to not have my entire dataset on cuda / mps, especially when training successive models in a notebook, so putting the dataset labels on the model's device in run_eval_step is helpful.

cybershiptrooper · 2024-08-23T23:46:23Z

Great! The changes look fine now.

Merging.

Thanks for the PR!

evanhanders added 7 commits August 12, 2024 15:56

improves training interface

6c5f4ed

A bunch of small changes to make training smoother; need to debug and…

1c089cf

… type-check.

New SIIT sampling protocol, Adds optimizer_kwargs to training_args

97c18d3

more training fixups (args, etc)

5e5eaa0

merge

855807a

fixes typing errors

7734db5

Defaults: use_single_step = True, betas = (0.9, 0.9); adds optimzier_…

18e467e

…cls to training_args

evanhanders requested a review from cybershiptrooper August 22, 2024 22:31

evanhanders marked this pull request as ready for review August 22, 2024 22:32

cybershiptrooper requested changes Aug 23, 2024

View reviewed changes

evanhanders added 2 commits August 23, 2024 09:09

cleans up PR per feedback

0d1b24f

reverts some indexing in IITBehaviorModelPair that caused circuits-be…

313309e

…nch to break

bug fix and adds back in one .to(device)

e4eba0a

cybershiptrooper merged commit e0be350 into cybershiptrooper:main Aug 23, 2024
3 checks passed

Provide feedback

Saved searches

Use saved searches to filter your results more quickly

Training improvements #17

Training improvements #17

evanhanders commented Aug 22, 2024

cybershiptrooper Aug 23, 2024

evanhanders Aug 23, 2024

cybershiptrooper Aug 23, 2024

evanhanders Aug 23, 2024

cybershiptrooper Aug 23, 2024

evanhanders Aug 23, 2024

cybershiptrooper Aug 23, 2024

evanhanders Aug 23, 2024

cybershiptrooper Aug 23, 2024

evanhanders Aug 23, 2024

cybershiptrooper Aug 23, 2024

evanhanders Aug 23, 2024

cybershiptrooper Aug 23, 2024

evanhanders Aug 23, 2024

cybershiptrooper Aug 23, 2024

evanhanders Aug 23, 2024

cybershiptrooper Aug 23, 2024

evanhanders Aug 23, 2024

cybershiptrooper commented Aug 23, 2024 •

edited

Loading

evanhanders commented Aug 23, 2024

evanhanders commented Aug 23, 2024

cybershiptrooper commented Aug 23, 2024 •

edited

Loading

Training improvements #17

Training improvements #17

Conversation

evanhanders commented Aug 22, 2024

Choose a reason for hiding this comment

Choose a reason for hiding this comment

Choose a reason for hiding this comment

Choose a reason for hiding this comment

Choose a reason for hiding this comment

Choose a reason for hiding this comment

Choose a reason for hiding this comment

Choose a reason for hiding this comment

Choose a reason for hiding this comment

Choose a reason for hiding this comment

Choose a reason for hiding this comment

Choose a reason for hiding this comment

Choose a reason for hiding this comment

Choose a reason for hiding this comment

Choose a reason for hiding this comment

Choose a reason for hiding this comment

Choose a reason for hiding this comment

Choose a reason for hiding this comment

cybershiptrooper commented Aug 23, 2024 • edited Loading

evanhanders commented Aug 23, 2024

evanhanders commented Aug 23, 2024

cybershiptrooper commented Aug 23, 2024 • edited Loading

cybershiptrooper commented Aug 23, 2024 •

edited

Loading

cybershiptrooper commented Aug 23, 2024 •

edited

Loading