fix the fx tracer leaf module logic

[TroyGarden]

Summary:

context

• in response to a TorchRec User Group post
• in torchrec.train_pipeline, the modified fx tracer
  the algorithm is based on "named_modules()" API from torch.nn.Module, there
  could be corner cases that a returned "leaf_module" is not actually called in the
  forward pass. for example the umbrella_module would be considered as the top-level
  leaf module but it actually won't appear in a fx-traced graph.

example

• traced graph

opcode         name          target                   args                          kwargs
-------------  ------------  -----------------------  ----------------------------  --------
placeholder    x             x                        ()                            {}
call_module    nested        nested                   (x,)                          {}
call_module    umbrella1_m1  umbrella1.m1             (x,)                          {}
call_module    umbrella1_m2  umbrella1.m2             (x,)                          {}
call_function  add           <built-in function add>  (umbrella1_m1, umbrella1_m2)  {}
call_function  add_1         <built-in function add>  (nested, add)                 {}
call_module    umbrella2_m1  umbrella2.m1             (x,)                          {}
call_function  add_2         <built-in function add>  (add_1, umbrella2_m1)         {}
call_module    umbrella2_m2  umbrella2.m2             (x,)                          {}
call_function  add_3         <built-in function add>  (add_2, umbrella2_m2)         {}
call_module    umbrella3     umbrella3                (x,)                          {}
call_function  add_4         <built-in function add>  (add_3, umbrella3)            {}
call_module    umbrella4_m1  umbrella4.m1             (x,)                          {}
call_function  add_5         <built-in function add>  (add_4, umbrella4_m1)         {}
call_module    umbrella4_m2  umbrella4.m2             (x,)                          {}
call_function  add_6         <built-in function add>  (add_5, umbrella4_m2)         {}
output         output        output                   (add_6,)                      {}

Differential Revision: D80858379

[meta-codesync]

@TroyGarden has exported this pull request. If you are a Meta employee, you can view the originating Diff in D80858379.