[Docs] Add reinforcement learning example illustrating gpu-to-gpu RDT and GRPO.

[crypdick]

Description

Example for first blog in the RDT series using NIXL for GPU-GPU tensor transfers.

[stephanie-wang]

Can you expand on the docstring a bit more? For people who are not as familiar with RL.

Would be good to say how sampling is done, when samples are kicked out, etc.

[crypdick]

updated the docstrings. lmk what you think

[stephanie-wang]

I thought original PPO paper suggested not using the KL divergence term?

I'd prefer to remove it assuming this particular example still works. Have you tried running without it to see if it will converge? If that doesn't work, the 1 weight update step behind seems simpler.

[cursor]

Bug: Generator Uses Stale Weights Due to Timing

Race condition in the training loop: generator.generate.remote(states) is called at line 385 before ray.wait([weights_updated_ref]) at line 390. This means that in each iteration (except the first), the generate call is queued on the single-threaded Generator actor before confirming that the previous weight update has completed. While the actor processes tasks sequentially, the generate task gets queued before the ray.wait confirms the update is done, potentially causing the generator to use stale weights. The generate call should be moved to after the ray.wait to ensure the generator always uses the most recent weights.

 

[cursor]

Bug: Unit Vector Function Restricts Single Batch Usage

The sample_unit_vector function's assertion batch_size > 1 is overly restrictive. The vector normalization logic works correctly for batch_size=1, unnecessarily limiting the function's utility and potentially causing failures in valid use cases.

 

[cursor]

Bug: Unit Vector Sampling Restriction Error

The sample_unit_vector function includes an assertion batch_size > 1 that is overly restrictive. The function's logic correctly handles batch_size=1 for sampling a single unit vector, so this assertion prevents valid usage without an apparent algorithmic or mathematical reason.

 

[daiping8]

Hello!I tested the code on a cluster with 1 head and 1 worker node, each equipped with an A800 GPU. After testing different transport backends, I obtained the following results. According to the data, there is no significant difference in execution time among the various transport backends. Is this situation normal? Thank you !!!

================================================================================
SUMMARY RESULTS
================================================================================
steps: 500
Transport       Init (s)     Steps (s)    Avg Step (s)    Total (s)   
--------------------------------------------------------------------------------
nixl            3.6700       145.9446     0.2919          149.6147    
nccl            3.6614       145.8298     0.2917          149.4911    
gloo            3.6215       144.9413     0.2899          148.5628    
object_store    3.5330       147.6611     0.2953          151.1940    
================================================================================

• Init (s) - Initialization Time (seconds): This is the time spent on system initialization, including operations such as creating various Ray Actors (ReplayBuffer, Learner, Scorer, Generator), initializing model weights, and pre-filling the replay buffer. This period is not included in the training steps and represents a one-time setup cost before training begins.
• Steps (s) - Total Training Steps Time (seconds): This is the total time spent across all training steps. It starts timing from the beginning of the first training step to the end of the last training step, covering the duration required to execute the specified number of complete training loops (as defined by the --steps parameter). This metric reflects the performance of the actual training process.
• Avg Step (s) - Average Time Per Step (seconds): This is the average time per training step, calculated by dividing the total training time (Steps (s)) by the total number of training steps. This metric can be used to compare efficiency differences per training step under different communication methods.
• Total (s) - Total Time (seconds): This is the end-to-end total duration of the entire training process, including both initialization time and training steps time. It covers the full time span from actor creation to completion of all training and resource cleanup.