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

[FAW] move coloparam setting in test code. (#1429)

Browse Source
This commit is contained in:
Jiarui Fang
2022-08-10 14:31:53 +08:00
committed by GitHub
parent cb98cf5558
commit 10b3df65c8
2 changed files with 6 additions and 6 deletions
Download Patch File Download Diff File Expand all files Collapse all files

3
colossalai/nn/_ops/cache_embedding/parallel_freq_aware_embedding.py
View File

@@ -67,9 +67,6 @@ class ParallelFreqAwareEmbeddingBag(BaseEmbeddingBag):
self.init_parameters()
else:
assert isinstance(_weight, ColoParameter), "initialized weight must in type of ColoParameter"
_weight.process_group = ProcessGroup(tp_degree=self.world_size)
_weight.set_tensor_spec(ShardSpec(dims=[-1], num_partitions=[self.world_size]),
ComputeSpec(ComputePattern.TP1D))
self._weight = _weight
@property