Running simulations in parallel

[SachinBM-CE]

Hi !
Is it possible to run multiple simulations in parallel, all updating the same policy / critic network using TorchFort?
If so can you please share some insights into this?
Currently I'm launching one simulation using one set of - (binary, case file)
What if there are multiple of these (binary, case file), but having the same yaml file and torchscript?

[romerojosh]

Hi @SachinBM-CE, running separate binaries in parallel will not update the same model as you describe. Each process would separately load the TorchScript model into memory and train it within the process independently. To train in parallel, you need to use the distributed RL system creation routines (e.g., https://nvidia.github.io/TorchFort/api/c_api.html#torchfort-rl-on-policy-create-distributed-system) and possibly also the _multi variants of the RL training routines if you want to have multiple "environments" per process (e.g., https://nvidia.github.io/TorchFort/api/c_api.html#torchfort-rl-on-policy-update-rollout-buffer-multi).

@azrael417 might be better able to handle follow up questions on this one but he is out of the office until next week.

[SachinBM-CE]

Hi @romerojosh !

Thanks for your reply. As a future step, I wanted to run multiple CFD simulations, each having multiple agents, to train the same RL model.

Otherwise, for training on a single CFD sim -

1. As per previous discussions, I have been using off_policy_update_replay_buffer_multi for training multiple agents in the same CFD simulation. Hope multiple environments and multi-agent RL mean the same as per TorchFort's lingo?

2. Is off_policy_create_distributed_system to be used when running with mpirun -np 4 ./binary? The other variant - off_policy_create_system cannot be used in this scenario?

[romerojosh]

1. I believe it is correct that "multi-env" and "multi-agent" are the same concept, but @azrael417 will have to follow up here as I might be missing some subtle difference between the two concepts.
2. Depends on what you are attempting to do. You need to use create_distributed_system for TorchFort to create a data-parallel model and share gradients across the ranks in the MPI communicator passed to that API. If for some reason you wanted to run multiple independent models concurrently, you can use the non-distributed API even in an MPI launched run with multiple processes.

[SachinBM-CE]

I will wait for a response from @azrael417 regarding the multi-env / multi-agent.

1. By running multiple independent models concurrently, do you mean each MPI process is having it's own copy of the neural network?
2. Is update_replay_buffer also affected by create_distributed_system ? Is the replay buffer having it's local replay buffer corresponding to each MPI Process?
3. How does the train_step vary with create_distributed_system and create_system?

The update_replay_buffer is checking if the batch_dimension is equal or not to n_envs parameter. The batch_dimension which is also equal to the number of agents may vary depending on the number of MPI processes and different zones of the domain. But what remains fixed is the total number of agents across all processes or zones.

For example: Suppose if 2 MPI processes are launched, following data is observed

----------------------------------------
Rank, boundary zone, number of agents
----------------------------------------
   0,             1,               m
   0,             2,               n
   1,             1,               m
   1,             2,               n

Should I do a MPI_Gather operation, to collect (s,a,r,s') from all the processes/zones whose sum total is always fixed, which can be used directly to set the n_envs parameter in the config.yaml? (OR) is there any functionality in TorchFort that can help in this scenario?

P.S.: Just for context, the dimension of s, a, r, s' are respectively - (number of agents, number of features), (number of agents, 1), (number of agents) and (number of agents, number of features). Here the number of agents can be equal to m or n depending on the boundary zone.

[azrael417]

Hello Sachin,

multi does not refer to multi agent but instead multi environment. It allows you to run multiple environments in a single task, but those apply to the same agents. In order to achieve multi agents, you should be able to take care of the independent agents in the backend. What is the actual setup you are using? Do you for example have two boundary layers and want to apply a different agent to each boundary layer? What is the actual model for each agent? An MLP acting on a feature vector and you want to have a different MLP for each wallpoint?

Concerning the gather you do not need to do that, if you initialized the distributed system accordingly all the gradient communication should be taken care of automatically. But before we go there, we need to understand what you are trying to achieve. For example for each agent, you might to create its own distributed system since inside a distributed system, all gradients are shared (which is what you want to avoid when running independent agents).

Best
Thorsten

[SachinBM-CE]

I am trying to reproduce https://www.nature.com/articles/s41467-022-28957-7 using TorchFort and https://github.com/ExtremeFLOW/neko.

From our previous discussions (#48), it was suggested to use MARL via the multi-environment mechanism.

The setup is for wall-modelled LES in a channel flow (two parallel walls). The near wall grid-points act as individual agents (number of agents = number of near-wall grid points). I want to train a single MLP / wall model using multi-agent RL. All the agents contribute to a single shared policy i.e. a single MLP / wall model.

Basically in this paper, they launch multiple CFD runs (with different Reynolds number) and each CFD run has the same number of agents together updating a single MLP / wall-model.

[SachinBM-CE]

I am using the following functions currently:
torchfort_rl_off_policy_create_system, predict, update_replay_buffer, train_step.
I am passing s,a,r,s' with their leading dimension equal to number of agents and giving n_envs parameter under replay_buffer in the yaml file. Is this not the way to work with MARL?

[azrael417]

If you want to share model information, then you should use the distributed system variant. In this case, you can use number of environments == number of agents and pass the arrays/tensors as you suggested accordingly. If you run multi-GPU, and your agents (which share the same weights) live on different GPUs, you need to use the distributed variant. Now you say that you are in a situation where m!=n for different ranks? That is a bit of a more complicated scenario. Is it possible to ensure that the number of environments is the same for each rank, i.e. for example (n+m)/2 or something like that?

[azrael417]

In this case, you can just use MPI to re-distribute the data beforehand, or, you can also split up the volume evenly so that each rank/GPU has the same number of wall points.

[SachinBM-CE]

Just to clarify, if I want to run on the HPC cluster with just CPUs cores,
Suppose, number of compute nodes per job = 1 (single node) and CPU cores per node = N, I should be using create_distributed_system, hope I am right...

Regarding the equal distribution of agents to the ranks, I need to try, but it is not so easy as ParMETIS does the mesh partitioning and I do not know how I can change that...

Additionally, when does the need arise to change the key variable used in these functions? As per my understanding of the documentation, key is a unique identifier for the neural network model...

[azrael417]

Hello Sachin,

if you cannot influence the partitioning then you can use all-to-all or gather/scatter to re-distribute the data between the nodes. However, I would probably not invest too much time in finding a good algorithm for that. I assume that the simulation cost is much higher than the cost for training the NN, so you can gather all data on rank 0 and run a non-distributed version for the RL algorithm there.

With key, you mean the first parameter which goes into all the torchfort functions? This should be unique to the system. For example in the distributed case, only systems with the same key communicate with each other and share the gradients. If the key is different on different ranks, it will not work.

In any case, I think it is easier to use a non distributed setup in your case and simply mpi gather all the data to rank 0 and train the model there. For agent step, just do a local step and broadcast/scatter the action prediction. I am almost certain it would not have a big negative impact on overall performance compared to the case where you balance it out, since that involves a lot of edge case handling depending on how ParMETIS decides to split things up. Does that make sense?

Best and thanks
Thorsten

[azrael417]

It is also easier to debug in case something does not work as expected, since you only have to print replay/rollout buffer entries from that rank then and can dump everything locally.

[SachinBM-CE]

Hi Thorsten,
The MPI_Gatherv implementation on rank 0 worked ! By collecting global (s,a,r,s') arrays from all ranks, I've eliminated the uneven agent distribution problem entirely. The replay buffer updates now handle the complete dataset regardless of how agents are spread across processes. Thanks for the guidance on this solution!

[SachinBM-CE]

I am collecting 100 timesteps of simulation data (s_1, a_1 , r_1, s'_1) ... (s_100, a_100, r_100, s'_100) which is equivalent to one episode (as terminal state is reached) and after which I want to train the MLP.

How do I set the epochs and number of gradient updates for training? Currently I have placed the call to train_step function inside a do loop. How is it ensured that entire episode data is trained?

[azrael417]

Which algorithm is that, some on or off policy one? In both cases, a call to train step only samples a single batch from the buffer (replay or rollout), with the difference that in the PPO case it will go through the episodes chronologically sorted and in case of off policy just sample randomly. If you want to make sure that you train a full episode, store exactly one episode (adjust buffer size) and then wrap it into a do loop with 100 train steps. However, for off policy that is not necessarily a good idea, I would definitely make the buffer big enough to store multiple episodes so that the algorithm can adjust also on past trajectories.
I am not an expert here, this is something you need to play around with. number of gradient updates is gradient accumulation, which is used to mimic larger batch sizes. I would just increase batch size instead.