ghstack-source-id: ac9d5c6
Pull Request resolved: #164783

ghstack-source-id: f9c6340
Pull Request resolved: #164783

ghstack-source-id: ca9908e
Pull Request resolved: #164783

ghstack-source-id: e813a82
Pull Request resolved: #164783

ghstack-source-id: 334a118
Pull Request resolved: #164783

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

ghstack-source-id: 334a118
Pull Request resolved: pytorch#164783

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

)

- 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: f029d2d
Pull Request resolved: pytorch#164783