diff --git a/setup.py b/setup.py
index 917174e7b..292cba063 100644
--- a/setup.py
+++ b/setup.py
@@ -52,6 +52,7 @@
 TRAINING_REQUIRES = [
     "trl == 0.11.4",
     "peft == 0.14.0",
+    "neuronx-distributed == 0.9.0",
 ]
 
 EXTRAS_REQUIRE = {