diff --git a/examples/mlperf/model_train.py b/examples/mlperf/model_train.py index 0d418938bc..6138c05dfa 100644 --- a/examples/mlperf/model_train.py +++ b/examples/mlperf/model_train.py @@ -2,6 +2,7 @@ import functools import os import time from tqdm import tqdm +import multiprocessing from tinygrad import Device, GlobalCounters, Tensor, TinyJit, dtypes from tinygrad.helpers import getenv, BEAM, WINO @@ -251,6 +252,7 @@ def train_maskrcnn(): pass if __name__ == "__main__": + multiprocessing.set_start_method('spawn') with Tensor.train(): for m in getenv("MODEL", "resnet,retinanet,unet3d,rnnt,bert,maskrcnn").split(","): nm = f"train_{m}" diff --git a/tinygrad/device.py b/tinygrad/device.py index 8f436edeb5..0a2a5e0385 100644 --- a/tinygrad/device.py +++ b/tinygrad/device.py @@ -1,7 +1,7 @@ from __future__ import annotations from collections import defaultdict from typing import TYPE_CHECKING, Any, List, Optional, Dict, Tuple, ClassVar, NamedTuple -import importlib, inspect, functools, pathlib, time, ctypes +import importlib, inspect, functools, pathlib, time, ctypes, os from tinygrad.helpers import ansilen, prod, getenv, colored, all_int, to_function_name, from_mv, flat_mv, diskcache_get, diskcache_put from tinygrad.helpers import DEBUG, CACHECOLLECTING, BEAM, NOOPT, GlobalCounters from tinygrad.shape.symbolic import Variable, sym_infer, sint @@ -23,6 +23,7 @@ class _Device: def __getitem__(self, ix:str) -> Compiled: return self.__get_canonicalized_item(self.canonicalize(ix)) @functools.lru_cache(maxsize=None) # this class is a singleton, pylint: disable=method-cache-max-size-none def __get_canonicalized_item(self, ix:str) -> Compiled: + if DEBUG >= 1: print(f"opening device {ix} from pid:{os.getpid()}") x = ix.split(":")[0].upper() return [cls for cname, cls in inspect.getmembers(importlib.import_module(f'tinygrad.runtime.ops_{x.lower()}')) if (cname.lower() == x.lower() + "device") and x in self._devices][0](ix) # noqa: E501 @functools.cached_property