github/ColossalAI
1
0
Fork 0
mirror of https://github.com/hpcaitech/ColossalAI.git synced 2025-09-15 22:19:38 +00:00
Code Issues Packages Projects Releases Wiki Activity

[fx] PoC of runtime shape consistency application (#1607)

Browse Source 
* [fx] PoC of runtime shape consistency application

* polish code
This commit is contained in:
YuliangLiu0306
2022-09-20 14:00:04 +08:00
committed by GitHub
parent 47b11c432c
commit 7d1bb71d5d
3 changed files with 197 additions and 1 deletions
Download Patch File Download Diff File Expand all files Collapse all files

3
colossalai/auto_parallel/solver/strategies_constructor.py
View File

@@ -408,7 +408,8 @@ class StrategiesConstructor:
sharding_strategy_attribute = ShardingStrategy(name,
output_sharding_spec,
memory_cost=memory_cost,
resharding_costs=resharding_costs)
resharding_costs=resharding_costs,
input_shardings=tuple(input_sharding_specs))
strategies_vector.append(sharding_strategy_attribute)
self.remove_duplicated_strategy(strategies_vector)