try out not sending the sentinel

EleutherAI · Apr 27, 2023 · d24ded8 · d24ded8
1 parent 44ec152
commit d24ded8
Show file tree

Hide file tree

Showing 2 changed files with 3 additions and 3 deletions.
diff --git a/elk/utils/fsdp.py b/elk/utils/fsdp.py
@@ -248,7 +248,7 @@ def _worker(
  out_queue.put(output_applied)
 
  # Indicate we're done with this dataset
- out_queue.put_nowait(None)
+ # out_queue.put_nowait(None)
 
  # Clean up the FSDP process group
  if fsdp_port is not None:

diff --git a/tests/test_inference_server.py b/tests/test_inference_server.py
@@ -36,7 +36,7 @@ def test_inference_server_normal():
  )
 
 
-def test_inference_server_fsdp():
+def test_inference_server_fsdp_single():
  model_str = "sshleifer/tiny-gpt2"
  server = InferenceServer(
  model_str=model_str, num_workers=2, fsdp=True, cpu_offload=True
@@ -51,7 +51,7 @@ def test_inference_server_fsdp():
  inputs = dict(input_ids=input_ids)
  # make the dict a dataset, while still making it a pytorch dataset
  input_dataset = Dataset.from_dict(inputs)
- input_dataset.set_format(type="torch")
+ input_dataset.set_format(tyxpe="torch")
  outputs = server.map(dataset=input_dataset, closure=lambda x: x)[0]
  assert (
  type(outputs) == transformers.modeling_outputs.CausalLMOutputWithCrossAttentions