Coverage for transformer_lens/train.py: 69%

1"""Train. 

2 

3Utilities for training :class:`transformer_lens.HookedTransformer` models on autoregressive language 

4modeling tasks. 

5""" 

6 

7import dataclasses 

8from dataclasses import dataclass 

9from typing import Optional, Union 

10 

11import torch 

12import torch.optim as optim 

13from torch.optim import Optimizer 

14from torch.utils.data import DataLoader, Dataset 

15from tqdm.auto import tqdm 

16 

17from transformer_lens import utilities as utils 

18from transformer_lens.HookedTransformer import HookedTransformer 

19from transformer_lens.utilities.library_utils import is_library_available 

20 

21 

22@dataclass 

23class HookedTransformerTrainConfig: 

24 """ 

25 Configuration class to store training hyperparameters for a training run of 

26 an HookedTransformer model. 

27 Args: 

28 num_epochs (int): Number of epochs to train for 

29 batch_size (int): Size of batches to use for training 

30 lr (float): Learning rate to use for training 

31 seed (int): Random seed to use for training 

32 momentum (float): Momentum to use for training 

33 max_grad_norm (float, *optional*): Maximum gradient norm to use for 

34 weight_decay (float, *optional*): Weight decay to use for training 

35 optimizer_name (str): The name of the optimizer to use 

36 device (str or torch.device, *optional*): Device to use for training 

37 warmup_steps (int, *optional*): Number of warmup steps to use for training 

38 save_every (int, *optional*): After how many batches should a checkpoint be saved 

39 save_dir, (str, *optional*): Where to save checkpoints 

40 wandb (bool): Whether to use Weights and Biases for logging 

41 wandb_project (str, *optional*): Name of the Weights and Biases project to use 

42 print_every (int, *optional*): Print the loss every n steps 

43 max_steps (int, *optional*): Terminate the epoch after this many steps. Used for debugging. 

44 """ 

45 

46 num_epochs: int 

47 batch_size: int 

48 lr: float = 1e-3 

49 seed: int = 0 

50 momentum: float = 0.0 

51 max_grad_norm: Optional[float] = None 

52 weight_decay: Optional[float] = None 

53 optimizer_name: str = "Adam" 

54 device: Optional[Union[str, torch.device]] = None 

55 warmup_steps: int = 0 

56 save_every: Optional[int] = None 

57 save_dir: Optional[str] = None 

58 wandb: bool = False 

59 wandb_project_name: Optional[str] = None 

60 print_every: Optional[int] = 50 

61 max_steps: Optional[int] = None 

62 

63 

64def train( 

65 model: HookedTransformer, 

66 config: HookedTransformerTrainConfig, 

67 dataset: Dataset, 

68) -> HookedTransformer: 

69 """ 

70 Trains an HookedTransformer model on an autoregressive language modeling task. 

71 Args: 

72 model: The model to train 

73 config: The training configuration 

74 dataset: The dataset to train on - this function assumes the dataset is set up for autoregressive language modeling. 

75 Returns: 

76 The trained model 

77 """ 

78 

79 # Work on a copy: mutating the caller's config (wandb_project_name/device 

80 # defaults below) was a silent side effect the caller never asked for. 

81 config = dataclasses.replace(config) 

82 

83 torch.manual_seed(config.seed) 

84 model.train() 

85 

86 if config.wandb: 86 ↛ 87line 86 didn't jump to line 87 because the condition on line 86 was never true

87 if not is_library_available("wandb"): 

88 raise ImportError("Wandb is not available") 

89 

90 import wandb 

91 

92 if config.wandb_project_name is None: 

93 config.wandb_project_name = "easy-transformer" 

94 wandb.init(project=config.wandb_project_name, config=vars(config)) 

95 

96 if config.device is None: 96 ↛ 100line 96 didn't jump to line 100 because the condition on line 96 was always true

97 config.device = utils.get_device() 

98 

99 optimizer: Optimizer 

100 if config.optimizer_name in ["Adam", "AdamW"]: 100 ↛ 113line 100 didn't jump to line 113 because the condition on line 100 was always true

101 # Weight decay in Adam is implemented badly, so use AdamW instead (see PyTorch AdamW docs) 

102 if config.weight_decay is not None: 102 ↛ 103line 102 didn't jump to line 103 because the condition on line 102 was never true

103 optimizer = optim.AdamW( 

104 model.parameters(), 

105 lr=config.lr, 

106 weight_decay=config.weight_decay, 

107 ) 

108 else: 

109 optimizer = optim.Adam( 

110 model.parameters(), 

111 lr=config.lr, 

112 ) 

113 elif config.optimizer_name == "SGD": 

114 optimizer = optim.SGD( 

115 model.parameters(), 

116 lr=config.lr, 

117 weight_decay=(config.weight_decay if config.weight_decay is not None else 0.0), 

118 momentum=config.momentum, 

119 ) 

120 else: 

121 raise ValueError(f"Optimizer {config.optimizer_name} not supported") 

122 

123 scheduler = None 

124 if config.warmup_steps > 0: 124 ↛ 125line 124 didn't jump to line 125 because the condition on line 124 was never true

125 scheduler = optim.lr_scheduler.LambdaLR( 

126 optimizer, 

127 lr_lambda=lambda step: min(1.0, step / config.warmup_steps), 

128 ) 

129 

130 dataloader = DataLoader(dataset, batch_size=config.batch_size, shuffle=True) 

131 

132 model.to(config.device) 

133 

134 for epoch in tqdm(range(1, config.num_epochs + 1)): 

135 samples = 0 

136 for step, batch in tqdm(enumerate(dataloader)): 

137 tokens = batch["tokens"].to(config.device) 

138 loss = model(tokens, return_type="loss") 

139 loss.backward() 

140 if config.max_grad_norm is not None: 140 ↛ 141line 140 didn't jump to line 141 because the condition on line 140 was never true

141 torch.nn.utils.clip_grad_norm_(model.parameters(), config.max_grad_norm) 

142 optimizer.step() 

143 if config.warmup_steps > 0: 143 ↛ 144line 143 didn't jump to line 144 because the condition on line 143 was never true

144 assert scheduler is not None 

145 scheduler.step() 

146 optimizer.zero_grad() 

147 

148 samples += tokens.shape[0] 

149 

150 if config.wandb: 150 ↛ 151line 150 didn't jump to line 151 because the condition on line 150 was never true

151 wandb.log({"train_loss": loss.item(), "samples": samples, "epoch": epoch}) 

152 

153 if config.print_every is not None and step % config.print_every == 0: 153 ↛ 156line 153 didn't jump to line 156 because the condition on line 153 was always true

154 print(f"Epoch {epoch} Samples {samples} Step {step} Loss {loss.item()}") 

155 

156 if ( 156 ↛ 161line 156 didn't jump to line 161 because the condition on line 156 was never true

157 config.save_every is not None 

158 and step % config.save_every == 0 

159 and config.save_dir is not None 

160 ): 

161 torch.save(model.state_dict(), f"{config.save_dir}/model_{step}.pt") 

162 

163 if config.max_steps is not None and step >= config.max_steps: 163 ↛ 164line 163 didn't jump to line 164 because the condition on line 163 was never true

164 break 

165 

166 return model