forked from EleutherAI/gpt-neox
-
Notifications
You must be signed in to change notification settings - Fork 0
/
mup_substitute.py
212 lines (196 loc) · 7.62 KB
/
mup_substitute.py
1
2
3
4
5
6
7
8
9
10
11
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
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
133
134
135
136
137
138
139
140
141
142
143
144
145
146
147
148
149
150
151
152
153
154
155
156
157
158
159
160
161
162
163
164
165
166
167
168
169
170
171
172
173
174
175
176
177
178
179
180
181
182
183
184
185
186
187
188
189
190
191
192
193
194
195
196
197
198
199
200
201
202
203
204
205
206
207
208
209
210
211
212
"""
Helper functions for performing coord check.
"""
import os
from copy import copy
from itertools import product
import numpy as np
import pandas as pd
import torch
import torch.nn.functional as F
from mup import coord_check as mup_coord_check
from megatron.training import train_step
def _get_coord_data(
neox_args,
timers,
lr_scheduler,
models,
dataloader,
optcls,
nsteps=3,
dict_in_out=False,
flatten_input=False,
flatten_output=False,
output_name="loss",
lossfn="xent",
filter_module_by_name=None,
fix_data=True,
cuda=True,
nseeds=1,
output_fdict=None,
input_fdict=None,
param_fdict=None,
show_progress=True,
one_hot_target=False,
):
df = []
for i in range(nseeds):
torch.manual_seed(i)
for width, model in models.items():
model = model()
model.train()
optimizer = optcls(model)
for step in range(nsteps + 1):
remove_hooks = []
# add hooks
for name, module in model.named_modules():
if filter_module_by_name and not filter_module_by_name(name):
continue
remove_hooks.append(
module.register_forward_hook(
mup_coord_check._record_coords(
df,
width,
name,
step + 1,
output_fdict=output_fdict,
input_fdict=input_fdict,
param_fdict=param_fdict,
)
)
)
# train for a step
loss_dict, skipped_iter = train_step(
neox_args=neox_args,
timers=timers,
data_iterator=dataloader,
model=model,
optimizer=optimizer,
lr_scheduler=lr_scheduler,
)
# remove hooks
for handle in remove_hooks:
handle.remove()
import gc
del model
gc.collect()
return pd.DataFrame(df)
def get_coord_data(
neox_args,
timers,
lr_scheduler,
models,
dataloader,
optimizer="sgd",
lr=None,
mup=True,
filter_trainable_by_name=None,
**kwargs
):
"""Get coord data for coord check.
Train the models in `models` with data from `dataloader` and optimizer
specified by `optimizer` and `lr` for `nsteps` steps, and record coordinate
statistics specified by `output_fdict`, `input_fdict`, `param_fdict`. By
default, only `l1` is computed for output activations of each module.
This function wraps around `_get_coord_data`, with the main difference being
user can specify common optimizers via a more convenient interface.
Inputs:
models:
a dict of lazy models, where the keys are numbers indicating width.
Each entry of `models` is a function that instantiates a model given
nothing.
dataloader:
an iterator whose elements are either Huggingface style dicts, if
`dict_in_out` is True, or (input, label). If `fix_data` is True
(which is the default), then only the first element of `dataloader`
is used in a loop and the rest of `dataloder` is ignored.
optimizer:
a string in `['sgd', 'adam', 'adamw']`, with default being `'sgd'`.
lr:
learning rate. By default is 0.1 for `'sgd'` and 1e-3 for others.
mup:
If True, then use the optimizer from `mup.optim`; otherwise, use the
one from `torch.optim`.
filter_trainable_by_name:
a function that returns a bool given module names (from
`model.named_modules()`), or None. If not None, then only modules
whose name yields True will be trained.
nsteps:
number of steps to train the model
dict_in_out:
whether the data loader contains Huggingface-style dict input and
output. Default: False
flatten_input:
if not `dict_in_out`, reshape the input to be
`input.view(input.shape[0], -1)`. Typically used for testing MLPs.
flatten_output:
if not `dict_in_out`, reshape the label to be `label.view(-1,
input.shape[-1])`.
output_name:
if `dict_in_out`, this is the key for the loss value if the output
is a dict. If the output is not a dict, then we assume the first
element of the output is the loss.
lossfn:
loss function to use if not `dict_in_out`. Can be either a string from
[`xent`, 'mse', 'nll', 'l1'] or a python `callable` such that
`lossfn(output, target)` returns the loss value. Examples of valid
`callable`s are `F.cross_entropy`, `F.mse_loss`, etc, where `F` is
`torch.nn.functional`. Default: 'xent'
filter_module_by_name:
a function that returns a bool given module names (from
`model.named_modules()`), or None. If not None, then only modules
whose name yields True will be recorded.
cuda:
whether to use cuda or not. Default: True
nseeds:
number of times to repeat the training, each with different seeds.
output_fdict, input_fdict, param_fdict:
function dicts to be used in `_record_coords`. By default, only `l1`
is computed for output activations of each module.
show_progress:
show progress using tqdm. Default: True
one_hot_target:
convert target label into a one-hot vector. This typically is only
used for `'mse'` or `'l1'` losses in classification tasks.
Default: False
Output:
a pandas DataFrame containing recorded results. The column names are
`'width', 'module', 't'` as well as names of statistics recorded, such
as `'l1'` (see `FDICT` for other premade statistics that can be
collected).
Breaking Changes:
In v1.0.0, when `lossfn=='mse'`, the target is automatically converted
to a one hot vector before loss computation. Starting in v1.1.0, this
behavior is turned off, and the user needs to explicitly turn on this
behavior by setting `one_hot_target=True`.
"""
if lr is None:
lr = 0.1 if optimizer == "sgd" else 1e-3
if mup:
from mup.optim import MuAdam as Adam
from mup.optim import MuAdamW as AdamW
from mup.optim import MuSGD as SGD
else:
from torch.optim import SGD, Adam, AdamW
def get_trainable(model):
params = model.parameters()
if filter_trainable_by_name is not None:
params = []
for name, p in model.named_parameters():
if filter_trainable_by_name(name):
params.append(p)
return params
if optimizer == "sgd":
optcls = lambda model: SGD(get_trainable(model), lr=lr)
elif optimizer == "adam":
optcls = lambda model: Adam(get_trainable(model), lr=lr)
elif optimizer == "adamw":
optcls = lambda model: AdamW(get_trainable(model), lr=lr)
elif optimizer is None:
raise ValueError("optimizer should be sgd|adam|adamw or a custom function")
data = _get_coord_data(
neox_args, timers, lr_scheduler, models, dataloader, optcls, **kwargs
)
data["optimizer"] = optimizer
data["lr"] = lr
return data