Coverage for transformer_lens/tools/training.py: 71%

1"""Train loop for TransformerLens models. 

2 

3Utilities for training models on autoregressive language modeling tasks. 

4Typed against the ``model_protocol`` surface (``__call__`` with 

5``return_type="loss"`` plus standard ``nn.Module`` parameter access), so any 

6conforming model — ``TransformerBridge`` foremost — works through this loop. 

7""" 

8 

9import dataclasses 

10from dataclasses import dataclass 

11from typing import Optional, Union 

12 

13import torch 

14import torch.optim as optim 

15from torch.optim import Optimizer 

16from torch.utils.data import DataLoader, Dataset 

17from tqdm.auto import tqdm 

18 

19from transformer_lens import utilities as utils 

20from transformer_lens.model_protocol import TrainableTransformerLensModel 

21from transformer_lens.utilities.library_utils import is_library_available 

22 

23 

24@dataclass 

25class TrainConfig: 

26 """Configuration class to store training hyperparameters for a training run. 

27 

28 Args: 

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

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

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

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

33 momentum (float): Momentum to use for training 

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

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

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

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

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

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

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

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

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

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

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

45 """ 

46 

47 num_epochs: int 

48 batch_size: int 

49 lr: float = 1e-3 

50 seed: int = 0 

51 momentum: float = 0.0 

52 max_grad_norm: Optional[float] = None 

53 weight_decay: Optional[float] = None 

54 optimizer_name: str = "Adam" 

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

56 warmup_steps: int = 0 

57 save_every: Optional[int] = None 

58 save_dir: Optional[str] = None 

59 wandb: bool = False 

60 wandb_project_name: Optional[str] = None 

61 print_every: Optional[int] = 50 

62 max_steps: Optional[int] = None 

63 

64 

65def train( 

66 model: TrainableTransformerLensModel, 

67 config: TrainConfig, 

68 dataset: Dataset, 

69) -> TrainableTransformerLensModel: 

70 """Train a model on an autoregressive language modeling task. 

71 

72 Args: 

73 model: The model to train (TrainableTransformerLensModel: callable with 

74 ``return_type="loss"`` and exposing torch parameters) 

75 config: The training configuration 

76 dataset: The dataset to train on - assumed set up for autoregressive language modeling. 

77 

78 Returns: 

79 The trained model 

80 """ 

81 

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

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

84 config = dataclasses.replace(config) 

85 

86 torch.manual_seed(config.seed) 

87 model.train() 

88 

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

90 if not is_library_available("wandb"): 

91 raise ImportError("Wandb is not available") 

92 

93 import wandb 

94 

95 if config.wandb_project_name is None: 

96 config.wandb_project_name = "easy-transformer" 

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

98 

99 if config.device is None: 

100 config.device = utils.get_device() 

101 

102 optimizer: Optimizer 

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

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

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

106 optimizer = optim.AdamW( 

107 model.parameters(), 

108 lr=config.lr, 

109 weight_decay=config.weight_decay, 

110 ) 

111 else: 

112 optimizer = optim.Adam( 

113 model.parameters(), 

114 lr=config.lr, 

115 ) 

116 elif config.optimizer_name == "SGD": 

117 optimizer = optim.SGD( 

118 model.parameters(), 

119 lr=config.lr, 

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

121 momentum=config.momentum, 

122 ) 

123 else: 

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

125 

126 scheduler = None 

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

128 scheduler = optim.lr_scheduler.LambdaLR( 

129 optimizer, 

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

131 ) 

132 

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

134 

135 model.to(config.device) 

136 

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

138 samples = 0 

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

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

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

142 loss.backward() 

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

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

145 optimizer.step() 

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

147 assert scheduler is not None 

148 scheduler.step() 

149 optimizer.zero_grad() 

150 

151 samples += tokens.shape[0] 

152 

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

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

155 

156 if config.print_every is not None and step % config.print_every == 0: 

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

158 

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

160 config.save_every is not None 

161 and step % config.save_every == 0 

162 and config.save_dir is not None 

163 ): 

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

165 

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

167 break 

168 

169 return model