make pp_data_dir configurable and add pp_vis dependencies (#183)

Author: lhzhang333

primus/modules/trainer/megatron/trainer.py

@@ -1732,7 +1732,7 @@ def get_e2e_base_metrics():
         if args.dump_pp_data:
             from .utils import dump_pp_data
 
-            pp_data_dir = "output/pp_data"
+            pp_data_dir = os.environ.get("DUMP_PP_DIR", "output/pp_data")
             dump_pp_data(args, get_num_microbatches(), pp_data_dir)
             log_rank_0(f"pp schedule data dumped to {pp_data_dir}")
 

tools/visualization/pp_vis/requirements.txt

@@ -1,2 +1,3 @@
 numpy
 matplotlib
+tornado