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 -*-
import functools
import gc
import os
import pickle
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)
def dl_worker_init(worker_id): # pylint: disable=unused-argument
gc.enable()
# Create the training data loader
train_dl = DataLoader(
dataset=train_ds,
@ -275,7 +271,6 @@ def get_train_data_loader(num_worker: int = 0, dataset_generate_func: Optional[C
pin_memory=True,
collate_fn=train_collate_fn,
persistent_workers=num_worker > 0,
worker_init_fn=dl_worker_init,
)
return train_dl, dataset_types