In colossalai, there are many ways to run DQN in a distributed way. However, for shallow networks like DQN, using tensor-parallel, pipeline or ZeRO to reduce memory use of GPUs is unnecessary. Therefore, we only present a naive DQN using data-parallel feature of colossalai. Readers could easily scale their training to multiple GPUs by changing the number of GPUs used.
In this example, we support Cart_Pole and Atari games such as Pong. We present two networks for different types
of environments. For Atari games, you should use networks with convolutional layers such as Atari_Network in this
example. For other simple gym environments, you could use similar network structure of Cart_Pole_Network.
If you want to use other environments in gym or Atari games, you could add your environment name in Gym_envs or
Atari_envs located in ./model.utils. Note for Atari environment, we provide some wrappers to deal with original
environment input. Therefore, if you don't want to change these wrappers, you could simply use Atari games with
NoFrameskip such as PongNoFrameskip-v4. If you want to use customize your environment input, you should adapt
init_atari_env, make_atari in ./model.utils and corresponding wrappers in ./model.wrappers to support the
environment you are going to use.
After you install all packages, you could simply start your training on single node by using:
torchrun --standalone --nproc_per_node=<num_gpus> main.py --config=./configs/<config_file> --from_torchyou need to replace <num_gpus> to number of GPUs used, and substitute <config_file> by config files you are
going to use. For customizing configuration files, you could check two config files for more information. We provide
detailed comments there. And we also provide logs and tensorboardX to record your DQN performance. you could make use
of them in your training.
When you are using multiple GPUs for data-parallel training, every GPU would create their own replay buffer and
sample transitions from it using batch size in your configuration files. Every GPU would hold an DQN agent with the same
network structure which forwards with their own transitions data but backwards with global all-reduced gradient.
Therefore, keeping batch size same but using multiple GPUs for data-parallel to train this example is similar to
enlarge batch size. If you want to use the same batch size, please reduce your batch size to batch_size / num_gpus
in your configs. However, if you are going to use larger batch size, please scale your learning rate and reduce
the total_step in your configs for better performance.
As I mentioned before, using tensor-parallel, pipeline or ZeRO in training DQN is less effective. However, if you want to
explore other features usage, you could check GPT example
for ZeRO and pipeline usage. For using tensor-parallel, you could replace torch.nn layers of your networks with layers
in colossalai.nn.layer.colossalai_layer, and set your parallel setting in your configuration file. And colossalai
would automatically use tensor-parallel. More details please check 1D_tensor_parallel feature.