Fix exception in ray_worker (#144)

richardsliu · web-flow · commit 9717eb91f7e2 · 2024-07-09T13:02:25.000-07:00
* Fix exception in ray_worker

* fix format
diff --git a/jetstream_pt/ray_worker.py b/jetstream_pt/ray_worker.py
@@ -350,15 +350,17 @@ def _call_model_generate(
     new_mask = mask.at[:, current_position].set(0)
     if self.env.quant_config.enable_kv_quantization:
       caches_obj = [
-          cache_manager.Int8KVCacheGenerate(k, v, ks, vs, input_indexes)
+          cache_manager.Int8KVCacheGenerate(
+              k, v, ks, vs, input_indexes, env=self.env
+          )
           for (k, v), (ks, vs) in torchjax.to_torch(
               list(zip(caches, cache_scales))
           )
       ]
     else:
       caches_obj = [
           cache_manager.KVCacheGenerate(
-              k, v, input_indexes, self.cache_sharding
+              k, v, input_indexes, self.cache_sharding, env=self.env
           )
           for k, v in torchjax.to_torch(caches)
       ]

Original file line number	Diff line number	Diff line change
`@@ -350,15 +350,17 @@ def _call_model_generate(`
`350`	`350`	`new_mask = mask.at[:, current_position].set(0)`
`351`	`351`	`if self.env.quant_config.enable_kv_quantization:`
`352`	`352`	`caches_obj = [`
`353`		`- cache_manager.Int8KVCacheGenerate(k, v, ks, vs, input_indexes)`
	`353`	`+ cache_manager.Int8KVCacheGenerate(`
	`354`	`+ k, v, ks, vs, input_indexes, env=self.env`
	`355`	`+ )`
`354`	`356`	`for (k, v), (ks, vs) in torchjax.to_torch(`
`355`	`357`	`list(zip(caches, cache_scales))`
`356`	`358`	`)`
`357`	`359`	`]`
`358`	`360`	`else:`
`359`	`361`	`caches_obj = [`
`360`	`362`	`cache_manager.KVCacheGenerate(`
`361`		`- k, v, input_indexes, self.cache_sharding`
	`363`	`+ k, v, input_indexes, self.cache_sharding, env=self.env`
`362`	`364`	`)`
`363`	`365`	`for k, v in torchjax.to_torch(caches)`
`364`	`366`	`]`