The problem
I am currently working on a project that I sadly can't share with you. The project is about hyper-parameter optimization for neural networks, and it requires that I train multiple neural network models (more than I can store on my GPU) in parallel. The network architectures stay the same, but the network parameters and hyper-parameters are subjected to change between each training interval. I am currently achieving this using PyTorch on a linux environment in order to allow my NVIDIA GTX 1660 (6GB RAM) to use the multiprocessing feature that PyTorch provides.
Code (simplified):
def training_function(checkpoint):
load(checkpoint)
train(checkpoint)
unload(checkpoint)
for step in range(training_steps):
trained_checkpoints = list()
for trained_checkpoint in pool.imap_unordered(training_function, checkpoints):
trained_checkpoints.append(trained_checkpoint)
for optimized_checkpoint in optimize(trained_checkpoints):
checkpoints.update(optimized_checkpoint)
I currently test with a population of 30 neural networks (i.e. 30 checkpoints) with the MNIST and FashionMNIST datasets which consists of 70 000 (50k training, 10k validation, 10k testing) 28x28 images with 1 channel each respectively. The network I train is a simple Lenet5 implementation.
I use a torch.multiprocessing pool and allow 7 processes to be spawned. Each process uses some of the GPU memory available just to initialize the CUDA environment in each process. After training, the checkpoints are adapted with my hyper-parameter optimization technique.
The load function in the training_function loads the model- and optimizer state (holds the network parameter tensors) from a local file into GPU memory using torch.load. The unload saves the newly trained states back to file using torch.save and deletes them from memory. I do this because PyTorch will only detach GPU tensors when no variable is referencing them. I have to do this because I have limited GPU memory.
The current setup works, but each CUDA initialization occupies over 700MB of GPU RAM, and so I am interested if there are other ways I could do this that may use less memory without a penalty to efficiency.
My attempts
I suspected I could use a thread pool in order to save some memory, and it did. By spawning 7 threads instead of 7 processes, CUDA was only initialized once, which saved almost half of my memory. However, this lead to a new problem in which the GPU only utilized approx. 30% utilization according to nvidia-smi that I am monitoring in a separate linux terminal. Without threads, I get around 85-90% utilization.
I also messed around with torch.multiprocessing.set_sharing_strategy which is currently set to 'file_descriptor', but with no luck.
My questions
Is there a better way to work with multiple model- and optimizer states without saving and loading them to files while training? I have tried to move the model to CPU using model.cpu() before saving the state_dict, but this did not work in my implementation (memory leaks).
Is there an efficient way I can train multiple neural networks at the same time that uses less GPU memory? When searching the web, I only find references to nn.DataParallel which trains the same model over multiple GPUs by copying it to each GPU. This does not work for my problem.
I will soon have access to multiple, more powerful GPUs with more memory, and I suspect this problem will be less annoying then, but I wouldn't be surprised if there is a better solution I am not getting.
Update (09.03.2020)
For any future readers, if you set out to do something similar to the pseudo code displayed above, and you plan on using multiple GPUs, please make sure to create one multiprocessing pool for each GPU device. Pools don't execute functions in order with the underlying processes it contains, and so you will end up initializing CUDA multiple times on the same process, wasting memory.
Another important note is that while you may be passing the device (e.g. 'cuda:1') to every torch.cuda-function, you may discover that torch does something with the default cuda device 'cuda:0' somewhere in the code, initializing CUDA on that device for every process, which wastes memory on an unwanted and non-needed CUDA initialization. I fixed this issue by using with torch.cuda.device(device_id) that encapsulate the entire training_function.
I ended up not using multiprocessing pools and instead defined my own custom process class that holds the device and training function. This means I have to maintain in-queues for each device-process, but they all share the same out-queue, meaning I can retrieve the results the moment they are available. I figured writing a custom process class was simpler than writing a custom pool class. I desperately tried to keep using pools as they are easily maintained, but I had to use multiple imap-functions, and so the results were not obtainable one at a time, which lead to a less efficient training-loop.
I am now successfully training on multiple GPUs, but my questions posted above still remains unanswered.
Update (10.03.2020)
I have implemented another way to store model- and optimizer statedicts outside of GPU RAM. I have written function that replaces every tensor in the dicts with it's .to('cpu') equivalent. This costs me some CPU memory, but it is more reliable than storing local files.
Update (11.06.2020)
I have still not found a different approach that leads to fewer CUDA initializations while maintaining the same processing speed. From what I've read and come to understand, PyTorch does not infer too much with how CUDA is operating, and leaves that up to NVIDIA.
I have ended up using a pool of custom, device-specific processes, called Workers, that is maintained by my custom pool class (more about this above). In addition, I let each of these Workers take in one or more checkpoints as well as the function that processes them (training, testing, hp optimization) via a Queue. These checkpoints are then processed simultaneously via a python multiprocessing ThreadPool in each Worker and the results are then returned one by one via the return Queue the moment they are ready.
This gives me the parallel procedure I was needing, but the memory issue is still there. Due to time constraints, I have come to terms with it for now.
Related
My python code has two steps. In each step, I train a neural network (primarily using from mesh_transformer.transformer_shard import CausalTransformer and delete the network before the next step that I train another network with the same function. The problem is that in some cases, I receive this error:
Resource exhausted: Failed to allocate request for 32.00MiB (33554432B) on device ordinal 0: while running replica 0 and partition 0 of a replicated computation (other replicas may have failed as well).
I think there is still some remaining stuff in the TPU memory I need to remove except that network. The point here is that both steps are independent, and they don't share any information or variable. But I have to do this sequentially to manage my storage on Google cloud. Also, when I run these two steps separately, it works fine. Is there any way to clean TPU memory thoroughly before going to the next step of my code? I think just removing the network is not enough.
Unfortunately, you can’t clean the TPU memory, but you can reduce memory usage by these options;
The most effective ways to reduce memory usage are to:
Reduce excessive tensor padding
Tensors in TPU memory are padded, that is, the TPU rounds up the sizes of tensors stored in memory to perform computations more efficiently. This padding happens transparently at the hardware level and does not affect results. However, in certain cases the padding can result in significantly increased memory use and execution time.
Reduce the batch size
Slowly reduce the batch size until it fits in memory, making sure that the total batch size is a multiple of 64 (the per-core batch size has to be a multiple of 8). Keep in mind that larger batch sizes are more efficient on the TPU. A total batch size of 1024 (128 per core) is generally a good starting point.
If the model cannot be run on the TPU even with a small batch size
(for example, 64), try reducing the number of layers or the layer
sizes.
You could read more about troubleshooting in this documentation
You can try to clean TPU state after each training and see if that helps with
tf.tpu.experimental.shutdown_tpu_system() call.
Another option is to restart the TPU to clean the memory using:
pip3 install cloud-tpu-client
import tensorflow as tf
from cloud_tpu_client import Client
print(tf.__version__)
Client().configure_tpu_version(tf.__version__, restart_type='always')
I use torch to build a model. In the training loop, I need to compute a dictionary of values which will happen on GPU. Then, every few iterations, I need to use the dictionary which was on GPU to perform some other tasks on CPU. So there is a back and forth between the two. Currently, I use torch.save and torch.load to save and load the said dictionary which is very slow and not thread-safe.
I am almost sure that there is a better way to accomplish what I am trying to do. I understand that copying data to/from GPU causes slowdown but I am looking for a strategy that does not involve the disk.
I am running a quantized TFLite model (in Linux PC) for inference using XNNPack backend. I am aware that TFLite models may suffer high latency for prediction and i'm trying to optimize its performance defining number of threads to TFLite.Interpreter(num_threads=X).
I made some trials using X=[4, 6, 8, None] and the best scenario was with X=4, but this doesn't make sense to me. How it is defined the optimal number of threads? And more, defining num_threads automatically works with multiple CPUs or do i have to use another library/package?
(other optimizations that could speed up inference are very welcome!). The model i'm using is a quantized google BERT.
Thanks.
It depends on your target environment. If the target is a single or dual core machine and you're not allowed to use multiple cores for your application, you should use num_threads=1.
Otherwise, you may use more threads to leverage multiple cores.
If your target only has 4 cores, using higher than 4 doesn't give a performance improvement but gives only memory and context switching overhead. (Also shape of inputs are related depends on op kernel implementation)
Regarding the performance improvement,
usually integer operation is faster than float. So you can optimize your model to use integer operations.
https://www.tensorflow.org/lite/performance/model_optimization
Also if your target has GPU, you could try GPU delegate.
https://www.tensorflow.org/lite/performance/gpu
When training, the model requires gradients and consumes much gpu memory, after training, I want to evaluate the model performance of different steps/epochs in parallel(multi-processing), where each of the subprocess loads different parameters for the model without gradients. In order to spawn as many subprocesses as there could be, I want to COMPLETELY free the gpu memory occupied by the previous model (in training phase).
I've tried del model and torch.empty_cache() as many suggested, but the nvidia-smi still shows the gpu memory is not released, which will prevent me from creating as many evaluation subprocesses as expected (raising the error of CUDA out of memory).
So anyone can help me completely delete the model from gpu? In addition, I don't mind loading the model repeatedly in each subprocess because the model in training is far bigger than the model in evaluating. Or tell me why my proposal is not realistic.
The tf.distribute.experimentalCentralStorageStrategy specifies that Variables are not mirrored, instead, they are placed on CPU and ops are replicated across all GPUs.
If I have a really big model that does not fit on any single GPU, could this be a solution since variables are stored on CPU? I know that there will be networking overhead and that's fine.
This Official TF Tutorial on Youtube states that this could be used to handle "large embeddings" that would not fit on one GPU. Could this also be the case for large variables and activations?
In the official documentation, it states that "if there is only one GPU, all variables and operations will be placed on that GPU." If I only used 1 GPU, it seems that CentralStorageStrategy would be automatically disabled even though storing large variables (embeddings for example) on the CPU instead of GPU could be very valuable since there might not exist a GPU that has enough memory to fit it on device. Is this a design oversight or intended behavior?