Skip to content

Conversation

@xpkcs
Copy link
Collaborator

@xpkcs xpkcs commented May 7, 2025

This PR creates a child class of CLTTrainer called CLTTrainerRay that enables training the model using Ray. It primarily modifies the train() func to report metrics to Ray and uses Ray's approach to checkpointing. Also, it adds a new script in the root dir called tune_clt_local_ray.py that mirrors train_clt_local.py but runs CLTTrainerRay in a function that can be used by a Ray Tuner to run training runs in parallel for different hyperparameter combinations.

When tune_clt_local_ray.py is run, it will create a working dir in the specified args.output_dir, within which, for each hyperparameter combination (a "trial" in Ray terminology), Ray will create a working dir for that trial, marked by the hyperparameters for that trial. In the trial working dir, checkpoints will be saved according to args.checkpoint_interval for the model and the activation store at that training step. Reported metrics for each trial can be viewed at progress.csv and result.json in the trial's working dir. See photo below:

image

Some notes:

  • Currently, hyperparameters cannot be passed through the command line - the hps dict in tune_clt_local_ray.py must be manually edited
  • Wandb is not enabled for CLTTrainerRay

This code was tested by first generating activations for 1000 tokens from the monology/pile-uncopyrighted dataset for the EleutherAI/pythia-70m model. Then, the script was run as follows:
python tune_clt_local_ray.py --activation-path ./tutorial_activations_local_1k_pythia/EleutherAI/pythia-70m/pile-uncopyrighted_train --model-name EleutherAI/pythia-70m --num-features 2048 --training-steps 5 --n-workers 2

To visualize results, you can use tensorboard like so:
tensorboard --logdir ./clt_train_local_1746647195

Sign up for free to join this conversation on GitHub. Already have an account? Sign in to comment

Labels

None yet

Projects

None yet

Development

Successfully merging this pull request may close these issues.

2 participants