Merge pull request #9 from LokiMetaSmith/fix-dataloader-typeerror

Fix TypeError in tokenizing_distributed_data_loader and robustness in…
This commit is contained in:
Lawrence R Kincheloe III 2025-11-20 23:12:47 -06:00 committed by GitHub
commit 104308cf78
No known key found for this signature in database
GPG Key ID: B5690EEEBB952194
2 changed files with 5 additions and 2 deletions

View File

@ -26,7 +26,9 @@ def print0(s="",**kwargs):
for arg in sys.argv[1:]:
if '=' not in arg:
# assume it's the name of a config file
assert not arg.startswith('--')
if arg.startswith('--'):
# ignore flags like --help or others without =
continue
config_file = arg
print0(f"Overriding config with {config_file}:")
with open(config_file) as f:

View File

@ -81,7 +81,8 @@ def tokenizing_distributed_data_loader_with_state(B, T, split, tokenizer_threads
state_dict = {"pq_idx": pq_idx, "rg_idx": rg_idx} # we need this in case we wish to approximately resume training
yield inputs, targets, state_dict
def tokenizing_distributed_data_loader(*args, **kwargs):
def tokenizing_distributed_data_loader(*args, device="cuda", **kwargs):
# helper function that only emits the inputs/targets and not the state_dict
kwargs["device"] = device
for inputs, targets, state_dict in tokenizing_distributed_data_loader_with_state(*args, **kwargs):
yield inputs, targets