github
/
InternLM
mirror of https://github.com/InternLM/InternLM
1
0
Fork 0
Code Issues Projects Releases Wiki Activity

revert worker_init_fn since we move gc disable after dataloader launching

pull/545/head
zigzagcai 2023-12-29 19:34:38 +08:00
parent df0acdee43
commit f8eaf618af
1 changed files with 0 additions and 5 deletions
Show Stats Download Patch File Download Diff File Expand all files Collapse all files

5
internlm/train/training_internlm.py
View File

@ -2,7 +2,6 @@
# -*- encoding: utf-8 -*- # -*- encoding: utf-8 -*-
import functools import functools
import gc
import os import os
import pickle import pickle
import time import time
@ -264,9 +263,6 @@ def get_train_data_loader(num_worker: int = 0, dataset_generate_func: Optional[C
) )
train_collate_fn = partial(packed_collate_fn, packed_length=data_cfg.packed_length) train_collate_fn = partial(packed_collate_fn, packed_length=data_cfg.packed_length)
def dl_worker_init(worker_id): # pylint: disable=unused-argument
gc.enable()
# Create the training data loader # Create the training data loader
train_dl = DataLoader( train_dl = DataLoader(
dataset=train_ds, dataset=train_ds,
@ -275,7 +271,6 @@ def get_train_data_loader(num_worker: int = 0, dataset_generate_func: Optional[C
pin_memory=True, pin_memory=True,
collate_fn=train_collate_fn, collate_fn=train_collate_fn,
persistent_workers=num_worker > 0, persistent_workers=num_worker > 0,
worker_init_fn=dl_worker_init,
) )
return train_dl, dataset_types return train_dl, dataset_types