Skip to content

Commit ec7c0ec

Browse files
Chia-Jung Changfacebook-github-bot
Chia-Jung Chang
authored andcommitted
Enable InputRecorder to run in offline mode and cuda device (#2067)
Summary: Enable `InputRecorder` to run in offline mode and cuda device. Specificlaly `InputRecorder` is called during Sandcastle CI tests, which do not allow us to access HuggingFace website. Differential Revision: D73179650
1 parent a96eeb1 commit ec7c0ec

File tree

1 file changed

+19
-9
lines changed

1 file changed

+19
-9
lines changed

torchao/_models/_eval.py

+19-9
Original file line numberDiff line numberDiff line change
@@ -183,18 +183,27 @@ def __init__(
183183
vocab_size=32000,
184184
pad_token=0,
185185
device="cpu",
186+
online_access=True,
186187
):
187-
try:
188-
super().__init__()
189-
except TypeError:
190-
# lm_eval 0.4.2 removed the default init
191-
super().__init__("gpt2", device="cpu")
188+
if online_access:
189+
try:
190+
super().__init__()
191+
except TypeError:
192+
# lm_eval 0.4.2 removed the default init
193+
super().__init__("gpt2", device=device)
194+
else:
195+
# Create a minimal implementation when run offline
196+
self._device = torch.device(device)
197+
self._max_length = 2048
198+
self._max_gen_toks = 256
199+
self._batch_size = 1
192200

193201
self.tokenizer = tokenizer
194202
self._device = torch.device(device)
195203
self.vocab_size = vocab_size
196204
self._max_seq_length = calibration_seq_length
197205
self.calibration_seq_length = calibration_seq_length
206+
self.online_access = online_access
198207

199208
# need to take inps and convert to corrent input
200209
# for model
@@ -257,10 +266,11 @@ def record_inputs(
257266
calibration_tasks,
258267
calibration_limit,
259268
):
260-
try:
261-
lm_eval.tasks.initialize_tasks()
262-
except:
263-
pass
269+
if self.online_access:
270+
try:
271+
lm_eval.tasks.initialize_tasks()
272+
except:
273+
pass
264274

265275
task_dict = get_task_dict(calibration_tasks)
266276
print("Obtaining GPTQ calibration inputs on: ", calibration_tasks)

0 commit comments

Comments
 (0)