LayerNorm OP draft

[lambda7xx]

Description of changes:

• how to handle int64_t effective_batch_size, effective_num_elements in LayerNormPerDeviceState
• how to handle gamma and beta in forward and backward
• leave some comment

Related Issues:

Linked Issues:

• Issue Update LayerNorm operator #1070

Issues closed by this PR:

• Closes #

---

This change is 

[reyna-abhyankar]

Can you apply the format script?

[lambda7xx]

ok

[lambda7xx]

Reviewable status: 0 of 5 files reviewed, 17 unresolved discussions (waiting on @lambda7xx and @reyna-abhyankar)

---

lib/runtime/src/ops/layer_norm.cc line 0 at r1 (raw file):

Previously, reyna-abhyankar (Reyna Abhyankar) wrote…

(Reviewable was unable to map this GitHub inline comment thread to the right spot — sorry!)

Keep these

using Legion::Context;
using Legion::PhysicalRegion;
using Legion::Runtime;
using Legion::Task;

Done.

---

lib/runtime/src/ops/layer_norm.cc line 0 at r1 (raw file):

Previously, reyna-abhyankar (Reyna Abhyankar) wrote…

(Reviewable was unable to map this GitHub inline comment thread to the right spot — sorry!)

@lambda7xx this is the original code that sets effective_num_elements and effective_batch_size. Seems like we can get this from input

Done.

---

lib/runtime/src/ops/layer_norm.cc line 0 at r1 (raw file):

Previously, reyna-abhyankar (Reyna Abhyankar) wrote…

Split the signature out of register task (see other merged PRs)

Done.

---

lib/runtime/src/ops/layer_norm.cc line 43 at r1 (raw file):

Previously, reyna-abhyankar (Reyna Abhyankar) wrote…

I think this is correct since gamma and beta are parameters

Done.

---

lib/runtime/src/ops/layer_norm.cc line 60 at r1 (raw file):

Previously, reyna-abhyankar (Reyna Abhyankar) wrote…

Based on original permissions, I think it's RW and for backward only the grads are RW

Done.

---

lib/runtime/src/ops/layer_norm.cc line 144 at r1 (raw file):

Previously, reyna-abhyankar (Reyna Abhyankar) wrote…

                                  InputParallelTensorDesc const &input_shape,

I think we should use InputParallelTensorDesc for input shape always because it also has the IsTrainable

Done.

---

lib/runtime/src/ops/layer_norm.cc line 162 at r1 (raw file):

Previously, reyna-abhyankar (Reyna Abhyankar) wrote…

I think it's the same shape for both

Done.

---

lib/runtime/src/ops/layer_norm.cc line 195 at r1 (raw file):

Previously, reyna-abhyankar (Reyna Abhyankar) wrote…

  fwd.add_weight_slot(GAMMA);
  fwd.add_weight_slot(BETA);

Done.

---

lib/runtime/src/ops/layer_norm.cc line 138 at r2 (raw file):

Previously, reyna-abhyankar (Reyna Abhyankar) wrote…

In the original code I don't think it's divided by num_replicas to get effective_batch_size. Just inputs->volume() / M

Done.

---

lib/runtime/src/ops/layer_norm.cc line 138 at r2 (raw file):

Previously, reyna-abhyankar (Reyna Abhyankar) wrote…

Done.

---

lib/runtime/src/ops/layer_norm.cc line 141 at r2 (raw file):

Previously, reyna-abhyankar (Reyna Abhyankar) wrote…

  effective_num_elements = M;
  effective_batch_size = input.shape.get_volume() / M;

Done.

---

lib/runtime/src/ops/layer_norm.cc line 187 at r2 (raw file):

Previously, reyna-abhyankar (Reyna Abhyankar) wrote…

fwd_binding.bind(GAMMA, input_shape);
fwd_binding.bind(BETA, input_shape);

Done.

---

lib/runtime/src/ops/layer_norm.cc line 217 at r4 (raw file):

Previously, reyna-abhyankar (Reyna Abhyankar) wrote…

OpTaskSignature bwd_signature<LAYERNORM_BWD_TASK_ID>()  {

Done.

---

lib/runtime/src/ops/layer_norm.cc line 223 at r4 (raw file):

Previously, reyna-abhyankar (Reyna Abhyankar) wrote…

OpTaskSignature init_signature<LAYERNORM_INIT_TASK_ID>()  {

Done.

---

lib/runtime/src/ops/layer_norm.cc line 236 at r4 (raw file):

Previously, reyna-abhyankar (Reyna Abhyankar) wrote…

  register_task(LAYERNORM_INIT_TASK_ID, "LayerNorm init", init_signature<LAYERNORM_INIT_TASK_ID>(), init_task);

Done.

---

lib/runtime/src/ops/layer_norm.cc line 247 at r4 (raw file):

Previously, reyna-abhyankar (Reyna Abhyankar) wrote…

      LAYERNORM_BWD_TASK_ID, "LayerNorm backward",  bwd_signature<AYERNORM_BWD_TASK_ID>() , backward_task);

Done.

---

deps/fmt line 0 at r2 (raw file):

Previously, reyna-abhyankar (Reyna Abhyankar) wrote…

Use most updated submodules

Done.