ghstack-source-id: e8c2ff8
Pull Request resolved: #164738

ghstack-source-id: 0da84d2
Pull Request resolved: #164738

This reverts commit ab01a0d.

Reverted #164738 on behalf of https://github.com/eellison due to merge sets makes this trickier ([comment](#164581 (comment)))

- Update the Memory Estimator to use node storages for analysis, which simplifies book keeping, as opposed to manually looking at operator schema. This will also allow me to reuse this component elsewhere.

- Factor out into separate class, so that this same logic can be used  in scheduling (node allocations / aliasing / uses)

- Adds Tests for correctness - right now only on fwd/bwd by itself, not with both.

Pull Request resolved: #164783
Approved by: https://github.com/ruisizhang123
ghstack dependencies: #164738

Respect max_coll_distance from overlap scheduler in bucketing, also, add an optimization in path searching.

Pull Request resolved: #164944
Approved by: https://github.com/IvanKobzarev
ghstack dependencies: #164738, #164783

…ted (#164945)

Pull Request resolved: #164945
Approved by: https://github.com/IvanKobzarev
ghstack dependencies: #164738, #164783, #164944

Add Memory Tracker utility, which will track live memory given alternate ordering of nodes.

Pull Request resolved: #165059
Approved by: https://github.com/ezyang, https://github.com/IvanKobzarev
ghstack dependencies: #164738, #164783, #164944, #164945

Bucketing a number of smallish improvements:

- Account for bucketing in overlap calculation: if an in-flight collective exists with the same bucket key, reduce new collectives estimated time by its latency time
-  Update compute domination so we are ordering based on compute idx, as opposed to compute depth, so we never reorder compute. this makes it a bit easier to reason about memory, and pre-fetching, although we can exploring reordering in the future.
- When we wait on a collective, force all collectives on the same process group as it that were enqueued prior to the collective to wait as well.

Better Memory Handling:
- Pre-fetch limiting - when scheduling collectives for overlap, only pre-fetch up to a certain distance, then schedule off-path collectives (which are typically memory reducing).
- When we are above peak memory, schedule waits.

TODO:
- for each compute node, we know its original memory in the graph. we could limit pre-fetching that goes across peak memory
- By scheduling off-path collectives for overlap, we reduce memory, but if there weren't enough compute for overlap, we need to proactively schedule them. not an issue yet on examples.
- config some hard coded constants, clean up enablement (can do in subsequent pr)

On small llama 2d backward :
578 of 618 potentially hideable collectives hidden
original mem 14.4GB, rescheduled mem, 15.9GB

on forward:
254/256 potentially hideable collectives hidden
original mem 5.8 gb, reshceduled mem 5.8GB

WIP: adding tests

Pull Request resolved: pytorch#165318
Approved by: https://github.com/ezyang, https://github.com/IvanKobzarev
ghstack dependencies: pytorch#164738, pytorch#164783, pytorch#164944, pytorch#164945, pytorch#165059

Original work by @ShatianWang, with lints applied. I am going to a few changes and add tests in subsequent prs but I want to preserve original commit first.

Pull Request resolved: pytorch#164738
Approved by: https://github.com/IvanKobzarev
ghstack dependencies: pytorch#164568, pytorch#164569, pytorch#164581

This reverts commit ab01a0d.

Reverted pytorch#164738 on behalf of https://github.com/eellison due to merge sets makes this trickier ([comment](pytorch#164581 (comment)))

Original work by @ShatianWang, with lints applied. I am going to a few changes and add tests in subsequent prs but I want to preserve original commit first.

Pull Request resolved: pytorch#164738
Approved by: https://github.com/IvanKobzarev

)

- Update the Memory Estimator to use node storages for analysis, which simplifies book keeping, as opposed to manually looking at operator schema. This will also allow me to reuse this component elsewhere.

- Factor out into separate class, so that this same logic can be used  in scheduling (node allocations / aliasing / uses)

- Adds Tests for correctness - right now only on fwd/bwd by itself, not with both.

Pull Request resolved: pytorch#164783
Approved by: https://github.com/ruisizhang123
ghstack dependencies: pytorch#164738

Respect max_coll_distance from overlap scheduler in bucketing, also, add an optimization in path searching.

Pull Request resolved: pytorch#164944
Approved by: https://github.com/IvanKobzarev
ghstack dependencies: pytorch#164738, pytorch#164783

…ted (pytorch#164945)

Pull Request resolved: pytorch#164945
Approved by: https://github.com/IvanKobzarev
ghstack dependencies: pytorch#164738, pytorch#164783, pytorch#164944

Add Memory Tracker utility, which will track live memory given alternate ordering of nodes.

Pull Request resolved: pytorch#165059
Approved by: https://github.com/ezyang, https://github.com/IvanKobzarev
ghstack dependencies: pytorch#164738, pytorch#164783, pytorch#164944, pytorch#164945

Bucketing a number of smallish improvements:

- Account for bucketing in overlap calculation: if an in-flight collective exists with the same bucket key, reduce new collectives estimated time by its latency time
-  Update compute domination so we are ordering based on compute idx, as opposed to compute depth, so we never reorder compute. this makes it a bit easier to reason about memory, and pre-fetching, although we can exploring reordering in the future.
- When we wait on a collective, force all collectives on the same process group as it that were enqueued prior to the collective to wait as well.

Better Memory Handling:
- Pre-fetch limiting - when scheduling collectives for overlap, only pre-fetch up to a certain distance, then schedule off-path collectives (which are typically memory reducing).
- When we are above peak memory, schedule waits.

TODO:
- for each compute node, we know its original memory in the graph. we could limit pre-fetching that goes across peak memory
- By scheduling off-path collectives for overlap, we reduce memory, but if there weren't enough compute for overlap, we need to proactively schedule them. not an issue yet on examples.
- config some hard coded constants, clean up enablement (can do in subsequent pr)

On small llama 2d backward :
578 of 618 potentially hideable collectives hidden
original mem 14.4GB, rescheduled mem, 15.9GB

on forward:
254/256 potentially hideable collectives hidden
original mem 5.8 gb, reshceduled mem 5.8GB

WIP: adding tests

Pull Request resolved: pytorch#165318
Approved by: https://github.com/ezyang, https://github.com/IvanKobzarev
ghstack dependencies: pytorch#164738, pytorch#164783, pytorch#164944, pytorch#164945, pytorch#165059

ghstack-source-id: afb2d8c
Pull Request resolved: pytorch#164738