Skip to content

Training

The training stage fits a fresh causal language model on tokenized timelines. All three trainers read a training.yaml configuration, pull the training and validation splits out of the processed data, and wrap HuggingFace's Trainer so that sequences of event tokens are modeled autoregressively (each token predicts the next). Time-aware position ids for time-based RoPE and a custom loss are wired in when the configuration requests them. On completion the model weights and the exact configuration used are written to output_home.

Trainer

The workhorse. Trainer builds a model from scratch from a HuggingFace architecture preset (its vocabulary, BOS, and EOS come from tokenizer.yaml), packs the training and tuning splits into fixed-length sequences, and runs standard next-token training. Its collate function ties labels to input_ids for causal language modeling and, when time-based RoPE is configured, derives per-token position_ids from elapsed time. Runs are tracked in Weights & Biases, and train() can resume from the latest checkpoint.

Bases: Configurable

the meds format dumps training (train), validation (tuning), and test (held_out) data into the same file; we need to start by fishing out training and validation data

Source code in src/cotorra/trainer.py 
 35
 36
 37
 38
 39
 40
 41
 42
 43
 44
 45
 46
 47
 48
 49
 50
 51
 52
 53
 54
 55
 56
 57
 58
 59
 60
 61
 62
 63
 64
 65
 66
 67
 68
 69
 70
 71
 72
 73
 74
 75
 76
 77
 78
 79
 80
 81
 82
 83
 84
 85
 86
 87
 88
 89
 90
 91
 92
 93
 94
 95
 96
 97
 98
 99
100
101
102
103
104
105
106
107
108
109
110
111
112
113
114
115
116
117
118
119
120
121
122
123
124
125
126
127
128
129
130
131
132
class Trainer(Configurable):
    """the meds format dumps training (train), validation (tuning), and test (held_out)
    data into the same file;
    we need to start by fishing out training and validation data"""

    default_file = "training.yaml"

    def __init__(
        self,
        training_cfg: pathlib.Path | str = None,
        processed_data_home: pathlib.Path | str = None,
        output_home: pathlib.Path | str = None,
        **kwargs,
    ):
        super().__init__(training_cfg, **kwargs)

        self.processed_data_home, self.output_home = map(
            lambda p: pathlib.Path(p).expanduser().resolve(),
            [processed_data_home, output_home],
        )

        self.tkzr_cfg = OmegaConf.load(self.processed_data_home / "tokenizer.yaml")
        self.loss = (
            Loss(self.cfg, self.tkzr_cfg).custom_loss if self.cfg.custom_loss else None
        )
        self.run_name = self.cfg.get("run_name", self.cfg.wandb.get("run_name", ""))
        self.loader = Loader(training_cfg, self.processed_data_home)

        self.trainer = TrainerWithCustomLoss(
            model_init=self.model_init,
            data_collator=self.collate_fn,
            compute_loss_func=self.loss,
            train_dataset=self.loader.get_train_data(),
            eval_dataset=self.loader.get_tuning_data(),
            args=TrainingArguments(
                output_dir=str(self.output_home), **self.cfg.training_args
            ),
        )
        self.model = self.trainer.model

        os.environ["WANDB_PROJECT"] = self.cfg.get("wandb", {}).get(
            "project", "cotorra"
        )
        os.environ["WANDB_NAME"] = self.cfg.get("wandb", {}).get("run_name", "cotorra")

    def model_init(self):
        conf_param = dict(
            vocab_size=len(self.tkzr_cfg.lookup),
            bos_token_id=self.tkzr_cfg.lookup.BOS,
            eos_token_id=self.tkzr_cfg.lookup.EOS,
        )
        config = AutoConfig.from_pretrained(
            self.cfg.model.model_name, **conf_param, **self.cfg.model.model_args
        )
        mdl = AutoModelForCausalLM.from_config(config)
        self.logger.info(
            "Loaded model {name} with {num} params ({dtype}).".format(
                name=self.cfg.model.model_name,
                num=sum(p.numel() for p in mdl.parameters()),
                dtype=next(mdl.parameters()).dtype,
            )
        )

        return mdl

    def collate_fn(self, batch):
        input_ids = t.stack([x["input_ids"] for x in batch])
        if "time_based_rope" not in self.cfg:
            return {"input_ids": input_ids, "labels": input_ids}
        else:
            p_ids = (
                t.stack([x["s_elapsed"] for x in batch])
                / self.cfg.time_based_rope.sec_per_pos_id
            )
            p_ids += t.arange(p_ids.shape[-1], device=p_ids.device, dtype=p_ids.dtype)
            return {"input_ids": input_ids, "labels": input_ids, "position_ids": p_ids}

    def train(self, resume_from_checkpoint: bool = False, verbose: bool = False):
        if resume_from_checkpoint:
            try:
                self.trainer.train(resume_from_checkpoint=True)
            except Exception as e:
                self.logger.warning(f"Encountered {e} on resume from checkpoint.")
                self.trainer.train()
        else:
            self.trainer.train()

        self.trainer.model.save_pretrained(self.output_home / f"mdl-{self.run_name}")

        with open(self.output_home / f"mdl-{self.run_name}-training.yaml", "w") as f:
            f.write(OmegaConf.to_yaml(self.cfg))

        if verbose:
            self.logger.summarize_trained_model(
                model=self.trainer.model,
                bos_token_id=self.tkzr_cfg.lookup["BOS"],
                reverse={v: k for k, v in self.tkzr_cfg.lookup.items()},
            )

TrainerDP (differential privacy)

TrainerDP extends Trainer with differentially private training via Opacus. It attaches a PrivacyEngine that replaces the model, optimizer, and data loader with DP-aware versions performing per-sample gradient clipping (max_grad_norm) and Gaussian noise injection (noise_multiplier); after training it reports the achieved (epsilon, delta) privacy guarantee. Because per-sample gradients are incompatible with gradient accumulation, DP training requires gradient_accumulation_steps=1, and the model is unwrapped from its GradSampleModule before being saved.

Bases: Trainer

Source code in src/cotorra/trainer_dp.py 
 67
 68
 69
 70
 71
 72
 73
 74
 75
 76
 77
 78
 79
 80
 81
 82
 83
 84
 85
 86
 87
 88
 89
 90
 91
 92
 93
 94
 95
 96
 97
 98
 99
100
101
102
103
104
105
106
107
108
109
110
111
112
113
114
115
116
117
118
119
120
121
122
123
124
125
126
127
128
129
130
131
132
133
134
135
136
137
138
139
140
141
142
143
144
145
146
147
148
149
150
151
152
153
154
class TrainerDP(Trainer):
    def __init__(
        self,
        training_cfg: pathlib.Path | str = None,
        processed_data_home: pathlib.Path | str = None,
        output_home: pathlib.Path | str = None,
        noise_multiplier: float | None = None,
        max_grad_norm: float | None = None,
        **kwargs,
    ):
        super().__init__(
            training_cfg=training_cfg,
            processed_data_home=processed_data_home,
            output_home=output_home,
            privacy_parameters={
                k: v
                for k, v in {
                    "noise_multiplier": noise_multiplier,
                    "max_grad_norm": max_grad_norm,
                }.items()
                if v is not None
            },
            **kwargs,
        )

        self.trainer = TrainerWithCustomLossDP(
            model=self.model,
            data_collator=self.collate_fn,
            compute_loss_func=self.loss,
            train_dataset=self.loader.get_train_data(),
            eval_dataset=self.loader.get_tuning_data(),
            args=TrainingArguments(
                output_dir=str(self.output_home), **self.cfg.training_args
            ),
        )

        self.trainer.create_optimizer()

        # HuggingFace Trainer calls get_train_dataloader() internally, so we must
        # ensure get_train_dataloader() returns the Opacus-wrapped DataLoader.
        privacy_engine = opacus.PrivacyEngine()
        noise_multiplier = self.cfg.get("privacy_parameters", {}).get(
            "noise_multiplier", 1.0
        )
        max_grad_norm = self.cfg.get("privacy_parameters", {}).get("max_grad_norm", 1.0)
        self.logger.info(
            f"Making private with {noise_multiplier=} and {max_grad_norm=}"
        )
        dp_model, dp_optimizer, dp_dataloader = privacy_engine.make_private(
            module=self.trainer.model,
            optimizer=self.trainer.optimizer,
            data_loader=self.trainer.get_train_dataloader(),
            noise_multiplier=noise_multiplier,
            max_grad_norm=max_grad_norm,
            poisson_sampling=False,
        )
        self.trainer.model = dp_model
        self.trainer.optimizer = dp_optimizer
        self.trainer.set_dp_train_dataloader(dp_dataloader)
        self.model = self.trainer.model
        self.privacy_engine = privacy_engine

    def train(self, verbose=False):
        """Override train to properly handle saving the DP-wrapped model."""
        self.trainer.train()

        self.logger.info("For (epsilon, delta)-differential privacy:")
        for delta in [1e-5, 1e-4, 1e-3]:
            self.logger.info(
                "delta={delta} gives epsilon={epsilon:.3e}".format(
                    delta=delta,
                    epsilon=self.privacy_engine.accountant.get_epsilon(delta=delta),
                )
            )

        # Unwrap the model from GradSampleModule before saving
        unwrapped_model = self.trainer.model._module
        unwrapped_model.save_pretrained(self.output_home / f"mdl-{self.run_name}")

        with open(self.output_home / f"mdl-{self.run_name}-training.yaml", "w") as f:
            f.write(OmegaConf.to_yaml(self.cfg))

        if verbose:
            self.logger.summarize_trained_model(
                model=unwrapped_model,
                bos_token_id=self.tkzr_cfg.lookup["BOS"],
                reverse={v: k for k, v in self.tkzr_cfg.lookup.items()},
            )

train(verbose=False)

Override train to properly handle saving the DP-wrapped model.

Source code in src/cotorra/trainer_dp.py 
129
130
131
132
133
134
135
136
137
138
139
140
141
142
143
144
145
146
147
148
149
150
151
152
153
154
def train(self, verbose=False):
    """Override train to properly handle saving the DP-wrapped model."""
    self.trainer.train()

    self.logger.info("For (epsilon, delta)-differential privacy:")
    for delta in [1e-5, 1e-4, 1e-3]:
        self.logger.info(
            "delta={delta} gives epsilon={epsilon:.3e}".format(
                delta=delta,
                epsilon=self.privacy_engine.accountant.get_epsilon(delta=delta),
            )
        )

    # Unwrap the model from GradSampleModule before saving
    unwrapped_model = self.trainer.model._module
    unwrapped_model.save_pretrained(self.output_home / f"mdl-{self.run_name}")

    with open(self.output_home / f"mdl-{self.run_name}-training.yaml", "w") as f:
        f.write(OmegaConf.to_yaml(self.cfg))

    if verbose:
        self.logger.summarize_trained_model(
            model=unwrapped_model,
            bos_token_id=self.tkzr_cfg.lookup["BOS"],
            reverse={v: k for k, v in self.tkzr_cfg.lookup.items()},
        )

Tuner

Tuner extends Trainer to run an Optuna hyperparameter search (over learning rate and gradient-accumulation steps) before the final training run. The best trial's hyperparameters are copied back onto the trainer, the model is retrained with them, and the weights and configuration are saved as usual.

Bases: Trainer

Source code in src/cotorra/tuner.py 
12
13
14
15
16
17
18
19
20
21
22
23
24
25
26
27
28
29
30
31
32
33
34
35
36
37
38
39
40
41
42
class Tuner(Trainer):
    def __init__(self, **kwargs):
        super().__init__(**kwargs)

    @staticmethod
    def optuna_hp_space(trial):
        return {
            "learning_rate": trial.suggest_float("learning_rate", 1e-4, 5e-4, log=True),
            "gradient_accumulation_steps": trial.suggest_int(
                "gradient_accumulation_steps", 1, 3
            ),
        }

    def train(self, verbose=False):
        best_trial = self.trainer.hyperparameter_search(
            hp_space=self.optuna_hp_space, **self.cfg.tuning_args
        )
        for n, v in best_trial.hyperparameters.items():
            setattr(self.trainer.args, n, v)
        self.trainer.train()
        self.trainer.model.save_pretrained(self.output_home / f"mdl-{self.run_name}")

        with open(self.output_home / f"mdl-{self.run_name}-tuning.yaml", "w") as f:
            f.write(OmegaConf.to_yaml(self.cfg))

        if verbose:
            self.logger.summarize_trained_model(
                model=self.trainer.model,
                bos_token_id=self.tkzr_cfg.lookup["BOS"],
                reverse={v: k for k, v in self.tkzr_cfg.lookup.items()},
            )